代码主要是参考bubbliiing的github YOLOv3的代码：github.com/bubbliiiing…
对于源代码的解读
预测部分
yolo.py文件
根据上面的预测的代码的部分我们，追踪到原始文件YOLO中的detect_image方法
总结：在预测中最最核心的代码
outputs = self.net(images) outputs = self.bbox_util.decode_box(outputs) #---------------------------------------------------------# # 将预测框进行堆叠，然后进行非极大抑制 #---------------------------------------------------------# results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)预测过程
首先的输入images[1,3,416,416],进入我们网络中输出一个list,这个list中有3个不同尺度的输出，第一个是[1,255,13,13],第二个是[1,255,26,26],第三个是[1,255,52,52],实际上就是我们的输出每个格子的3个框的每个框的85个预测值，13 * 13 * 3 * 85更好理解一点，然后我们使用decode_box对我们的每种分辨率的特征进行解码得到修正后的归一化坐标,并将维度进行修改成[1,507,85].然后使用torch.cat方法将三个分辨率的预测结果放到一起[1,10647,85]，然后我们对这些预测框进行非极大值抑制，来得到真正的预测框，并通过yolo_correct_boxes，将归一化坐标转化成真正的图片上真实的坐标信息，最终得到的结果是[框数,数据维度]，数据维度包括[ymin,xmin,ymax,xmax,conf,class_pro,class_label_num].
'''
训练自己的数据集必看注释！
'''
class YOLO(object):
_defaults = {
#--------------------------------------------------------------------------#
# 使用自己训练好的模型进行预测一定要修改model_path和classes_path！
# model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
#
# 训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
# 验证集损失较低不代表mAP较高，仅代表该权值在验证集上泛化性能较好。
# 如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
#--------------------------------------------------------------------------#
"model_path" : 'model_data/yolo_weights.pth',
"classes_path" : 'model_data/coco_classes.txt',
#---------------------------------------------------------------------#
# anchors_path代表先验框对应的txt文件，一般不修改。
# anchors_mask用于帮助代码找到对应的先验框，一般不修改。
#---------------------------------------------------------------------#
"anchors_path" : 'model_data/yolo_anchors.txt',
"anchors_mask" : [[6, 7, 8], [3, 4, 5], [0, 1, 2]],
#---------------------------------------------------------------------#
# 输入图片的大小，必须为32的倍数。
#---------------------------------------------------------------------#
"input_shape" : [416, 416],
#---------------------------------------------------------------------#
# 只有得分大于置信度的预测框会被保留下来
#---------------------------------------------------------------------#
"confidence" : 0.5,
#---------------------------------------------------------------------#
# 非极大抑制所用到的nms_iou大小
#---------------------------------------------------------------------#
"nms_iou" : 0.3,
#---------------------------------------------------------------------#
# 该变量用于控制是否使用letterbox_image对输入图像进行不失真的resize，
# 在多次测试后，发现关闭letterbox_image直接resize的效果更好
#---------------------------------------------------------------------#
"letterbox_image" : False,
#-------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
#-------------------------------#
"cuda" : False,
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化YOLO
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
for name, value in kwargs.items():
setattr(self, name, value)
#---------------------------------------------------#
# 获得种类和先验框的数量
#---------------------------------------------------#
self.class_names, self.num_classes = get_classes(self.classes_path)
self.anchors, self.num_anchors = get_anchors(self.anchors_path)
self.bbox_util = DecodeBox(self.anchors, self.num_classes, (self.input_shape[0], self.input_shape[1]), self.anchors_mask)
#---------------------------------------------------#
# 画框设置不同的颜色
#---------------------------------------------------#
hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
self.generate()
#---------------------------------------------------#
# 生成模型
#---------------------------------------------------#
def generate(self):
#---------------------------------------------------#
# 建立yolov3模型，载入yolov3模型的权重
#---------------------------------------------------#
self.net = YoloBody(self.anchors_mask, self.num_classes)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.net.load_state_dict(torch.load(self.model_path, map_location=device))
self.net = self.net.eval()
print('{} model, anchors, and classes loaded.'.format(self.model_path))
if self.cuda:
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image, crop = False):
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
image = cvtColor(image)
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
# 也可以直接resize进行识别
#---------------------------------------------------------#
image_data = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
#---------------------------------------------------------#
# 添加上batch_size维度
#---------------------------------------------------------#
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
with torch.no_grad():
images = torch.from_numpy(image_data)
if self.cuda:
images = images.cuda()
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
outputs = self.net(images)
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,
image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
if results[0] is None:
return image
top_label = np.array(results[0][:, 6], dtype = 'int32')
top_conf = results[0][:, 4] * results[0][:, 5]
top_boxes = results[0][:, :4]
#---------------------------------------------------------#
# 设置字体与边框厚度
#---------------------------------------------------------#
font = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
thickness = int(max((image.size[0] + image.size[1]) // np.mean(self.input_shape), 1))
#---------------------------------------------------------#
# 是否进行目标的裁剪
#---------------------------------------------------------#
if crop:
for i, c in list(enumerate(top_label)):
top, left, bottom, right = top_boxes[i]
top = max(0, np.floor(top).astype('int32'))
left = max(0, np.floor(left).astype('int32'))
bottom = min(image.size[1], np.floor(bottom).astype('int32'))
right = min(image.size[0], np.floor(right).astype('int32'))
dir_save_path = "img_crop"
if not os.path.exists(dir_save_path):
os.makedirs(dir_save_path)
crop_image = image.crop([left, top, right, bottom])
crop_image.save(os.path.join(dir_save_path, "crop_" + str(i) + ".png"), quality=95, subsampling=0)
print("save crop_" + str(i) + ".png to " + dir_save_path)
#---------------------------------------------------------#
# 图像绘制
#---------------------------------------------------------#
for i, c in list(enumerate(top_label)):
predicted_class = self.class_names[int(c)]
box = top_boxes[i]
score = top_conf[i]
top, left, bottom, right = box
top = max(0, np.floor(top).astype('int32'))
left = max(0, np.floor(left).astype('int32'))
bottom = min(image.size[1], np.floor(bottom).astype('int32'))
right = min(image.size[0], np.floor(right).astype('int32'))
label = '{} {:.2f}'.format(predicted_class, score)
draw = ImageDraw.Draw(image)
label_size = draw.textsize(label, font)
label = label.encode('utf-8')
print(label, top, left, bottom, right)
if top - label_size[1] >= 0:
text_origin = np.array([left, top - label_size[1]])
else:
text_origin = np.array([left, top + 1])
for i in range(thickness):
draw.rectangle([left + i, top + i, right - i, bottom - i], outline=self.colors[c])
draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=self.colors[c])
draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
del draw
return image
def get_FPS(self, image, test_interval):
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
image = cvtColor(image)
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
# 也可以直接resize进行识别
#---------------------------------------------------------#
image_data = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
#---------------------------------------------------------#
# 添加上batch_size维度
#---------------------------------------------------------#
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
with torch.no_grad():
images = torch.from_numpy(image_data)
if self.cuda:
images = images.cuda()
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
outputs = self.net(images)
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,
image_shape, self.letterbox_image, conf_thres=self.confidence, nms_thres=self.nms_iou)
t1 = time.time()
for _ in range(test_interval):
with torch.no_grad():
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
outputs = self.net(images)
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,
image_shape, self.letterbox_image, conf_thres=self.confidence, nms_thres=self.nms_iou)
t2 = time.time()
tact_time = (t2 - t1) / test_interval
return tact_time
def get_map_txt(self, image_id, image, class_names, map_out_path):
f = open(os.path.join(map_out_path, "detection-results/"+image_id+".txt"),"w")
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
image = cvtColor(image)
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
# 也可以直接resize进行识别
#---------------------------------------------------------#
image_data = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
#---------------------------------------------------------#
# 添加上batch_size维度
#---------------------------------------------------------#
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
with torch.no_grad():
images = torch.from_numpy(image_data)
if self.cuda:
images = images.cuda()
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
outputs = self.net(images)
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,
image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
if results[0] is None:
return
top_label = np.array(results[0][:, 6], dtype = 'int32')
top_conf = results[0][:, 4] * results[0][:, 5]
top_boxes = results[0][:, :4]
for i, c in list(enumerate(top_label)):
predicted_class = self.class_names[int(c)]
box = top_boxes[i]
score = str(top_conf[i])
top, left, bottom, right = box
if predicted_class not in class_names:
continue
f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))
f.close()
return
针对与上面的代码进行详细的解读：
_default
我们可修改的模型加载参数从这个部分进行修改，从指定包括
- 可修改
模型参数文件物体类别文件图片resize的大小(32倍数)置信度分数nms阈值是否进行不失真的resize是否使用GPU- 不修改
先验框宽高文件先验框索引文件
_defaults = {
#--------------------------------------------------------------------------#
# 使用自己训练好的模型进行预测一定要修改model_path和classes_path！
# model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
#
# 训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
# 验证集损失较低不代表mAP较高，仅代表该权值在验证集上泛化性能较好。
# 如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
#--------------------------------------------------------------------------#
"model_path" : 'model_data/yolo_weights.pth',
"classes_path" : 'model_data/coco_classes.txt',
#---------------------------------------------------------------------#
# anchors_path代表先验框对应的txt文件，一般不修改。
# anchors_mask用于帮助代码找到对应的先验框，一般不修改。
#---------------------------------------------------------------------#
"anchors_path" : 'model_data/yolo_anchors.txt',
"anchors_mask" : [[6, 7, 8], [3, 4, 5], [0, 1, 2]],
#---------------------------------------------------------------------#
# 输入图片的大小，必须为32的倍数。
#---------------------------------------------------------------------#
"input_shape" : [416, 416],
#---------------------------------------------------------------------#
# 只有得分大于置信度的预测框会被保留下来
#---------------------------------------------------------------------#
"confidence" : 0.5,
#---------------------------------------------------------------------#
# 非极大抑制所用到的nms_iou大小
#---------------------------------------------------------------------#
"nms_iou" : 0.3,
#---------------------------------------------------------------------#
# 该变量用于控制是否使用letterbox_image对输入图像进行不失真的resize，
# 在多次测试后，发现关闭letterbox_image直接resize的效果更好
#---------------------------------------------------------------------#
"letterbox_image" : False,
#-------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
#-------------------------------#
"cuda" : False,
}
初始化
比较常规的代码：直接下载下面的注释中
# 读取上面的_defaults的信息，将其作为类属性存储
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
# 初始化
#---------------------------------------------------#
# 初始化YOLO
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
# 将实例化时候的参数也加入我们的属性字典中
for name, value in kwargs.items():
setattr(self, name, value)
#---------------------------------------------------#
# 获得种类和先验框的数量
#---------------------------------------------------#
# 这个部分和我们加载过的classes_path有关，pascal voc是20类，coco是80类
# self.class_names = [`person`,`car`,....] 类别列表
# self.num_classes = 80
self.class_names, self.num_classes = get_classes(self.classes_path)
# self.anchors = array([[10,13],[16,30],....]) 注意这个地方还是框的真实的宽高
# self.num_anchors = 9
self.anchors, self.num_anchors = get_anchors(self.anchors_path)
# 初始化一个decodebox对象，为了后面对box解析使用
self.bbox_util = DecodeBox(self.anchors, self.num_classes, (self.input_shape[0], self.input_shape[1]), self.anchors_mask)
#---------------------------------------------------#
# 画框设置不同的颜色
#---------------------------------------------------#
# 这个地方是为每一个类生成一个不同颜色的框，用于后面展示使用(不做深究)
hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
# 调用generate生成模型
self.generate()
#---------------------------------------------------#
# 生成模型
#---------------------------------------------------#
def generate(self):
#---------------------------------------------------#
# 建立yolov3模型，载入yolov3模型的权重
#---------------------------------------------------#
# 初始化我们的网络
self.net = YoloBody(self.anchors_mask, self.num_classes)
# 选择设备cpu or gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 加载模型预训练的参数，并放到指定devices上
self.net.load_state_dict(torch.load(self.model_path, map_location=device))
self.net = self.net.eval()
print('{} model, anchors, and classes loaded.'.format(self.model_path))
# 使用GPU的话
if self.cuda:
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
比较简单，不解读
- 追踪函数1
get_classes
# 获取类别列表、和类别数
def get_classes(classes_path):
with open(classes_path, encoding='utf-8') as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names, len(class_names)
- 追踪函数2
get_classes
# 获取框宽高array和框的数量
def get_anchors(anchors_path):
'''loads the anchors from a file'''
with open(anchors_path, encoding='utf-8') as f:
anchors = f.readline()
anchors = [float(x) for x in anchors.split(',')]
anchors = np.array(anchors).reshape(-1, 2)
return anchors, len(anchors)
- 追踪类1
DecodeBox
对于DecodeBox的解读：juejin.cn/post/707891… 详细的分析都放在这个部分
- 追踪类2
YoloBody
对于YoloBody的解读：juejin.cn/post/707892… 详细的分析都放在这个部分
- 直接对detect_image进行解读
针对于不失真的图片的blog的解读：blog.csdn.net/weixin_4479…
def detect_image(self, image, crop = False):
# 获取图片的宽高(用于最后的真实框的宽高计算)
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
# 对图片进行色彩转化(将图像转换成RGB图像，防止灰度图在预测时报错,代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB)
image = cvtColor(image)
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
# 也可以直接resize进行识别
#---------------------------------------------------------#
# 进行重置大小，如果letterbox_image是True的话，就会进行不失真的resize
image_data = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
#---------------------------------------------------------#
# 添加上batch_size维度
#---------------------------------------------------------#
# 因为我们需要将图片输入到网络中，我们是需要batch_size的这一个维度的
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
# np.array(image_data, dtype='float32') 首先将图片变成ndarray的格式，类型变成float32
# preprocess_input(np.array(image_data, dtype='float32')) 将图片数据进行归一化
# np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)) # 通道数提前，pytorch是 NHW 图片是HWN
# np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0) # 在前面扩展出来一维batch
# 开启计算
with torch.no_grad():
# 将图片从ndarray的格式变成tensor
images = torch.from_numpy(image_data)
# 使用gpu的话，将数据放置到gpu上
if self.cuda:
images = images.cuda()
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
# 使用网络对我们的图片进行预测 这个部分我们将得到[[batch_size,13,13,3 * (5+num_class)],[batch_size,26,26,3 * (5+num_class)],[batch_size,52,52,3 * (5+num_class)]]
outputs = self.net(images)
# 我们对输出出来的数据进行decode，将[batch_size,13,13,3 * (5+num_class)]-->[batch_size,3 * 13 * 13,(5+num_class)] = [batch_size,507,(5+num_class)]
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制，最终我们得到的是results--->[[537.0076,77.90824,953.0772,270.35666,0.99966,0.999975,0],....]
# 分别是y1,x1,y2,x2,conf,class_pro,class_label 七个值
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
# 假设经过极大值抑制之后，我们得到的数据中已经没有了信息(就是什么都没有搜索到)，直接返回图片即可
if results[0] is None:
return image
# 分别得到label、conf置信度、boxes的左上和右下点的坐标
top_label = np.array(results[0][:, 6], dtype = 'int32')
top_conf = results[0][:, 4] * results[0][:, 5]
top_boxes = results[0][:, :4]
#---------------------------------------------------------#
# 设置字体与边框厚度
#---------------------------------------------------------#
font = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
thickness = int(max((image.size[0] + image.size[1]) // np.mean(self.input_shape), 1))
#---------------------------------------------------------#
# 是否进行目标的裁剪
#---------------------------------------------------------#
# 下面的代码就是看我们是不是有crop参数，有的话，我们就在图片中从上面得到boxes的位置信息中循环进行裁剪，并将裁剪后的图片保存在指定文件夹中
if crop:
for i, c in list(enumerate(top_label)):
top, left, bottom, right = top_boxes[i]
top = max(0, np.floor(top).astype('int32'))
left = max(0, np.floor(left).astype('int32'))
bottom = min(image.size[1], np.floor(bottom).astype('int32'))
right = min(image.size[0], np.floor(right).astype('int32'))
dir_save_path = "img_crop"
if not os.path.exists(dir_save_path):
os.makedirs(dir_save_path)
crop_image = image.crop([left, top, right, bottom])
crop_image.save(os.path.join(dir_save_path, "crop_" + str(i) + ".png"), quality=95, subsampling=0)
print("save crop_" + str(i) + ".png to " + dir_save_path)
#---------------------------------------------------------#
# 图像绘制
#---------------------------------------------------------#
# 我们通过对label的循环，可以分别拿到conf和box
for i, c in list(enumerate(top_label)):
predicted_class = self.class_names[int(c)]
box = top_boxes[i]
score = top_conf[i]
# 哈哈，注意这个地方我们的[y1,x1,y2,x2]恰好对应top,left,bottom,right的次序方向
top, left, bottom, right = box
# 获取左上、右下的坐标
top = max(0, np.floor(top).astype('int32'))
left = max(0, np.floor(left).astype('int32'))
bottom = min(image.size[1], np.floor(bottom).astype('int32'))
right = min(image.size[0], np.floor(right).astype('int32'))
label = '{} {:.2f}'.format(predicted_class, score)
# 在图片上绘制
draw = ImageDraw.Draw(image)
label_size = draw.textsize(label, font)
label = label.encode('utf-8')
print(label, top, left, bottom, right)
# 假设我们的框没有贴到了图片的上边缘，那我们文字的起始位置就是从top向上一个label的距离
# 假设我们的框贴到了图片的上边缘，那我们文字的起始位置就是从top + 1的位置，放在方框的内部
if top - label_size[1] >= 0:
text_origin = np.array([left, top - label_size[1]])
else:
text_origin = np.array([left, top + 1])
# 指定的相框的厚度、我们起始是一个像素一个像素的向外进行绘制直线，最终形成一个指定厚度的矩形框
for i in range(thickness):
draw.rectangle([left + i, top + i, right - i, bottom - i], outline=self.colors[c])
# 绘制文字框
draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=self.colors[c])
# 写上文字
draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
# 删除绘画体对象
del draw
# 最后将我们绘制好的图片返回回去
return image
- 调用其他代码
- 解读cvtColor
def cvtColor(image): # 图片的shape是3且第三维的维度是3，可以确定是RGB的图片 if len(np.shape(image)) == 3 and np.shape(image)[2] == 3: return image else: # 非RGB图片就是灰度图片，需要转化成RGB,才能进行预测 image = image.convert('RGB') return image- 解读resize_image
# 不失真完成图片的resize def resize_image(image, size, letterbox_image): # 传入图片的宽高 iw, ih = image.size # 要生成的图片的大小 w, h = size # 是否需要进行填充 if letterbox_image: # 这个地方是我们需要将图片扩展的倍数，之所以取最小值是因为，假设我们取其中较大的值，则当这个值到达指定长度的数，另外一个已经超出了边界 # iw = 240 ih = 411 w=h=416,ih只需要1.01倍即可达到416，而iw需要1.73倍，所以当iw平铺达到最大的时候，ih已经超过了416 # 所以选择小的倍数，让长宽等比放达到其中一个先达到最大值 scale = min(w/iw, h/ih) nw = int(iw*scale) nh = int(ih*scale) # 我们使用Image.BICUBIC的方式进行插值图片放大 image = image.resize((nw,nh), Image.BICUBIC) # 构造一张新图片mask，大小是指定大小 new_image = Image.new('RGB', size, (128,128,128)) # 将我们的image图片复制到mask中，指定左上角位置 new_image.paste(image, ((w-nw)//2, (h-nh)//2)) else: # 不用填充，图片看起来会变形 new_image = image.resize((w, h), Image.BICUBIC) return new_image- 解读preprocess_input
# 将图片像素进行归一化 def preprocess_input(image): image /= 255.0 return image- 解读DecodeBox.decode_box、non_max_suppression
- 对于DecodeBox的解读：juejin.cn/post/707891… 详细的分析都放在这个部分
对于输入的数据的理解
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
很重要
outputs是一个len为3的列表，这三个尺度13 * 13 个格子每个格子3个框，一共507个框，同理我们的26 * 26 * 3=2028个框，52 * 52 * 3 = 8112个框
经过了torch.cat(outputs,1) 进行第一个维度的相加之后变成了[1,10647,85],也就是说我们一共预测出来了10647个框，其中后面是位置信息和置信度和种类信息
top_label展示
top_conf的展示
top_boxes
- 对
get_FPS进行解读(非重要，可跳过)- 主要功能：得到单位时间可以完成多少张图片的预测。test_interval的值越大越准
def get_FPS(self, image, test_interval):
# 这个部分和detect_image的前半部分一致，不进行解读
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
image = cvtColor(image)
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
# 也可以直接resize进行识别
#---------------------------------------------------------#
image_data = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
#---------------------------------------------------------#
# 添加上batch_size维度
#---------------------------------------------------------#
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
with torch.no_grad():
images = torch.from_numpy(image_data)
if self.cuda:
images = images.cuda()
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
outputs = self.net(images)
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,image_shape, self.letterbox_image, conf_thres=self.confidence, nms_thres=self.nms_iou)
# 主要是从这个地方开始不同，我们进行时间统计，在前面指定的参数test_interval=100，进行100轮的预测-box解码-非极大值抑制过程，最后得到运行总时间，去除测试的轮次，可以得到单位时间可以完成多少张图片的预测
t1 = time.time()
for _ in range(test_interval):
# 循环（预测-box解码-非极大值抑制）
with torch.no_grad():
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
outputs = self.net(images)
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, image_shape, self.letterbox_image, conf_thres=self.confidence, nms_thres=self.nms_iou)
t2 = time.time()
# 得到单位时间可以完成多少张图片的预测
tact_time = (t2 - t1) / test_interval
return tact_time
- 对
get_map_txt的解读- 获得每个图片的多个
预测种类、置信度、左、上、右、下的信息存储到我们的上面的信息文件中，主要是为了得到map的计算
- 获得每个图片的多个
def get_map_txt(self, image_id, image, class_names, map_out_path):
# 打开一个文件用于记录数据
f = open(os.path.join(map_out_path, "detection-results/"+image_id+".txt"),"w")
# 这个部分和detect_image的前半部分一致，不进行解读
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
image = cvtColor(image)
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
# 也可以直接resize进行识别
#---------------------------------------------------------#
image_data = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
#---------------------------------------------------------#
# 添加上batch_size维度
#---------------------------------------------------------#
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
with torch.no_grad():
images = torch.from_numpy(image_data)
if self.cuda:
images = images.cuda()
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
outputs = self.net(images)
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
if results[0] is None:
return
top_label = np.array(results[0][:, 6], dtype = 'int32')
top_conf = results[0][:, 4] * results[0][:, 5]
top_boxes = results[0][:, :4]
for i, c in list(enumerate(top_label)):
predicted_class = self.class_names[int(c)]
box = top_boxes[i]
score = str(top_conf[i])
top, left, bottom, right = box
if predicted_class not in class_names:
continue
# 我们将上面得到的预测种类、置信度、左、上、右、下的信息存储到我们的上面的信息文件中，主要是为了得到map的计算
f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))
f.close()
return
