【目标检测】YOLOv3代码实现之Pytorch
代码目录
- 一、资源下载
- 1.1 代码下载
- 1.2 权重下载
- 二、源码解析
- 2.1 config文件夹
- 2.1.1 coco.data
- 2.1.2 create_custom_model.sh
- 2.1.3 custom.data
- 2.1.4 yolov3.cfg
- 2.1.5 yolov3-custom.cfg
- 2.1.6 yolov3-tiny.cfg
- 2.2 data文件夹
- 2.2.1 coco文件夹
- 2.2.2 custom文件夹
- 2.2.3 samples文件夹
- 2.2.4 coco.names
- 2.2.5 get_coco_dataset.sh
- 2.3 utils文件夹
- 2.3.1 augmentations.py
- 2.3.2 datasets.py
- 2.3.3 logger.py
- 2.3.4 parse_config.py
- 2.3.5 utils.py
- 2.4 weights文件夹
- 2.5 detect.py
- 2.6 models.py
- 2.7 test.py
- 2.8 train.py
一、资源下载
1.1 代码下载
项目代码：https://github.com/eriklindernoren/PyTorch-YOLOv3?from=singlemessage
整个项目的结构如下图：
1.2 权重下载
预训练权重：https://pjreddie.com/media/files/yolov3.weights
将下载的权重放入weight文件夹,如下图：
二、源码解析
2.1 config文件夹
2.1.1 coco.data
coco数据集的信息：类别数量，训练集路径、验证集路径、类别名称路径…
2.1.2 create_custom_model.sh
脚本文件：用户自定义自己的模型，运行此文件用来生成自定义模型的配置文件yolov3-custom.cfg。可对比yolov3.cfg
2.1.3 custom.data
自己数据集的信息,用来训练自己的检测任务：类别数量，训练集路径、验证集路径、类别名称路径…。可对比coco.data。
2.1.4 yolov3.cfg
yolov3网络模型的配置信息：卷积层（卷积核数、卷积核尺寸、步长…）、yolo层及其他层的配置信息。
2.1.5 yolov3-custom.cfg
自定义的网络模型的配置信息，由create_custom_model.sh脚本文件生成。
2.1.6 yolov3-tiny.cfg
yolov3的tiny版本网络模型的配置信息。
2.2 data文件夹
2.2.1 coco文件夹
是coco训练集、验证集的数据集，是运行get_coco_dataset.sh脚本文件后的结果。
2.2.2 custom文件夹
是自定义数据集的信息。
1）images文件夹：所有训练集、验证集的图片。一部分图片如图
2）labels文件夹：使用图片标记软件对images文件夹里的图片进行标注得到对应的标签文件。一部分标签文件如图
每个标签文件为一个txt文件，txt文件的每一行数据为一个groundthuth信息：类别序号，边界框坐标信息。如图示例，0代表类别索引号，后面为边界框坐标信息。
3）classes.names是自定义数据集的类别名称文件。例
4）train.txt文件是训练集图片路径的集合，如图，每行数据是训练集某图像的路径。
5）valid.txt文件是验证集图片路径的集合，如图，每行数据是训练集某图片的路径。
2.2.3 samples文件夹
是模型测试图片所在的文件夹，用来看模型的检测结果。
2.2.4 coco.names
coco数据的类别信息，类似classes.names。如图部分截图
2.2.5 get_coco_dataset.sh
脚本文件，用来获取coco数据，生成coco文件夹及其内容。
2.3 utils文件夹
2.3.1 augmentations.py
进行数据增强的文件，本项目只是进行水平翻转的数据增强，图像进行翻转的时候，对应标注信息也进行了修改，最终返回的是翻转后的图片和翻转后的图片对应的标签。
import torch
def horisontal_flip(images, targets):
images = torch.flip(images, [-1])
#target=[n,索引，x,y,w,h]
targets[:, 2] = 1 - targets[:, 2]#图像aug时，标签target 也要同时进行改变，yolo的标签是比例关系，水平翻转图像，标签也要跟着改变，这里只需要改变box 中心点相对于图像左边界的距离。
return images, targets
2.3.2 datasets.py
对数据集进行操作的py文件，包含图像的填充、图像大小的调整、测试数据集的加载类、评估数据集的加载类。整个文件包含3个函数和2个类,如下
import glob
import random
import os
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as F
from utils.augmentations import horisontal_flip
from torch.utils.data import Dataset
import torchvision.transforms as transforms
'''图片填充函数：
把图片用pad_value填充成一个正方形，返回填充后的图片以及填充的位置信息 pad = (0, 0, pad1, pad2) if h <= w else (pad1, pad2, 0, 0)'''
def pad_to_square(img, pad_value):
c, h, w = img.shape#获取图片形状
dim_diff = np.abs(h - w)#图片长和宽的差值
# (upper / left) padding and (lower / right) padding
pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
# 填充方式，如果高小于宽则上下填充，如果高大于宽，左右填充
pad = (0, 0, pad1, pad2) if h <= w else (pad1, pad2, 0, 0)
# 图片填充，参数img是原图，pad是填充方式（0,0,pad1,pad2）或（pad1，pad2,0，0），value是填充的值
img = F.pad(img, pad, "constant", value=pad_value)
return img, pad#返回填充后的图片，和pad(填充的方式，这是为了方便以后还原原图尺寸)
'''图片调整大小：将正方形图片使用插值方法，改变到固定size大小'''
def resize(image, size):
image = F.interpolate(image.unsqueeze(0), size=size, mode="nearest").squeeze(0)
return image
"""随机裁剪函数：将图片随机裁剪到某个尺寸（使用插值法）"""
def random_resize(images, min_size=288, max_size=448):
new_size = random.sample(list(range(min_size, max_size + 1, 32)), 1)[0]
images = F.interpolate(images, size=new_size, mode="nearest")
return images
'''数据集加载类1：加载并处理图片，返回的是图片路径，和经过处理后的图片'''
#用于预测：在detect.py中加载数据集时使用
class ImageFolder(Dataset):
def __init__(self, folder_path, img_size=416):#初始化的参数为：测试图片所在的文件夹的路径、图片的尺寸（用于输入到网络的图片的大小）
#获取文件夹下图片的路径，files是图片路径组成的列表
self.files = sorted(glob.glob("%s/*.*" % folder_path))#例在detect.py中folder_path=data/samples
#print(type(self.files))#2.3.3 logger.py
用来将监控数据写入文件系统（日志），保存训练的某些信息。如损失等。这个logger类在train.py中使用，在训练过程中保存一些信息到日志文件。
import tensorflow as tf
class Logger(object):
def __init__(self, log_dir):
"""创建监控对象，记录到log_dir文件夹下."""
self.writer = tf.summary.FileWriter(log_dir)
def scalar_summary(self, tag, value, step):#将监控数据写入日志
"""Log a scalar variable."""
summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
self.writer.add_summary(summary, step)
def list_of_scalars_summary(self, tag_value_pairs, step):#将监控数据批量写入日志
"""Log scalar variables."""
summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value) for tag, value in tag_value_pairs])
self.writer.add_summary(summary, step)
2.3.4 parse_config.py
包含两个解析器：
1.模型配置解析器：返回一个列表model_defs，列表的每一个元素为一个字典，字典代表模型某一个层（模块）的信息 。
2.数据配置解析器：返回一个字典，每一个键值对描述了，数据的名称路径，或其他信息。
'''模型配置解析器：解析yolo-v3层配置文件函数，并返回模块定义module_defs，path就是yolov3.cfg路径'''
def parse_model_config(path):
'''
看此函数，一定要先看config文件夹下的yolov3.cfg文件，如下是yolov3。cfg的一部分内容展示：
[convolutional]
batch_normalize=1
filters=32
size=3
stride=1
pad=1
activation=leaky
# Downsample
[convolutional]
batch_normalize=1
filters=64
size=3
stride=2
pad=1
activation=leaky
。。。
:param path: 模型配置文件路径，yolov3.cfg的路径
:return: 模型定义，列表类型，列表中的元素是字典，字典包含了每一个模块的定义参数
'''
# 打开yolov3.cfg文件,并将文件内容存入列表，列表的每一个元素为文件的一行数据。
file = open(path, 'r')
lines = file.read().split('\n')
lines = [x for x in lines if x and not x.startswith('#')]#不读取注释
lines = [x.rstrip().lstrip() for x in lines] #去除边缘空白
#定义一个列表modle_defs
module_defs = []
#读取cfg的每一行内容：
# 1.如果该行内容以[开头:代表是模型的一个新块的开始，给module_defs列表新增一个字典
# 字典的‘type’=[]内的内容，如果[]内的内容是convolution,则字典添加'batch_normalize'：0
# 2.如果该行内容不以[开头，代表是块的具体内容
# 等号前的值为字典的key，等号后的值为字典的value
for line in lines:#读取yolov3.cfg文件的每一行
#如果一行内容以[开头说明是一个模型的开始,[]里的内容是模块的名称，如[convolutional][convolutional][shortcut]。。。。
if line.startswith('['): # This marks the start of a new block
# 将一个空字典添加到模型定义module_defs列表中
module_defs.append({})
# 给该字典内容赋值：例{’type‘：’convolutional‘}
module_defs[-1]['type'] = line[1:-1].rstrip()
# 如果当前的模块是convolutional模块，给字典的内容赋值：{’type‘：’convolutional‘，'batch_normalize'：0}
if module_defs[-1]['type'] == 'convolutional':
module_defs[-1]['batch_normalize'] = 0
#如果一行内容不以[开头说明是模块里的具体内容
else:
key, value = line.split("=")
value = value.strip()#strip()删除头尾空格，rstrip()删除结尾空格
# 将该行内容添加到字典中，key为等式左边的内容，value为等式右边的内容
module_defs[-1][key.rstrip()] = value.strip()
return module_defs#模型定义，是一个列表，列表每一个元素为一个字典，字典包含一个模块的具体信息
'''数据配置解析器：参数path为配置文件的路径'''
def parse_data_config(path):
"""
数据配置包含的信息：
classes= 80
train=data/coco/trainvalno5k.txt
valid=data/coco/5k.txt
names=data/coco.names
backup=backup/
eval=coco
"""
#创建一个字典
options = dict()
#为字典添加元素
options['gpus'] = '0,1,2,3'
options['num_workers'] = '10'
#读取数据配置文件的每一行，并将每一行的信息以键值对的形式存入字典中
with open(path, 'r') as fp:
lines = fp.readlines()
for line in lines:
line = line.strip()
if line == '' or line.startswith('#'):
continue
key, value = line.split('=')
options[key.strip()] = value.strip()
return options#返回一个字典，字典的key为名称（train，valid,names..），value为路径或其他信息
2.3.5 utils.py
是项目的工具文件
from __future__ import division
import tqdm
import torch
import numpy as np
def to_cpu(tensor):
return tensor.detach().cpu()
'''加载数据集类别信息：返回类别组成的列表'''
def load_classes(path):#参数为类别名称文件的路径。例coco.names或classes.names的路径
fp = open(path, "r")
names = fp.read().split("\n")[:-1]#将文件的每一行数据存入列表，这使得数据集的每个类别的名称存入到一个列表
return names#返回类别名称构成的列表
'''权重初始化函数'''
def weights_init_normal(m):
classname = m.__class__.__name__
if classname.find("Conv") != -1:#卷积层权重初始化设置
torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
elif classname.find("BatchNorm2d") != -1:#批量归一化层权重初始化设置
torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
torch.nn.init.constant_(m.bias.data, 0.0)
'''改变预测边界框的尺寸函数：参数为，边界框、当前的图片尺寸（标量）、原始图片尺寸。因为网络预测的边界框信息是，
对图像填充、调整大小后的图片进行预测的结果，因此需要对预测的边界框进行调整使其适应于原图的目标'''
def rescale_boxes(boxes, current_dim, original_shape):
#原始图片的高和宽
orig_h, orig_w = original_shape
#原始图片的填充信息：根据原图的宽高的差值来计算。
#pad_x为宽天长的像素数量， pad_y为高填充的像素数量
pad_x = max(orig_h - orig_w, 0) * (current_dim / max(original_shape))# 原图的高大于宽。改变后图片的大小/原图的最长边的尺寸=缩放比率
pad_y = max(orig_w - orig_h, 0) * (current_dim / max(original_shape))
#将预测的边界框信息，调整为适应于原图
unpad_h = current_dim - pad_y
unpad_w = current_dim - pad_x
# 改变预测边界框的尺寸，使其是适用于原图片
boxes[:, 0] = ((boxes[:, 0] - pad_x // 2) / unpad_w) * orig_w#左上x的坐标
boxes[:, 1] = ((boxes[:, 1] - pad_y // 2) / unpad_h) * orig_h#左上y的坐标
boxes[:, 2] = ((boxes[:, 2] - pad_x // 2) / unpad_w) * orig_w
boxes[:, 3] = ((boxes[:, 3] - pad_y // 2) / unpad_h) * orig_h
return boxes#返回调整后的预测边界框的信息/
'''将边界框信息转换为左上右下坐标表示函数'''
def xywh2xyxy(x):
y = x.new(x.shape)
y[..., 0] = x[..., 0] - x[..., 2] / 2
y[..., 1] = x[..., 1] - x[..., 3] / 2
y[..., 2] = x[..., 0] + x[..., 2] / 2
y[..., 3] = x[..., 1] + x[..., 3] / 2
return y
"""度量计算：参数为true_positive（值为0或1,list）、预测置信度(list)，预测类别(list)，真实类别(list)
返回：p, r, ap, f1, unique_classes.astype("int32")"""
def ap_per_class(tp, conf, pred_cls, target_cls):#参数：true_positives, pred_scores, pred_labels 、图片真实标签信息
# 按照置信度排序，后的tp, conf, pred_cls
i = np.argsort(-conf)
#print('所有预测框的个数为',len(i))
tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]#按照置信度排序后的tp(值为0,1), conf, pred_cls
#print('tp[i]',tp[i])
# 获取图片中真实框所包含的类别（类别不重复）
unique_classes = np.unique(target_cls)
#print('unique_classes',unique_classes)
# Create Precision-Recall curve and compute AP for each class
ap, p, r = [], [], []
for c in tqdm.tqdm(unique_classes, desc="Computing AP"):#为每一个类别计算AP
# i:对于所有预测边界框的类pred_cls，判断与当前c类是否相同，相同则该位置为true否则为false,得到与pred_class形状相同的布尔列表
i = pred_cls == c
# ground truth 中类别为c的数量
n_gt = (target_cls == c).sum()
#预测边界框中类别为c的数量
n_p = i.sum()
if n_p == 0 and n_gt == 0:
continue
elif n_p == 0 or n_gt == 0:
ap.append(0)
r.append(0)
p.append(0)
else:
# 计算FP和TP
fpc = (1 - tp[i]).cumsum()#i列表记录着索引对应位置是否是c类别的边界框，tp记录着索引对应位置是否是正例框
tpc = (tp[i]).cumsum()
# print('tp[i]',tp[i],len(tp[i]))#tp[i]是所有框中类别为c的预测框的true_positive信息（值为0或1，1代表与真值框iou大于阈值）
# print('fpc',fpc,len(fpc))#fpc为类别为c的预测框中为正例的预测框
# print('tpc', tpc,len(tpc))#tpc为类别为c的预测框中为负例的预测框
#计算召回率
recall_curve = tpc / (n_gt + 1e-16)
#print('recall_curve',recall_curve)
r.append(recall_curve[-1])
#print('r',r)
#计算精度
precision_curve = tpc / (tpc + fpc)
#print('precision_curve',precision_curve)
p.append(precision_curve[-1])
#print('p',p)
# 计算AP：AP from recall-precision curve
ap.append(compute_ap(recall_curve, precision_curve))
# Compute F1 score (harmonic mean of precision and recall)
p, r, ap = np.array(p), np.array(r), np.array(ap)
f1 = 2 * p * r / (p + r + 1e-16)
return p, r, ap, f1, unique_classes.astype("int32")
"""计算AP"""
def compute_ap(recall, precision):#参数精度和召回率
# correct AP calculation
# 给Precision-Recall曲线添加头尾
mrec = np.concatenate(([0.0], recall, [1.0]))
mpre = np.concatenate(([0.0], precision, [0.0]))
# compute the precision envelope
# 简单的应用了一下动态规划，实现在recall=x时，precision的数值是recall=[x, 1]范围内的最大precision
for i in range(mpre.size - 1, 0, -1):
mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
# to calculate area under PR curve, look for points
# where X axis (recall) changes value
# 寻找recall[i]!=recall[i+1]的所有位置，即recall发生改变的位置，方便计算PR曲线下的面积，即AP
i = np.where(mrec[1:] != mrec[:-1])[0]
# and sum (\Delta recall) * prec
# 用积分法求PR曲线下的面积，即AP
ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
return ap
'''统计信息计算：参数，模型预测输出（NMS处理后的结果），真实标签（适应于原图的x,y,x,y）,iou阈值。
返回，true_positive（值为0/1，如果预测边界框与真实边界框重叠度大则值为1，否则为0）,预测置信度，预测类别'''
def get_batch_statistics(outputs, targets, iou_threshold):
# outputs为非极大值抑制后的结果(x,y,x,y,object_confs,class_confs,class_preds)长度为7
batch_metrics = []
for sample_i in range(len(outputs)):#遍历每个output的边界框，因为是批量操作的，每个批量有很多图片，每个图片对应一个output,所以遍历每个output
if outputs[sample_i] is None:
continue
'''图片的预测信息：'''
output = outputs[sample_i]#取第sample_i个output信息，每个output里面包含很多边界框
pred_boxes = output[:, :4]#预测边界框的坐标信息
pred_scores = output[:, 4]#预测边界框的置信度
pred_labels = output[:, -1]#预测边界框的类别
true_positives = np.zeros(pred_boxes.shape[0])#true_positive的长度为pre_boxes的个数
'''图片的标注信息（groundtruth）：'''
#坐标信息，格式为（xyxy）
annotations = targets[targets[:, 0] == sample_i][:, 1:]#这句把对应ID下的target和图像进行匹配，dataset.py里的ListDataset类里的collate_fn函数给target赋予ID
#类别信息
target_labels = annotations[:, 0] if len(annotations) else []
if len(annotations):
detected_boxes = []#创建空列表
target_boxes = annotations[:, 1:]#真实边界框（groundtruth）坐标
for pred_i, (pred_box, pred_label) in enumerate(zip(pred_boxes, pred_labels)):#遍历预测框：坐标和类别
if len(detected_boxes) == len(annotations):
break
# Ignore if label is not one of the target labels
if pred_label not in target_labels:
continue
# 计算预测框和真实框的IOU
iou, box_index = bbox_iou(pred_box.unsqueeze(0), target_boxes).max(0)
#如果预测框和真实框的IOU大于阈值，那么可以认为该预测边界框预测’正确‘，并把该边界框的true_positives值设置为1
if iou >= iou_threshold and box_index not in detected_boxes:
true_positives[pred_i] = 1
detected_boxes += [box_index]
batch_metrics.append([true_positives, pred_scores, pred_labels])
return batch_metrics#true_positive,预测置信度，预测类别
"""未用到"""
def bbox_wh_iou(wh1, wh2):
wh2 = wh2.t()
w1, h1 = wh1[0], wh1[1]
w2, h2 = wh2[0], wh2[1]
inter_area = torch.min(w1, w2) * torch.min(h1, h2)
union_area = (w1 * h1 + 1e-16) + w2 * h2 - inter_area
return inter_area / union_area
"""计算两个边界框的IOU值"""
def bbox_iou(box1, box2, x1y1x2y2=True):
#获取边界框的左上右下坐标值
if not x1y1x2y2:
#如果边界框的表示方式为（center_x,center_y,width,height）则转换表示格式为（x，y，x，y）
b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
else:
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]#box1的左上右下坐标
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]#box1的左上右下坐标
#相交矩形的左上右下坐标
inter_rect_x1 = torch.max(b1_x1, b2_x1)
inter_rect_y1 = torch.max(b1_y1, b2_y1)
inter_rect_x2 = torch.min(b1_x2, b2_x2)
inter_rect_y2 = torch.min(b1_y2, b2_y2)
# 相交矩形的面积
inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(
inter_rect_y2 - inter_rect_y1 + 1, min=0
)
#并集的面积
b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)
iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)
return iou#返回重叠度IOU的值
'''非极大值抑制函数：返回边界框【x1,y1,x2,y2,conf,class_conf,class_pred】，参数为，模型预测，置信度阈值，nms阈值'''
def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
"""
Removes detections with lower object confidence score than 'conf_thres' and performs Non-Maximum Suppression to further filter detections.
Returns detections with shape:
(x1, y1, x2, y2, object_conf, class_score, class_pred)
"""
"""（1）模型预测坐标格式转变： (center x, center y, width, height) to (x1, y1, x2, y2)"""
#三个yolo层,有三个尺寸的输出分别为13,26,52，所以对于一张图片，
# 模型输出的shape是(10647,85),(13*13+26*26+52*52)*3=10647,后面的85是(x,y,w,h, conf, cls) xywh加一个置信度加80个分类。
#prediction的形状为[1, 10647, 85]，85的前4个信息为坐标信息（center x, center y, width, height）
# 第5个信息为目标置信度，第6-85的信息为80个类的置信度
prediction[..., :4] = xywh2xyxy(prediction[..., :4])# 将模型预测的坐标信息由(center x, center y, width, height) 格式转变为 (x1, y1, x2, y2)格式
output = [None for _ in range(len(prediction))]
#遍历每个图片，每张图片的预测image_pred：
for image_i, image_pred in enumerate(prediction):#遍历预测边界框
"""（2）边界框筛选：去除目标置信度低于阈值的边界框"""
image_pred = image_pred[image_pred[:, 4] >= conf_thres]#筛选每幅图片预测边界框中目标置信度大于阈值的边界框
# If none are remaining => process next image
if not image_pred.size(0):#判断本图片经过目标置信度阈值的赛选是否还存在边界框，如果没有边界框则执行下一个图片的NMS
continue
"""(3)非极大值抑制：根据score进行排序得到最大值，找到和这个score最大的预测类别相同的计算iou值，通过加权计算，得到最终的预测框(xyxy),最后从prediction中去掉iou大于设置的iou阈值的边界框。"""
# 分数=目标置信度*80个类别得分的最大值。
score = image_pred[:, 4] * image_pred[:, 5:].max(1)[0]
# 根据score为图片中的预测边界框进行排序
image_pred = image_pred[(-score).argsort()]#形状【经过置信度阈值筛选后的边界框数量，85】
#类别置信度最大值和类别置信度最大值所在位置（索引，也就是预测的类别）
class_confs, class_preds = image_pred[:, 5:].max(1, keepdim=True)#
detections = torch.cat((image_pred[:, :5], class_confs.float(), class_preds.float()), 1)#(x,y,x,y,object_confs,class_confs,class_preds)长度为7
keep_boxes = []
while detections.size(0):
# 将当前第一个边界框（当前分数最高的边界框）与剩余边界框计算IoU，并且大于NMS阈值的边界框
#第一个bbx与其余bbx的iou大于nms_thres的判别(0, 1)， 1为大于，0为小于
large_overlap = bbox_iou(detections[0, :4].unsqueeze(0), detections[:, :4]) > nms_thres
# 判断他们的类别是否相同，只有相同时才进行nms， 相同时为1， 不同时为0
label_match = detections[0, -1] == detections[:, -1]
# invalid 为Indices of boxes with lower confidence scores, large IOUs and matching labels
# 只有在两个bbx的iou大于thres，且类别相同时，invalid为True，其余为False
invalid = large_overlap & label_match
# weights为对应的权值, 其格式为：将True bbx中的confidence连成一个Tensor
weights = detections[invalid, 4:5]
# Merge overlapping bboxes by order of confidence
# 这里得到最后的bbx它是跟他满足IOU大于threshold，并且相同label的一些bbx，根据confidence重新加权得到
# 并不是原始bbx的保留。
detections[0, :4] = (weights * detections[invalid, :4]).sum(0) / weights.sum()
keep_boxes += [detections[0]]
## 去掉这些invalid，即iou大于阈值且预测同一类
detections = detections[~invalid]
if keep_boxes:
output[image_i] = torch.stack(keep_boxes)
return output#返回NMS后的边界框(x,y,x,y,object_confs,class_confs,class_preds)长度为7、
def build_targets(pred_boxes, pred_cls, target, anchors, ignore_thres):
ByteTensor = torch.cuda.ByteTensor if pred_boxes.is_cuda else torch.ByteTensor
FloatTensor = torch.cuda.FloatTensor if pred_boxes.is_cuda else torch.FloatTensor
nB = pred_boxes.size(0)
nA = pred_boxes.size(1)
nC = pred_cls.size(-1)
nG = pred_boxes.size(2)
# Output tensors
obj_mask = ByteTensor(nB, nA, nG, nG).fill_(0)
noobj_mask = ByteTensor(nB, nA, nG, nG).fill_(1)
class_mask = FloatTensor(nB, nA, nG, nG).fill_(0)
iou_scores = FloatTensor(nB, nA, nG, nG).fill_(0)
tx = FloatTensor(nB, nA, nG, nG).fill_(0)
ty = FloatTensor(nB, nA, nG, nG).fill_(0)
tw = FloatTensor(nB, nA, nG, nG).fill_(0)
th = FloatTensor(nB, nA, nG, nG).fill_(0)
tcls = FloatTensor(nB, nA, nG, nG, nC).fill_(0)
# Convert to position relative to box
target_boxes = target[:, 2:6] * nG
gxy = target_boxes[:, :2]
gwh = target_boxes[:, 2:]
# Get anchors with best iou
ious = torch.stack([bbox_wh_iou(anchor, gwh) for anchor in anchors])
best_ious, best_n = ious.max(0)
# Separate target values
b, target_labels = target[:, :2].long().t()
gx, gy = gxy.t()
gw, gh = gwh.t()
gi, gj = gxy.long().t()
# Set masks
obj_mask[b, best_n, gj, gi] = 1
noobj_mask[b, best_n, gj, gi] = 0
# Set noobj mask to zero where iou exceeds ignore threshold
for i, anchor_ious in enumerate(ious.t()):
noobj_mask[b[i], anchor_ious > ignore_thres, gj[i], gi[i]] = 0
# Coordinates
tx[b, best_n, gj, gi] = gx - gx.floor()
ty[b, best_n, gj, gi] = gy - gy.floor()
# Width and height
tw[b, best_n, gj, gi] = torch.log(gw / anchors[best_n][:, 0] + 1e-16)
th[b, best_n, gj, gi] = torch.log(gh / anchors[best_n][:, 1] + 1e-16)
# One-hot encoding of label
tcls[b, best_n, gj, gi, target_labels] = 1
# Compute label correctness and iou at best anchor
class_mask[b, best_n, gj, gi] = (pred_cls[b, best_n, gj, gi].argmax(-1) == target_labels).float()
iou_scores[b, best_n, gj, gi] = bbox_iou(pred_boxes[b, best_n, gj, gi], target_boxes, x1y1x2y2=False)
tconf = obj_mask.float()
return iou_scores, class_mask, obj_mask, noobj_mask, tx, ty, tw, th, tcls, tconf
2.4 weights文件夹
下载的预训练权重存放的文件夹
2.5 detect.py
模型训练完成后，进行检测测试的文件。验证数据集在data/samples文件夹下，验证结果保存在本py文件自动创建的文件夹output文件夹下。
from __future__ import division
from models import *
from utils.utils import *
from utils.datasets import *
import os
import time
import datetime
import argparse
from PIL import Image
import torch
from torch.utils.data import DataLoader
from torch.autograd import Variable
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.ticker import NullLocator
if __name__ == "__main__":
##########################################################################################################################
'''(1)参数解析'''
parser = argparse.ArgumentParser()
# 测试文件夹路径
parser.add_argument("--image_folder", type=str, default="data/samples", help="path to dataset")
#yolov3的模型信息（网络层，每层的卷积核数量，尺寸，步长。。。）
parser.add_argument("--model_def", type=str, default="config/yolov3.cfg", help="path to model definition file")
#预训练模型路径
parser.add_argument("--weights_path", type=str, default="weights/yolov3.weights", help="path to weights file")
#类名字
parser.add_argument("--class_path", type=str, default="data/coco.names", help="path to class label file")
#目标置信度阈值
parser.add_argument("--conf_thres", type=float, default=0.8, help="object confidence threshold")
#NMS的IoU阈值
parser.add_argument("--nms_thres", type=float, default=0.4, help="iou thresshold for non-maximum suppression")
#批量大小
parser.add_argument("--batch_size", type=int, default=1, help="size of the batches")
#CPU线程
parser.add_argument("--n_cpu", type=int, default=0, help="number of cpu threads to use during batch generation")
#图片维度
parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
#checkpoint_model
parser.add_argument("--checkpoint_model", type=str, help="path to checkpoint model")
opt = parser.parse_args()
print(opt)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs("output", exist_ok=True)#创建预测图片的输出位置
##########################################################################################################################
'''(2)模型构建'''
# 加载模型：这条语句加载darkent模型结构，即YOLOv3模型。Darknet模型在model.py中进行定义。
#将模型设置为评估模式
model = Darknet(opt.model_def, img_size=opt.img_size).to(device)#根据模型的配置文件，搭建模型的结构
#为模型结构加载训练的权重（模型参数）
if opt.weights_path.endswith(".weights"):
# Load darknet weights
model.load_darknet_weights(opt.weights_path)
else:
model.load_state_dict(torch.load(opt.weights_path))
model.eval() # 设置模型为评估模式，不然只要输入数据就会进行参数更新、优化
##########################################################################################################################
'''(3)数据集加载、类别加载'''
#加载测试的图片：
# dataloader本质是一个可迭代对象，使用iter()访问，不能使用next()访问；
#也可以使用`for inputs, labels in dataloaders`进行可迭代对象的访问
#一般我们实现一个datasets对象，传入到dataloader中；然后内部使用yeild返回每一次batch的数据
dataloader = DataLoader(
ImageFolder(opt.image_folder, img_size=opt.img_size),#评估数据集，ImageFolder在datasets.py中定义，返回的是图片路径，和经过处理（填充、调整大小）的图片
batch_size=opt.batch_size,
shuffle=False,
num_workers=opt.n_cpu,
)
#加载类别名，classes是一个列表
classes = load_classes(opt.class_path) # Extracts class labels from file
#创建保存图片路径和图片检测信息的列表
imgs = []
img_detections = []
##########################################################################################################################
"""(3)模型预测：将图片路径、图片预测结果存入imgs和img_detections列表中"""
print("\nPerforming object detection:")
prev_time = time.time()
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
# 测试图片的检测：并将图片路径和检测结果信息保存
# 算出batch中图片的地址img_paths和检测结果detections
for batch_i, (img_paths, input_imgs) in enumerate(dataloader):#使用dataloader加载数据，加载的数据为一批量的数据
# 把输入图像转换为tensor并变为变量
input_imgs = Variable(input_imgs.type(Tensor))
# 目标检测：使用模型检测图像，检测结果为一个张量，
# 对检测结果进行非极大值抑制，得到最终结果
with torch.no_grad():
detections = model(input_imgs)
#print(detections.shape)#[：, 10647, 85]
##非极大值抑制：将边界框信息，转变为左上右下坐标，并且去除置信度低的坐标. (x1, y1, x2, y2, object_conf, class_score, class_pred)
detections = non_max_suppression(detections, opt.conf_thres, opt.nms_thres)#非极大值抑制[:,:,7]
# 打印:检测时间，检测的批次
current_time = time.time()
inference_time = datetime.timedelta(seconds=current_time - prev_time)
prev_time = current_time
print("\t+ Batch %d, Inference Time: %s" % (batch_i, inference_time))
# 保存图片路径，图片的检测信息(经过NMS处理后)
imgs.extend(img_paths)
img_detections.extend(detections)#长度为7
##########################################################################################################################
"""（4）将检测结果绘制到图片，并保存"""
#边界框颜色
cmap = plt.get_cmap("tab20b") # Bounding-box colors
colors = [cmap(i) for i in np.linspace(0, 1, 20)]
#遍历图片
for img_i, (path, detections) in enumerate(zip(imgs, img_detections)):
print("(%d) Image: '%s'" % (img_i, path))
#读取图片并将图片绘制在plt.figure
img = np.array(Image.open(path))#读取图片
plt.figure()#创建图片画布
fig, ax = plt.subplots(1)
ax.imshow(img)#将读取的图片绘制到画布
#将图片对应的检测的边界框和标签绘制到图片上
if detections is not None:
# 将检测的边界框（对填充、调整大小的原图的预测），重新设置尺寸，使其与原图目标能匹配
detections = rescale_boxes(detections, opt.img_size, img.shape[:2])
#获取检测结果的类标签，并为每一个类指定一种颜色
unique_labels = detections[:, -1].cpu().unique()#返回参数数组中所有不同的值,并按照从小到大排序可选参数
n_cls_preds = len(unique_labels)
bbox_colors = random.sample(colors, n_cls_preds)#为每一类分配一个边界框颜色
#遍历图片对应检测结果的每一个边界框
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:#检测结果为左上和右下坐标
print("\t+ Label: %s, Conf: %.5f" % (classes[int(cls_pred)], cls_conf.item()))
#边界框宽和高
box_w = x2 - x1
box_h = y2 - y1
#将边界框写入图片中，并设置颜色
color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]
# 创建一个矩形边界框
bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=color, facecolor="none")
# 吧矩形边界框写入画布
ax.add_patch(bbox)
# 为检测边界框添加类别信息
plt.text( x1,y1,s=classes[int(cls_pred)],color="white",verticalalignment="top",bbox={"color": color, "pad": 0} )
#将绘制好边界框的图片保存
plt.axis("off")
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
filename = path.split("/")[-1].split(".")[0]
plt.savefig(f"output/{filename}.png", bbox_inches="tight", pad_inches=0.0)
plt.close()
2.6 models.py
定义模型结构的文件，根据模型的配置文件信息，来构建模型结构。
from __future__ import division
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from utils.parse_config import *
from utils.utils import build_targets, to_cpu
'''构建网络函数：通过获取的模型定义module_defs来构建YOLOv3模型结构,根据module_defs中的模块配置构造层块的模块列表'''
def create_modules(module_defs):
'''构建模型结构'''
'''（1）解析模型超参数，获取模型的输入通道数'''
#从model_def获取net的配置信息组成的字典hyperparams。model_def是由parse_config函数解析出来的列表，每个元素为一个字典，每一个字典包含了某层、模块的参数信息
hyperparams = module_defs.pop(0)#hyperparams为module_defs的第一个字典元素，是模型的超参数信息{'type': 'net',...}
output_filters = [int(hyperparams["channels"])]
'''(2)构建nn.ModuleList()，用来存放创建的网络层、模块'''
module_list = nn.ModuleList()
'''(3)遍历模型定义列表的每个字典元素，创建相应的层、模块，添加到nn.ModuleList()中'''
#遍历 module_defs的每个字典，根据字典内容，创建相应的层或模块。其中字典的type的值有一下几种："convolutional"，"maxpool"
#"upsample"， "route"，"shortcut"， "yolo"
for module_i, module_def in enumerate(module_defs):
#创建一个 nn.Sequential()
modules = nn.Sequential()
#卷积层构建，并添加到nn.Sequential()
if module_def["type"] == "convolutional":
#获取convolutional层的参数信息
bn = int(module_def["batch_normalize"])
filters = int(module_def["filters"])
kernel_size = int(module_def["size"])
pad = (kernel_size - 1) // 2
#创建convolution层：根据convolutional层的参数信息，创建convolutional层，并将改层加入到nn.Sequential()中
modules.add_module(f"conv_{module_i}",#层在模型中的名字
nn.Conv2d(#层
in_channels=output_filters[-1],#输入的通道数
out_channels=filters,#输出的通道数
kernel_size=kernel_size,#卷结核大小
stride=int(module_def["stride"]),#步长
padding=pad,#填充
bias=not bn,
),
)
if bn:
#添加BatchNorm2d层
modules.add_module(f"batch_norm_{module_i}", nn.BatchNorm2d(filters, momentum=0.9, eps=1e-5))
if module_def["activation"] == "leaky":
#添加激活层LeakyReLU
modules.add_module(f"leaky_{module_i}", nn.LeakyReLU(0.1))
#池化层构建，并添加到nn.Sequential()
elif module_def["type"] == "maxpool":
# 获取maxpool层的参数信息
kernel_size = int(module_def["size"])
stride = int(module_def["stride"])
# 根据maxpool层的参数信息，创建maxpool层，并将改层加入到 nn.Sequential()中
if kernel_size == 2 and stride == 1:
modules.add_module(f"_debug_padding_{module_i}", nn.ZeroPad2d((0, 1, 0, 1)))
#创建maxpool层
modules.add_module(f"maxpool_{module_i}",
nn.MaxPool2d(
kernel_size=kernel_size, #卷积核大小
stride=stride, #步长
padding=int((kernel_size - 1) // 2))#填充
)
#上采样层构建，并添加到nn.Sequential()
#上采样层是自定义的层，需要实例化Upsample为一个对象，将对象层添加到模型列表中
elif module_def["type"] == "upsample":
#上采样的配置例，如下
# [upsample]
# stride = 2
# 构建upsample层，上采样层类，重写了forward函数
upsample = Upsample(scale_factor=int(module_def["stride"]), mode="nearest")
#层添加到模型
modules.add_module(f"upsample_{module_i}", upsample)
elif module_def["type"] == "route":
#youte信息，例
# [route]
# layers = -1, 36
# 获取route层的参数信息
layers = [int(x) for x in module_def["layers"].split(",")]
filters = sum([output_filters[1:][i] for i in layers])
modules.add_module(f"route_{module_i}", EmptyLayer())#EmptyLayer()为“路线”和“快捷方式”层的占位符
elif module_def["type"] == "shortcut":
filters = output_filters[1:][int(module_def["from"])]
modules.add_module(f"shortcut_{module_i}", EmptyLayer())#EmptyLayer()为“路线”和“快捷方式”层的占位符
elif module_def["type"] == "yolo":
#例：假设yolo的配置信息如下
# [yolo]
# mask = 3,4,5
# anchors = 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326
# classes=80
# num=9
# jitter=.3
# ignore_thresh = .7
# truth_thresh = 1
# random=1
#获取anchor的索引，上例为3,4,5
anchor_idxs = [int(x) for x in module_def["mask"].split(",")]
#提取anchor尺寸信息,放入列表
anchors = [int(x) for x in module_def["anchors"].split(",")]
anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)]
anchors = [anchors[i] for i in anchor_idxs]
num_classes = int(module_def["classes"])
#print('anchors1:', anchors)#上例为anchors1: [(30, 61), (62, 45), (59, 119)]
#获取图片的输入尺寸
img_size = int(hyperparams["height"])
#定义yolo检测层：实例化yolo类，创建yolo层，传入的参数为三个anchor的尺寸，类别的数量，图像的大小
yolo_layer = YOLOLayer(anchors, num_classes, img_size)
#将YOLO层加入到模型列表
modules.add_module(f"yolo_{module_i}", yolo_layer)
module_list.append(modules) #将创建的nn.Sequential()即创建的层，添加到 nn.ModuleList()中
output_filters.append(filters)#将创建的层的输出通道数添加到filters列表中，作为下次创建层的输入通道数
return hyperparams, module_list#返回网络的参数、网络结构即层组成的列表
'''上采样层'''
class Upsample(nn.Module):
""" nn.Upsample 被重写 """
def __init__(self, scale_factor, mode="nearest"):
super(Upsample, self).__init__()
self.scale_factor = scale_factor#上采样步长
self.mode = mode
def forward(self, x):
x = F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)#上采样方法，插值
return x#返回上采样结果
'''emptylayer定义'''
class EmptyLayer(nn.Module):
"""Placeholder for 'route' and 'shortcut' layers"""
def __init__(self):
super(EmptyLayer, self).__init__()
'''yolo层定义：检测层'''
class YOLOLayer(nn.Module):
"""Detection layer"""
def __init__(self, anchors, num_classes, img_dim=416):#参数为三个anchor的尺寸，类别的数量，图像的大小
super(YOLOLayer, self).__init__()
#基础设置
self.anchors = anchors#anchor的尺寸信息，例某一层yolo尺寸为[(30, 61), (62, 45), (59, 119)]
self.num_anchors = len(anchors)#anchor的数量
self.num_classes = num_classes#类别的数量
self.ignore_thres = 0.5
self.mse_loss = nn.MSELoss()
self.bce_loss = nn.BCELoss()
self.obj_scale = 1
self.noobj_scale = 100
self.metrics = {}
self.img_dim = img_dim
self.grid_size = 0 # grid size
#计算网格单元偏移
def compute_grid_offsets(self, grid_size, cuda=True):
#获取网格尺寸（几×几）
self.grid_size = grid_size
g = self.grid_size
# print('g',g) g可能的取值为13/26/52，对应不同yolo层的特征图的尺寸
FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
#获取网格单元大小
self.stride = self.img_dim / self.grid_size#网格单元的尺寸
# Calculate offsets for each grid，假设g取13,
#torch.arange(g) 为tensor([0,1,2,3,4,5,6,7,8,9,10,11,12])
#torch.arange(g).repeat(g, 1) 为由tensor([0,1,2,3,4,5,6,7,8,9,10,11,12])组成的13行一列的张量
#torch.arange(g).repeat(g, 1).view([1, 1, g, g]) 改变视图为【1,1,13,13】
self.grid_x = torch.arange(g).repeat(g, 1).view([1, 1, g, g]).type(FloatTensor)#
self.grid_y = torch.arange(g).repeat(g, 1).t().view([1, 1, g, g]).type(FloatTensor)
#把anchor的宽和高转变为相对于网格单元大小的度量
self.scaled_anchors = FloatTensor([(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors])#例某一层yolo尺寸为[(30, 61), (62, 45), (59, 119)]
self.anchor_w = self.scaled_anchors[:, 0:1].view((1, self.num_anchors, 1, 1))#获取anchor的宽
self.anchor_h = self.scaled_anchors[:, 1:2].view((1, self.num_anchors, 1, 1))#获取anchor的高
def forward(self, x, targets=None, img_dim=None):
#yolo层的前向传播，参数为yolo层来自上层的输出作为输入x
FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
#图片的大小
self.img_dim = img_dim
#获取x的形状
num_samples = x.size(0)
grid_size = x.size(2)
prediction = (
x.view(num_samples, self.num_anchors, self.num_classes + 5, grid_size, grid_size)#（num_samples,3,85,gride_size,grid_size）
.permute(0, 1, 3, 4, 2)#permute是用来做维度换位置的，（num_samples,3,gride_size,grid_size,85）
.contiguous()#调用contiguous()时，会强制拷贝一份tensor，让它的布局和从头创建的一毛一样。而不是与原数据公用一份内存。
)
# 得到outputs
x = torch.sigmoid(prediction[..., 0]) # Center x
y = torch.sigmoid(prediction[..., 1]) # Center y
w = prediction[..., 2] # Width
h = prediction[..., 3] # Height
pred_conf = torch.sigmoid(prediction[..., 4]) # Conf
pred_cls = torch.sigmoid(prediction[..., 5:]) # Cls pred.
# If grid size does not match current we compute new offsets
if grid_size != self.grid_size:
self.compute_grid_offsets(grid_size, cuda=x.is_cuda)
# Add offset and scale with anchors
pred_boxes = FloatTensor(prediction[..., :4].shape)
pred_boxes[..., 0] = x.data + self.grid_x
pred_boxes[..., 1] = y.data + self.grid_y
pred_boxes[..., 2] = torch.exp(w.data) * self.anchor_w
pred_boxes[..., 3] = torch.exp(h.data) * self.anchor_h
output = torch.cat(
(
pred_boxes.view(num_samples, -1, 4) * self.stride,
pred_conf.view(num_samples, -1, 1),
pred_cls.view(num_samples, -1, self.num_classes),
),
-1,
)
if targets is None:
return output, 0
else:
iou_scores, class_mask, obj_mask, noobj_mask, tx, ty, tw, th, tcls, tconf = build_targets(
pred_boxes=pred_boxes,
pred_cls=pred_cls,
target=targets,
anchors=self.scaled_anchors,
ignore_thres=self.ignore_thres,
)
# Loss : Mask outputs to ignore non-existing objects (except with conf. loss)
loss_x = self.mse_loss(x[obj_mask], tx[obj_mask])
loss_y = self.mse_loss(y[obj_mask], ty[obj_mask])
loss_w = self.mse_loss(w[obj_mask], tw[obj_mask])
loss_h = self.mse_loss(h[obj_mask], th[obj_mask])
loss_conf_obj = self.bce_loss(pred_conf[obj_mask], tconf[obj_mask])
loss_conf_noobj = self.bce_loss(pred_conf[noobj_mask], tconf[noobj_mask])
loss_conf = self.obj_scale * loss_conf_obj + self.noobj_scale * loss_conf_noobj
loss_cls = self.bce_loss(pred_cls[obj_mask], tcls[obj_mask])
total_loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls
# Metrics
cls_acc = 100 * class_mask[obj_mask].mean()
conf_obj = pred_conf[obj_mask].mean()
conf_noobj = pred_conf[noobj_mask].mean()
conf50 = (pred_conf > 0.5).float()
iou50 = (iou_scores > 0.5).float()
iou75 = (iou_scores > 0.75).float()
detected_mask = conf50 * class_mask * tconf
precision = torch.sum(iou50 * detected_mask) / (conf50.sum() + 1e-16)
recall50 = torch.sum(iou50 * detected_mask) / (obj_mask.sum() + 1e-16)
recall75 = torch.sum(iou75 * detected_mask) / (obj_mask.sum() + 1e-16)
self.metrics = {
"loss": to_cpu(total_loss).item(),
"x": to_cpu(loss_x).item(),
"y": to_cpu(loss_y).item(),
"w": to_cpu(loss_w).item(),
"h": to_cpu(loss_h).item(),
"conf": to_cpu(loss_conf).item(),
"cls": to_cpu(loss_cls).item(),
"cls_acc": to_cpu(cls_acc).item(),
"recall50": to_cpu(recall50).item(),
"recall75": to_cpu(recall75).item(),
"precision": to_cpu(precision).item(),
"conf_obj": to_cpu(conf_obj).item(),
"conf_noobj": to_cpu(conf_noobj).item(),
"grid_size": grid_size,
}
return output, total_loss
"""Darknet类：YOLOv3模型"""
class Darknet(nn.Module):
"""YOLOv3 object detection model"""
def __init__(self, config_path, img_size=416):
super(Darknet, self).__init__()
# parse_model_config（）模型配置的解析器:用来解析yolo-v3层配置文件(yolov3.cfg)并返回模块定义
#（模型定义module_defs是一个列表，每一个元素是一个字典，该字典描绘了网络每一个模块/层的信息）
self.module_defs = parse_model_config(config_path)
#通过获取的模型定义module_defs，来构建YOLOv3模型
self.hyperparams,self.module_list = create_modules(self.module_defs)#模型参数和模型结构
self.yolo_layers = [layer[0] for layer in self.module_list if hasattr(layer[0], "metrics")]
self.img_size = img_size
self.seen = 0
self.header_info = np.array([0, 0, 0, self.seen, 0], dtype=np.int32)
def forward(self, x, targets=None):
img_dim = x.shape[2]
loss = 0
layer_outputs, yolo_outputs = [], []
for i, (module_def, module) in enumerate(zip(self.module_defs, self.module_list)):
if module_def["type"] in ["convolutional", "upsample", "maxpool"]:
x = module(x)
elif module_def["type"] == "route":
x = torch.cat([layer_outputs[int(layer_i)] for layer_i in module_def["layers"].split(",")], 1)
elif module_def["type"] == "shortcut":
layer_i = int(module_def["from"])
x = layer_outputs[-1] + layer_outputs[layer_i]
elif module_def["type"] == "yolo":
x, layer_loss = module[0](x, targets, img_dim)
loss += layer_loss
yolo_outputs.append(x)
layer_outputs.append(x)
yolo_outputs = to_cpu(torch.cat(yolo_outputs, 1))
return yolo_outputs if targets is None else (loss, yolo_outputs)
def load_darknet_weights(self, weights_path):
"""Parses and loads the weights stored in 'weights_path'"""
# Open the weights file
with open(weights_path, "rb") as f:
header = np.fromfile(f, dtype=np.int32, count=5) # First five are header values
self.header_info = header # Needed to write header when saving weights
self.seen = header[3] # number of images seen during training
weights = np.fromfile(f, dtype=np.float32) # The rest are weights
# Establish cutoff for loading backbone weights
cutoff = None
if "darknet53.conv.74" in weights_path:
cutoff = 75
ptr = 0
for i, (module_def, module) in enumerate(zip(self.module_defs, self.module_list)):
if i == cutoff:
break
if module_def["type"] == "convolutional":
conv_layer = module[0]
if module_def["batch_normalize"]:
# Load BN bias, weights, running mean and running variance
bn_layer = module[1]
num_b = bn_layer.bias.numel() # Number of biases
# Bias
bn_b = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(bn_layer.bias)
bn_layer.bias.data.copy_(bn_b)
ptr += num_b
# Weight
bn_w = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(bn_layer.weight)
bn_layer.weight.data.copy_(bn_w)
ptr += num_b
# Running Mean
bn_rm = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(bn_layer.running_mean)
bn_layer.running_mean.data.copy_(bn_rm)
ptr += num_b
# Running Var
bn_rv = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(bn_layer.running_var)
bn_layer.running_var.data.copy_(bn_rv)
ptr += num_b
else:
# Load conv. bias
num_b = conv_layer.bias.numel()
conv_b = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(conv_layer.bias)
conv_layer.bias.data.copy_(conv_b)
ptr += num_b
# Load conv. weights
num_w = conv_layer.weight.numel()
conv_w = torch.from_numpy(weights[ptr : ptr + num_w]).view_as(conv_layer.weight)
conv_layer.weight.data.copy_(conv_w)
ptr += num_w
def save_darknet_weights(self, path, cutoff=-1):
"""
@:param path - path of the new weights file
@:param cutoff - save layers between 0 and cutoff (cutoff = -1 -> all are saved)
"""
fp = open(path, "wb")
self.header_info[3] = self.seen
self.header_info.tofile(fp)
# Iterate through layers
for i, (module_def, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])):
if module_def["type"] == "convolutional":
conv_layer = module[0]
# If batch norm, load bn first
if module_def["batch_normalize"]:
bn_layer = module[1]
bn_layer.bias.data.cpu().numpy().tofile(fp)
bn_layer.weight.data.cpu().numpy().tofile(fp)
bn_layer.running_mean.data.cpu().numpy().tofile(fp)
bn_layer.running_var.data.cpu().numpy().tofile(fp)
# Load conv bias
else:
conv_layer.bias.data.cpu().numpy().tofile(fp)
# Load conv weights
conv_layer.weight.data.cpu().numpy().tofile(fp)
fp.close()
2.7 test.py
用来评估模型性能的文件。
from __future__ import division
from models import *
from utils.utils import *
from utils.datasets import *
from utils.parse_config import *
import argparse
import tqdm
import torch
from torch.utils.data import DataLoader
from torch.autograd import Variable
"""模型评估函数：参数为模型、valid数据集路径、iou阈值。nms阈值、网络输入大小、批量大小"""
def evaluate(model, path, iou_thres, conf_thres, nms_thres, img_size, batch_size):
#加上model.eval(). 否则的话，有输入数据，即使不训练，它也会改变权值
model.eval()
'''(1)获取评估数据集：变为batch组成的数据集'''
# dataset（验证集图片路径集、验证集图片集，验证集标签集）
# dataloader获取批量batch,验证集图片路径batch、验证集图片batch，验证集标签batch）
dataset = ListDataset(path, img_size=img_size, augment=False, multiscale=False)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=1,
collate_fn=dataset.collate_fn)#collate_fn参数，实现自定义的batch输出
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
labels = []
sample_metrics = [] # List of tuples (TP, confs, pred)
for batch_i, (_, imgs, targets) in enumerate(tqdm.tqdm(dataloader, desc="Detecting objects")):#tqdm进度条
'''(2) batch标签处理'''
labels += targets[:, 1].tolist()#将targets的类别信息转变为list存到label列表中
# Rescale target
targets[:, 2:] = xywh2xyxy(targets[:, 2:])#将targets的坐标变为（xyxy）形式，此时的坐标也是归一化的形式
targets[:, 2:] *= img_size#适应于原图的比target形式
'''(3)batch图片预测，并进行NMS处理'''
# 图片输入模型，并对模型输出进行非极大值抑制
imgs = Variable(imgs.type(Tensor), requires_grad=False)
with torch.no_grad():
outputs = model(imgs)
outputs = non_max_suppression(outputs, conf_thres=conf_thres, nms_thres=nms_thres)
'''（4）预测信息统计：得到经过NMS处理后，预测边界框的true_positive（值为或1）、预测置信度，预测类别信息'''
sample_metrics += get_batch_statistics(outputs, targets, iou_threshold=iou_thres)#参数：模型输出，真实标签（适应于原图的x,y,x,y）,iou阈值
# 这里需要注意,github上面的代码有错误,需要添加if条件语句，训练才能正常运行
if len(sample_metrics) == 0:
return np.array([]), np.array([]), np.array([]), np.array([]), np.array([])
# sample_metrics信息解析，获取独立的 true_positive（值为或1）、预测置信度，预测类别 信息
true_positives, pred_scores, pred_labels = [np.concatenate(x, 0) for x in list(zip(*sample_metrics))]
#计算 precision, recall, AP, f1, ap_class，这里调用了utils.py中的函数进行计算
precision, recall, AP, f1, ap_class = ap_per_class(true_positives, pred_scores, pred_labels, labels)#pred_labels, labels的长度是不同的
return precision, recall, AP, f1, ap_class
if __name__ == "__main__":
'''(1)参数解析'''
parser = argparse.ArgumentParser()
parser.add_argument("--batch_size", type=int, default=8, help="size of each image batch")
parser.add_argument("--model_def", type=str, default="config/yolov3.cfg", help="path to model definition file")
parser.add_argument("--data_config", type=str, default="config/custom.data", help="path to data config file")
parser.add_argument("--weights_path", type=str, default="checkpoints/yolov3_ckpt_9.pth", help="path to weights file")#"weights/yolov3.weights"
parser.add_argument("--class_path", type=str, default="data/coco.names", help="path to class label file")
parser.add_argument("--iou_thres", type=float, default=0.5, help="iou threshold required to qualify as detected")
parser.add_argument("--conf_thres", type=float, default=0.001, help="object confidence threshold")
parser.add_argument("--nms_thres", type=float, default=0.5, help="iou thresshold for non-maximum suppression")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
opt = parser.parse_args()
#print(opt)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
"""（2）数据解析"""
# 调用parse_config。py中的数据解析桉树，返回值 data_config 为字典{class:80,train:路径，valid:路径。。。}
data_config = parse_data_config(opt.data_config)
valid_path = data_config["valid"]#验证集路径valid=data/custom/valid.txt
class_names = load_classes(data_config["names"])#类别路径
"""（3）模型构建：构建模型，加载模型参数"""
model = Darknet(opt.model_def).to(device)
if opt.weights_path.endswith(".weights"):
# Load darknet weights
model.load_darknet_weights(opt.weights_path)#
else:
model.load_state_dict(torch.load(opt.weights_path))#自定义的函数
print("Compute mAP...")
"""(4)模型评估"""
precision, recall, AP, f1, ap_class = evaluate(
model,#模型
path=valid_path,#验证集路径
iou_thres=opt.iou_thres,
conf_thres=opt.conf_thres,#置信度阈值
nms_thres=opt.nms_thres,#nms阈值
img_size=opt.img_size,#网路输入尺寸
batch_size=8,#批量
)
print(precision, recall, AP, f1, ap_class)
print("Average Precisions:")
for i, c in enumerate(ap_class):
print(f"+ Class '{c}' ({class_names[c]}) - AP: {AP[i]}")
print(f"mAP: {AP.mean()}")
2.8 train.py
模型训练的文件夹，训练会生成：
（1）checkpoint文件夹，用来保存某epoch训练后的模型参数
（2）logs文件夹，用来保存日志信息
from __future__ import division
from models import *
from utils.logger import *
from utils.utils import *
from utils.datasets import *
from utils.parse_config import *
from terminaltables import AsciiTable
import os
from test import evaluate
import time
import datetime
import argparse
import torch
from torch.utils.data import DataLoader
from torch.autograd import Variable
if __name__ == "__main__":
'''（1）参数解析'''
parser = argparse.ArgumentParser()
parser.add_argument("--epochs", type=int, default=10, help="number of epochs")
parser.add_argument("--batch_size", type=int, default=1, help="size of each image batch")
#梯度累加数
parser.add_argument("--gradient_accumulations", type=int, default=2, help="number of gradient accums before step")
parser.add_argument("--model_def", type=str, default="config/yolov3.cfg", help="path to model definition file")
parser.add_argument("--data_config", type=str, default="config/custom.data", help="path to data config file")
parser.add_argument("--pretrained_weights", type=str, help="if specified starts from checkpoint model")
parser.add_argument("--n_cpu", type=int, default=1, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
parser.add_argument("--checkpoint_interval", type=int, default=1, help="interval between saving model weights")
parser.add_argument("--evaluation_interval", type=int, default=1, help="interval evaluations on validation set")
parser.add_argument("--compute_map", default=False, help="if True computes mAP every tenth batch")
parser.add_argument("--multiscale_training", default=True, help="allow for multi-scale training")
parser.add_argument("--weights_path", type=str, default="checkpoints/yolov3_ckpt_9.pth", help="path to weights file")
opt = parser.parse_args()
print(opt)
'''(2)实例化日志类'''
logger = Logger("logs")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
'''(3)文件夹创建'''
os.makedirs("output", exist_ok=True)
os.makedirs("checkpoints", exist_ok=True)
"""(4)初始化模型：模型构建，模型参数装载"""
model = Darknet(opt.model_def).to(device)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if opt.pretrained_weights:
if opt.pretrained_weights.endswith(".pth"):
model.load_state_dict(torch.load(opt.pretrained_weights))
else:
model.load_darknet_weights(opt.pretrained_weights)
"""(5)数据集加载"""
data_config = parse_data_config(opt.data_config)#调用parse_config.py文件的数据配置解析函数，获取data_config为一个字典
train_path = data_config["train"]#训练集路径
valid_path = data_config["valid"]#验证集路径
class_names = load_classes(data_config["names"])#调用utils.py内的load_classes函数用于获取数据集包含的类别名称
#dataset是数据集中，图片的路径和、图片、标签（归一化的格式x,y,w,h）的集合
dataset = ListDataset(train_path, augment=True, multiscale=opt.multiscale_training)
#dataloader是dataset装载成批量形式
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.n_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn,
)
"""(7)优化器"""
optimizer = torch.optim.Adam(model.parameters())
"""(8)模型训练"""
metrics = [
"grid_size",
"loss",
"x",
"y",
"w",
"h",
"conf",
"cls",
"cls_acc",
"recall50",
"recall75",
"precision",
"conf_obj",
"conf_noobj",
]
for epoch in range(opt.epochs):#迭代epoch次训练
model.train()#设置模型为训练模式
start_time = time .time()
print('start_time',start_time)
for batch_i, (_, imgs, targets) in enumerate(dataloader):#每一epoch的批量迭代
#批量的累计迭代数
batches_done = len(dataloader) * epoch + batch_i
#图片、标签的变量化处理
imgs = Variable(imgs.to(device))#把图像变为变量，可以记录梯度
targets = Variable(targets.to(device), requires_grad=False)#把标签变为变量，不记录梯度
# 获取模型的输出与损失，损失反向传播
loss, outputs = model(imgs, targets)#将图片和标签输入模型，获取输出
loss.backward()
#计算梯度
if batches_done % opt.gradient_accumulations:
# 在每一步之前计算梯度Accumulates gradient before each step
optimizer.step()
optimizer.zero_grad()
#训练的epoch及batch信息
log_str = "\n---- [Epoch %d/%d, Batch %d/%d] ----\n" % (epoch+1, opt.epochs, batch_i+1, len(dataloader))
#print('log_str',log_str)#例---- [Epoch 1/10, Batch 1/10] ----
#创建行索引
metric_table = [["Metrics", *[f"YOLO Layer {i}" for i in range(len(model.yolo_layers))]]]#创建训练过程中的表格，行索引
#print(metric_table)# [['Metrics', 'YOLO Layer 0', 'YOLO Layer 1', 'YOLO Layer 2']]
# 在每一个 YOLO layer的各项指标信息
for i, metric in enumerate(metrics):#metrics为各项指标名称组成的列表，上面已经定义
#获取metrics各个项的数值类型
formats = {m: "%.6f" for m in metrics}#将所有的metrics中的输出数值类型定义，这一步把全部的输出类型全部定义保留6位小数
formats["grid_size"] = "%2d"
formats["cls_acc"] = "%.2f%%"
#print(' formats', formats)#{'grid_size': '%2d', 'loss': '%.6f', 'x': '%.6f', 'y': '%.6f', 'w': '%.6f', 'h': '%.6f', 'conf': '%.6f', 'cls': '%.6f', 'cls_acc': '%.2f%%', 'recall50': '%.6f', 'recall75': '%.6f', 'precision': '%.6f', 'conf_obj': '%.6f', 'conf_noobj': '%.6f'}
#表格赋值
row_metrics = [formats[metric] % yolo.metrics.get(metric, 0) for yolo in model.yolo_layers]#？？？？？？？？？？？？？
#print('row_metrics',row_metrics)
metric_table += [[metric, *row_metrics]]
# Tensorboard 日志信息
tensorboard_log = []
for j, yolo in enumerate(model.yolo_layers):
for name, metric in yolo.metrics.items():
if name != "grid_size":
tensorboard_log += [(f"{name}_{j+1}", metric)]#把除grid_size的其余信息，添加到日志中
tensorboard_log += [("loss", loss.item())]#把损失也添加到日志信息中
#把日志信息列表写入创建的日志对象
logger.list_of_scalars_summary(tensorboard_log, batches_done)
#log_str打印各项指标参数：
log_str += AsciiTable(metric_table).table
log_str += f"\nTotal loss {loss.item()}"
# 计算该epoch剩余需要的大概时间
epoch_batches_left = len(dataloader) - (batch_i + 1)
time_left = datetime.timedelta(seconds=epoch_batches_left * (time.time() - start_time) / (batch_i + 1))
log_str += f"\n---- ETA {time_left}"
print(log_str)
model.seen += imgs.size(0)
'''(9)训练时评估'''
if epoch % opt.evaluation_interval == 0:
print("\n---- Evaluating Model ----")
# 在评估数据集上对当前模型进行评估，具体评估细节可以看test.py
precision, recall, AP, f1, ap_class = evaluate(
model,
path=valid_path,
iou_thres=0.5,
conf_thres=0.5,
nms_thres=0.5,
img_size=opt.img_size,
batch_size=8,
)
evaluation_metrics = [
("val_precision", precision.mean()),
("val_recall", recall.mean()),
("val_mAP", AP.mean()),
("val_f1", f1.mean()),
]
logger.list_of_scalars_summary(evaluation_metrics, epoch)
# Print class APs and mAP
ap_table = [["Index", "Class name", "AP"]]
for i, c in enumerate(ap_class):
ap_table += [[c, class_names[c], "%.5f" % AP[i]]]
print(AsciiTable(ap_table).table)
print(f"---- mAP {AP.mean()}")
'''(10)模型保存'''
if epoch % opt.checkpoint_interval == 0:
torch.save(model.state_dict(), f"checkpoints/yolov3_ckpt_%d.pth" % epoch)