Tensorflow版本yolo v3源码阅读笔记(6)

在这一篇中，我们开始分析测试函数的代码。

实话实说，在神经网络的代码中，测试部分是最有趣的，因为和前面枯燥的神经网络训练相比，测试可以马上见到效果。

下面开始代码的分析吧。

导入需要的库

import cv2
import os
import shutil
import numpy as np
import tensorflow as tf
import core.utils as utils
from core.config import cfg
from core.yolov3 import YOLOv3, decode

设置参数

INPUT_SIZE   = 416#yolov3网络的图片输入大小
NUM_CLASS    = len(utils.read_class_names(cfg.YOLO.CLASSES))#yolov3网络可以检测到的物体的种类个数
CLASSES      = utils.read_class_names(cfg.YOLO.CLASSES)#yolov3检测到的物体的名称

创建文件的保存路径

predicted_dir_path = '../mAP/predicted'#图片真实框的保存信息
ground_truth_dir_path = '../mAP/ground-truth'#图片预测框的保存信息
if os.path.exists(predicted_dir_path): shutil.rmtree(predicted_dir_path)#循环删除预测框文件目录
if os.path.exists(ground_truth_dir_path): shutil.rmtree(ground_truth_dir_path)#循环删除真实框文件目录
if os.path.exists(cfg.TEST.DECTECTED_IMAGE_PATH): shutil.rmtree(cfg.TEST.DECTECTED_IMAGE_PATH)#循环删除带框图像的文件目录

os.mkdir(predicted_dir_path)#新建预测框文件目录
os.mkdir(ground_truth_dir_path)#新建真实框文件目录
os.mkdir(cfg.TEST.DECTECTED_IMAGE_PATH)#新建带框图片的文件目录

构建模型

# Build Model
#构建模型的输入
input_layer  = tf.keras.layers.Input([INPUT_SIZE, INPUT_SIZE, 3])
feature_maps = YOLOv3(input_layer)#模型输入特征图

bbox_tensors = []
for i, fm in enumerate(feature_maps):
    bbox_tensor = decode(fm, i)#处理模型输出的特征图
    bbox_tensors.append(bbox_tensor)#保存特征图信息

#构建模型，指定输入和输出
model = tf.keras.Model(input_layer, bbox_tensors)
model.load_weights("./yolov3")#加载模型的预处理参数

测试

1、首先打开测试图片文件夹。

with open(cfg.TEST.ANNOT_PATH, 'r') as annotation_file:

2、然后遍历测试文件夹中的每一张图片信息，包括图片本身以及图片的真实框、类别信息等。

for num, line in enumerate(annotation_file):

假设此时是第200张图片，那么此时的line为：

line = 'E:\\Pycharm\\code\\Jupyter\\tensorflow2.0\\My_net\\YOLO_v3\\data\\dataset\\test\\000200.jpg 141,305,197,361,7 87,23,143,79,0 46,114,102,170,1 16,55,44,83,5 31,253,87,309,5 257,305,341,389,6 317,163,401,247,4 164,118,248,202,0\n'

3、使用strip()去除首尾空格，用split()分割字符串line。

annotation = line.strip().split()

分割后的annotation格式如下：

annotation = ['E:\\Pycharm\\code\\Jupyter\\tensorflow2.0\\My_net\\YOLO_v3\\data\\dataset\\test\\000200.jpg',
              '141,305,197,361,7',
              '87,23,143,79,0',
              '46,114,102,170,1',
              '16,55,44,83,5',
              '31,253,87,309,5',
              '257,305,341,389,6',
              '317,163,401,247,4',
              '164,118,248,202,0']

4、获取图片绝对路径

image_path = annotation[0]

5、获取图片名字

image_name = image_path.split('/')[-1]

6、读取图片并将其从BGR转成RGB格式

image = cv2.imread(image_path)
 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

7、将图片的真实框坐标以及类别信息转化为整数

bbox_data_gt = np.array([list(map(int, box.split(','))) for box in annotation[1:]])

if len(bbox_data_gt) == 0:
      bboxes_gt=[]
      classes_gt=[]
else:
      bboxes_gt, classes_gt = bbox_data_gt[:, :4], bbox_data_gt[:, 4]#提取4个坐标信息和类别信息

8、将图片的真实框信息写入文件

        print('=> ground truth of %s:' % image_name)
        num_bbox_gt = len(bboxes_gt)
        with open(ground_truth_path, 'w') as f:#打开真实框文件
            for i in range(num_bbox_gt):#遍历真实框信息
                class_name = CLASSES[classes_gt[i]]#获取真实框类别
                xmin, ymin, xmax, ymax = list(map(str, bboxes_gt[i]))#获取真实框坐标
                bbox_mess = ' '.join([class_name, xmin, ymin, xmax, ymax]) + '\n'
                f.write(bbox_mess)#写入文件
                print('\t' + str(bbox_mess).strip())

此时第200张图片对应的199.txt中的内容是：

7 141 305 197 361
0 87 23 143 79
1 46 114 102 170
5 16 55 44 83
5 31 253 87 309
6 257 305 341 389
4 317 163 401 247
0 164 118 248 202

9、生成预测框的文件路径信息

print('=> predict result of %s:' % image_name)
predict_result_path = os.path.join(predicted_dir_path, str(num) + '.txt')

10、通过yolov3网络进行预测

        image_size = image.shape[:2]#获取原图的大小
        image_data = utils.image_preporcess(np.copy(image), [INPUT_SIZE, INPUT_SIZE])#将原图大小处理为yolov3需要的图片尺寸，这其中的操作包括放大缩小、填充操作、归一化等。具体流程可以参见utils文件夹中的image_preporcess函数
        image_data = image_data[np.newaxis, ...].astype(np.float32)#对图片添加一个维度，未处理前image_data的shape=[416,416,3]，处理后image_data的shape=[1,416,416,3]

        pred_bbox = model.predict(image_data)#开始进行预测，输出3个尺度的特征图[1,52,52,3,85]，[1,26,26,3,85]，[1,13,13,3,85]
        #下面2步是对3个尺度的特征图进行reshape和concat操作，最终pred_bbox的shape=[1*52*52*3+1*26*26*3+1*13*13*3，85]
        pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
        pred_bbox = tf.concat(pred_bbox, axis=0)
        #把pred_bbox送入postprocess_boxes进行一些过滤操作，排除一些在原图之外的预测框，面积=0的预测框，得分<cfg.TEST.SCORE_THRESHOLD的预测框
        bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE, cfg.TEST.SCORE_THRESHOLD)
        #经过postprocess_boxes处理过后，基本上余下的预测框都是合格的，但还有一个问题，就是可能有多个预测框都对应的是同一个物体，按照要求，一个物体只需要一个预测框就够了，所以这个nms函数就是来干这件事的，具体的代码流程参考utils文件夹的nms函数。
        bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms')

11、在原图上把预测框绘制出来，并写入文件

     if cfg.TEST.DECTECTED_IMAGE_PATH is not None:
            image = utils.draw_bbox(image, bboxes)#绘制预测框
            cv2.imwrite(cfg.TEST.DECTECTED_IMAGE_PATH+image_name, image)#将绘有框的原图写入文件

12、将预测框的信息写入文件

with open(predict_result_path, 'w') as f:#打开预测框文件
            for bbox in bboxes:#遍历预测框
                coor = np.array(bbox[:4], dtype=np.int32)#获取预测框坐标
                score = bbox[4]#获取预测框分数
                class_ind = int(bbox[5])#获取预测框类别
                class_name = CLASSES[class_ind]#获取预测框的类别名称
                score = '%.4f' % score
                xmin, ymin, xmax, ymax = list(map(str, coor))
                bbox_mess = ' '.join([class_name, score, xmin, ymin, xmax, ymax]) + '\n'
                f.write(bbox_mess)#写入文件
                print('\t' + str(bbox_mess).strip())

以上就是yolov3中测试代码的全部流程。