YOLO V1代码

参考:YOLOV1(pytorch) 代码讲解

1、数据集分析和处理

        首先分析要处理的数据,由于个人电脑内存有限,处理数据,只选取VOC 2007数据集中的6张图片。主要目的是理解YOLO V1的工作原理。

YOLO V1代码_第1张图片

 下面是000015.xml的内容:

     
   VOC2007   
   
   000015.jpg
   
   
      The VOC2007 Database  
      PASCAL VOC2007
       
      flickr
      328420077
   
   
   
      toxan
      ?
   
    
   
      500
      375
      3
   
    
   0
   
   
      
      bicycle
      Left
      
      0
      
      0
       
      
         77
         136
         360
         358
      
   

目标框的定义方式

目标检测中,目标的标签需要包含两类信息:

  • 目标类别
  • 目标的位置信息(也就是目标的外接矩形bounding box, 简称为bbox)

用来表达bbox的格式通常有两种,(x1, y1, x2, y2) 和 (c_x, c_y, w, h) ,

(x1,y1)表示左上角坐标,(x2,y2)表示右下角坐标

(c_x,c_y)表示bbox中心点坐标,(w,h)表示bbox的宽高

2、画出数据集中某一个标注好的图片及其标注框

内有详细注释

YOLO V1代码_第2张图片

import mxnet as mx
import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
# 在图像中绘制方框
from matplotlib.patches import Rectangle
import random
import cv2

# 测试图像的标注文件和图像
xml_path = 'VOCdevkit/VOC2007/Annotations/000015.xml'
img_path = 'VOCdevkit/VOC2007/JPEGImages/000015.jpg'

bbox = []
tree = ET.parse(xml_path) #生成一个总目录名为tree
root = tree.getroot()
objects = root.findall('object') #在总目录下找到object,名为objects
for object in objects:
    # 根据找出的object再进一步找出图像中的类别名
    name = object.find('name').text
    # xml中的标记信息中的bndbox
    bndbox = object.find('bndbox')

    # 对应坐标
    xmin = int(bndbox.find('xmin').text)
    ymin = int(bndbox.find('ymin').text)
    xmax = int(bndbox.find('xmax').text)
    ymax = int(bndbox.find('ymax').text)

    # 格式:类名:name+[坐标]
    bbox_i = [name, xmin, ymin, xmax, ymax]
    bbox.append(bbox_i)

print(bbox)

#将作为二进制字符串加载的图像转换为numpy数组
image_string = open(img_path, 'rb').read() #读取的是二进制字符串
#cv2.imdecode()函数从指定的内存缓存中读取数据,
#并把数据转换(解码)成图像格式;主要用于从网络传输数据中恢复出图像。
#展示的是原始图片的大小
image = mx.image.imdecode(image_string, flag=1).asnumpy()
# 显示原图
#它是用来创建 总画布/figure“窗口”的,
# 有figure就可以在上边(或其中一个子网格/subplot上)作图了,(fig:是figure的缩写)。
fig, ax = plt.subplots()
plt.imshow(image)

#colors为字典形式
colors = dict()
for bbox_i in bbox:
    class_name = bbox_i[0]
    if class_name not in colors:
        # 随机生成颜色或者自己指定也可以
        #{'bicycle': (0.46746658758487813, 0.38011471681747944, 0.038329821431872024)}
        colors[class_name] = (random.random(), random.random(), random.random())
    xmin = bbox_i[1]
    ymin = bbox_i[2]
    xmax = bbox_i[3]
    ymax = bbox_i[4]

    #rectangle函数是用来绘制一个矩形框的,通常用在图片的标记上。
    #Rectangle(xy=(77, 136), width=283, height=222, angle=0)
    rect = Rectangle(xy=(xmin, ymin), width=xmax - xmin,
                     height=ymax - ymin,
                     edgecolor=colors[class_name],
                     facecolor='None',
                     linewidth=3.5  #线条宽度
                     )
    #给图形添加标签
    plt.text(xmin, ymin - 6, '{:s}'.format(class_name),
             bbox=dict(facecolor=colors[class_name],
                       alpha=0.5))
    ##添加到图版中
    ax.add_patch(rect)
    plt.axis('off')
    plt.savefig('VOC_image_draw/000015_groundtruth')

3、文件

 此文件的主要任务是,读取文件夹内所有的xml文件及其信息,然后将其信息(name,bbox,class)写入一个txt文件,在此阶段训练集以及测试集被划分开来这里是按照7:3的比例进行划分,后续的数据集处理需要用到这些信息。

import xml.etree.ElementTree as ET
import os
import random

VOC_CLASSES = (  # 定义所有的类名
    'aeroplane', 'bicycle', 'bird', 'boat',
    'bottle', 'bus', 'car', 'cat', 'chair',
    'cow', 'diningtable', 'dog', 'horse',
    'motorbike', 'person', 'pottedplant',
    'sheep', 'sofa', 'train', 'tvmonitor')   # 使用其他训练集需要更改

# 定义一些参数
#训练数据集
train_set = open('voctrain.txt', 'w')
#测试数据集
test_set = open('voctest.txt', 'w')
#数据集存放路径
Annotations = 'VOCdevkit//VOC2007//Annotations//'
#方法用于返回指定的文件夹包含的文件或文件夹的名字的列表
#['000015.xml', '000018.xml', '000022.xml', '000025.xml', '000027.xml', '000028.xml']
xml_files = os.listdir(Annotations)
random.shuffle(xml_files)  # 打乱数据集
train_num = int(len(xml_files) * 0.7)  # 训练集数量
#['000025.xml', '000027.xml', '000018.xml', '000028.xml']
train_lists = xml_files[:train_num]   # 训练列表
#['000015.xml', '000022.xml']
test_lists = xml_files[train_num:]    # 测试列表

#函数的功能:读取xml文件信息,
# 在"object"目录下查看"difficult"值是否为1,
# 若不为1则在名为"obj_struct"的字典中存入"bbox"和"name"的信息,
# 再将这个字典作为名为"objects"的列表的元素,最终输出这个列表。
# 所以这个名为"objects"的列表中的每一个元素都是一个字典。

#'VOCdevkit//VOC2007//Annotations//000027.xml'
def parse_rec(filename):  # 输入xml文件名
    tree = ET.parse(filename) # 生成一个总目录名为tree
    objects = []
    for obj in tree.findall('object'):  ## 在总目录下找到'object'目录 名为 obj
        obj_struct = {} #obi_struct={'name':'cow'}
        difficult = int(obj.find('difficult').text) #若此目录没有下一层目录则用find而不用findall
        if difficult == 1:  # 若为1则跳过本次循环
            continue
        obj_struct['name'] = obj.find('name').text #obj_struct['name']='cat'
        bbox = obj.find('bndbox')
        #当前类别预测框的信息
        obj_struct['bbox'] = [int(float(bbox.find('xmin').text)),
                              int(float(bbox.find('ymin').text)),
                              int(float(bbox.find('xmax').text)),
                              int(float(bbox.find('ymax').text))]
        objects.append(obj_struct)

    return objects

#在划分好训练集以及测试集后,
# 生成voctrain.txt以及voctest.txt,
# 并且输入以下信息: 图片文件名(带后缀),bbox的四个角坐标,
# 类名代号(数字0~19),这四个信息,
# 图片内的信息都在一行内输入完成,所以有多少张图片文件就有多少行。
def write_txt():
    count = 0
    #文件名称 train_list= '000025.xml'
    for train_list in train_lists: # 生成训练集txt
        count += 1
        # image_name='000025.jpg',图片名称
        image_name = train_list.split('.')[0] + '.jpg'  # 图片文件名
        #result= [{'name': 'cow', 'bbox': [2, 84, 59, 248]}, {'name': 'cow', 'bbox': [68, 115, 233, 279]}, {'name': 'cow', 'bbox': [64, 173, 377, 373]}, {'name': 'person', 'bbox': [320, 2, 496, 375]}, {'name': 'person', 'bbox': [221, 4, 341, 374]}, {'name': 'person', 'bbox': [135, 14, 220, 148]}, {'name': 'cow', 'bbox': [69, 43, 156, 177]}, {'name': 'person', 'bbox': [58, 54, 104, 139]}, {'name': 'person', 'bbox': [279, 1, 331, 86]}, {'name': 'person', 'bbox': [320, 22, 344, 96]}]
        results = parse_rec(Annotations + train_list) #Annotations + train_list是xml文件名
        if len(results) == 0:
            print(train_list)
            continue
        #000025.jpg
        train_set.write(image_name)
        #result={'name': 'cow', 'bbox': [2, 84, 59, 248]}
        for result in results:
            class_name = result['name'] #class_name=cow,20种类别
            bbox = result['bbox']  #bbox=[2,84,59,248]
            class_name = VOC_CLASSES.index(class_name) #class_name = 9
            train_set.write(' ' + str(bbox[0]) +
                            ' ' + str(bbox[1]) +
                            ' ' + str(bbox[2]) +
                            ' ' + str(bbox[3]) +
                            ' ' + str(class_name))
        train_set.write('\n')
    train_set.close()

    for test_list in test_lists:   # 生成测试集txt
        count += 1
        image_name = test_list.split('.')[0] + '.jpg'  # 图片文件名
        results = parse_rec(Annotations + test_list) #Annotations + train_list是xml文件名
        if len(results) == 0:
            print(test_list)
            continue
        test_set.write(image_name)
        for result in results:
            class_name = result['name']
            bbox = result['bbox']
            class_name = VOC_CLASSES.index(class_name)
            test_set.write(' ' + str(bbox[0]) +
                            ' ' + str(bbox[1]) +
                            ' ' + str(bbox[2]) +
                            ' ' + str(bbox[3]) +
                            ' ' + str(class_name))
        test_set.write('\n')
    test_set.close()


if __name__ == '__main__':
    write_txt()

YOLO V1代码_第3张图片

最终的部分结果如上图所示,一张图片可能含有多个框所以按照顺序依次排列,以000018.jpg为例30,31,358,279 为bbox信息,11为类别信息按照定义的元组11代表dog,从下图可以证实代码的正确性。

YOLO V1代码_第4张图片

 4、 文件

此文件主要任务就是根据txt文件内的信息制作ground truth,并且还会进行一定的数据增强。最终输出一个7*7*30的张量。在(train.py中会用到)

以下为文件的主要结构

class yoloDataset(Dataset):
    def __init__(self, img_root, list_file, train, transform):
    def __getitem__(self, idx):
    def __len__(self):
    def encoder(self, boxes, labels):  # 输出ground truth (一个7*7*30的张量)
    # 以下都是数据增强操作
    def random_flip(self, im, boxes): # 随机翻转
    def randomScale(self, bgr, boxes): # 随机伸缩变换
    def randomBlur(self, bgr): # 随机模糊处理
    def RandomBrightness(self, bgr): # 随机调整图片亮度
    def randomShift(self, bgr, boxes, labels): # 平移变换

首先是 __init__ 部分,大致的操作为:逐行读取生成的文本文件的内容,然后对其进行分类,将信息保存在fnames,boxes,labels三个列表中,以下为需要用到的函数

file_txt = open(list_file)     # os模块函数,用于打开文件
lines = file_txt.readlines()   # 读取txt文件每一行
splited = line.strip().split()  
'''
strip()  # 移除首位的换行符号
split()  # 以空格为分界线,将所有元素组成一个列表
'''
fnames = splited[0]
x_max = splited[1 + 5*i]  # i 表示bbox的个数,隔5个数,即为下一个bbox的横坐标
y_max = splited[2 + 5*i]
x_min = splited[3 + 5*i]
y_min = splited[4 + 5*i]

下面是完整的def __init__部分:

def __init__(self, img_root, list_file, train, transform):   # list_file为txt文件  img_root为图片路径
    self.root = img_root
    self.train = train
    self.transform = transform
    # 后续要提取txt文件信息,分类后装入以下三个列表
    self.fnames = []
    self.boxes = []
    self.labels = []

    self.S = 7   # YOLOV1
    self.B = 2   # 相关
    self.C = CLASS_NUM  # 参数
    self.mean = (123, 117, 104)  # RGB

    #首先是 __init__ 部分,大致的操作为:逐行读取生成的文本文件的内容,
    # 然后对其进行分类,将信息保存在fnames,boxes,labels三个列表中,
    # 以下为需要用到的函数

    file_txt = open(list_file)
    lines = file_txt.readlines()   # 读取txt文件每一行
    for line in lines:   # 逐行开始操作
        splited = line.strip().split() # 移除首位的换行符号再生成一张列表
        self.fnames.append(splited[0])  # 存储图片的名字
        num_boxes = (len(splited) - 1) // 5  # 每一幅图片里面有多少个bbox
        box = []
        label = []
        for i in range(num_boxes): # bbox四个角的坐标
            x = float(splited[1 + 5 * i])  #隔5个数,即为下一个bbox的横坐标
            y = float(splited[2 + 5 * i])
            x2 = float(splited[3 + 5 * i])
            y2 = float(splited[4 + 5 * i])
            c = splited[5 + 5 * i]  # 代表物体的类别,即是20种物体里面的哪一种  值域 0-19
            box.append([x, y, x2, y2])
            label.append(int(c))
        self.boxes.append(torch.Tensor(box))
        self.labels.append(torch.LongTensor(label))
    self.num_samples = len(self.boxes)

 其次是 __getitem__ 部分,此为python类中专有的方法,其功能是按照索引获取值,大致的操作为:先对图片进行一些数据增强,再对坐标进行归一化处理(归一化处理是横坐标除以宽度,纵坐标除以长度),最后通过 encoder 方法得到ground truth。最后对每一张图片都返回预处理后的 图片以及他的ground truth。

def __getitem__(self, idx):
    fname = self.fnames[idx]
    img = cv2.imread(os.path.join(self.root + fname))
    boxes = self.boxes[idx].clone()
    labels = self.labels[idx].clone()
    if self.train:  # 数据增强里面的各种变换用torch自带的transform是做不到的,因为对图片进行旋转、随即裁剪等会造成bbox的坐标也会发生变化,所以需要自己来定义数据增强
        img, boxes = self.random_flip(img, boxes)
        img, boxes = self.randomScale(img, boxes)
        img = self.randomBlur(img)
        img = self.RandomBrightness(img)
        # img = self.RandomHue(img)
        # img = self.RandomSaturation(img)
        img, boxes, labels = self.randomShift(img, boxes, labels)
        # img, boxes, labels = self.randomCrop(img, boxes, labels)
    h, w, _ = img.shape
    boxes /= torch.Tensor([w, h, w, h]).expand_as(boxes)  # 坐标归一化处理,为了方便训练(归一化处理是横坐标除以宽度,纵坐标除以长度)
    img = self.BGR2RGB(img)  # because pytorch pretrained model use RGB
    img = self.subMean(img, self.mean)  # 减去均值
    img = cv2.resize(img, (self.image_size, self.image_size))  # 将所有图片都resize到指定大小
    target = self.encoder(boxes, labels)  # 将图片标签编码到7x7*30的向量

    for t in self.transform:
        img = t(img)

    return img, target

最后是encoder方法,它需要输入bbox以及labels信息最终输出一个7*7*30的tensor作为ground truth。值得注意的是bbox的坐标信息中的 x 与 y 含义是左上角顶点坐标相对与中心点坐标(xc,yc)的偏移是一个小于1的数,因此在进行预测的时候还需要进行解码。(这也是比较重要的一部分)

-----------------------------------------------

5、文件

   文件的任务如文件名所示,即定义损失函数供训练时使用。本文件结构如下

YOLO V1代码_第5张图片

 __init__ 部分没什么好说的,定义一些必要的参数。    

forward :在创建对象时会自动调用,这部分计算全部损失函数,此函数会自动进行前向传播因为此类继承了(nn.Module),所以去需要在对象中传入两个参数(对应train.py),下面解释以下代码的逻辑:传入的两个参数格式为(batch_size*7*7*30)的张量,前者将图片输入神经网络得到的输出值,后者就是上面制作的target也就是ground truth。需要提取ground truth与pred_target的bbox信息,置信度信息以及类别信息,求取损失函数。可以参考此博客
目标检测模型YOLO-V1损失函数详解

def forward(self, pred_tensor, target_tensor):
    '''
    pred_tensor: (tensor) size(batchsize,7,7,30)
    target_tensor: (tensor) size(batchsize,7,7,30) --- ground truth
    '''
    N = pred_tensor.size()[0]  # batchsize
    coo_mask = target_tensor[:, :, :, 4] > 0  # 具有目标标签的索引值 true batchsize*7*7
    noo_mask = target_tensor[:, :, :, 4] == 0  # 不具有目标的标签索引值 false batchsize*7*7
    coo_mask = coo_mask.unsqueeze(-1).expand_as(target_tensor)  # 得到含物体的坐标等信息,复制粘贴 batchsize*7*7*30
    noo_mask = noo_mask.unsqueeze(-1).expand_as(target_tensor)  # 得到不含物体的坐标等信息 batchsize*7*7*30

    coo_pred = pred_tensor[coo_mask].view(-1, int(CLASS_NUM + 10))  # view类似于reshape
    box_pred = coo_pred[:, :10].contiguous().view(-1, 5)  # 塑造成X行5列(-1表示自动计算),一个box包含5个值
    class_pred = coo_pred[:, 10:]  # [n_coord, 20]

    coo_target = target_tensor[coo_mask].view(-1, int(CLASS_NUM + 10))
    box_target = coo_target[:, :10].contiguous().view(-1, 5)
    class_target = coo_target[:, 10:]

    # 不包含物体grid ceil的置信度损失
    noo_pred = pred_tensor[noo_mask].view(-1, int(CLASS_NUM + 10))
    noo_target = target_tensor[noo_mask].view(-1, int(CLASS_NUM + 10))
    noo_pred_mask = torch.ByteTensor(noo_pred.size()).bool()
    noo_pred_mask.zero_()
    noo_pred_mask[:, 4] = 1
    noo_pred_mask[:, 9] = 1
    noo_pred_c = noo_pred[noo_pred_mask]  # noo pred只需要计算 c 的损失 size[-1,2]
    noo_target_c = noo_target[noo_pred_mask]
    nooobj_loss = F.mse_loss(noo_pred_c, noo_target_c, size_average=False)  # 均方误差

    # compute contain obj loss
    coo_response_mask = torch.ByteTensor(box_target.size()).bool()  # ByteTensor 构建Byte类型的tensor元素全为0
    coo_response_mask.zero_()  # 全部元素置False                            bool:将其元素转变为布尔值

    no_coo_response_mask = torch.ByteTensor(box_target.size()).bool()  # ByteTensor 构建Byte类型的tensor元素全为0
    no_coo_response_mask.zero_()  # 全部元素置False                            bool:将其元素转变为布尔值

    box_target_iou = torch.zeros(box_target.size())

    # box1 = 预测框  box2 = ground truth
    for i in range(0, box_target.size()[0], 2):  # box_target.size()[0]:有多少bbox,并且一次取两个bbox
        box1 = box_pred[i:i + 2]  # 第一个grid ceil对应的两个bbox
        box1_xyxy = Variable(torch.FloatTensor(box1.size()))
        box1_xyxy[:, :2] = box1[:, :2] / float(self.S) - 0.5 * box1[:, 2:4]  # 原本(xc,yc)为7*7 所以要除以7
        box1_xyxy[:, 2:4] = box1[:, :2] / float(self.S) + 0.5 * box1[:, 2:4]
        box2 = box_target[i].view(-1, 5)
        box2_xyxy = Variable(torch.FloatTensor(box2.size()))
        box2_xyxy[:, :2] = box2[:, :2] / float(self.S) - 0.5 * box2[:, 2:4]
        box2_xyxy[:, 2:4] = box2[:, :2] / float(self.S) + 0.5 * box2[:, 2:4]
        iou = self.compute_iou(box1_xyxy[:, :4], box2_xyxy[:, :4])
        max_iou, max_index = iou.max(0)
        max_index = max_index.data
        coo_response_mask[i + max_index] = 1  # IOU最大的bbox
        no_coo_response_mask[i + 1 - max_index] = 1  # 舍去的bbox
        # confidence score = predicted box 与 the ground truth 的 IOU
        box_target_iou[i + max_index, torch.LongTensor([4])] = max_iou.data

    box_target_iou = Variable(box_target_iou)
    # 置信度误差(含物体的grid ceil的两个bbox与ground truth的IOU较大的一方)
    box_pred_response = box_pred[coo_response_mask].view(-1, 5)
    box_target_response_iou = box_target_iou[coo_response_mask].view(-1, 5)
    # IOU较小的一方
    no_box_pred_response = box_pred[no_coo_response_mask].view(-1, 5)
    no_box_target_response_iou = box_target_iou[no_coo_response_mask].view(-1, 5)
    no_box_target_response_iou[:, 4] = 0  # 保险起见置0(其实原本就是0)

    box_target_response = box_target[coo_response_mask].view(-1, 5)

    # 含物体grid ceil中IOU较大的bbox置信度损失
    contain_loss = F.mse_loss(box_pred_response[:, 4], box_target_response_iou[:, 4], size_average=False)
    # 含物体grid ceil中舍去的bbox损失
    no_contain_loss = F.mse_loss(no_box_pred_response[:, 4], no_box_target_response_iou[:, 4], size_average=False)
    # bbox坐标损失
    loc_loss = F.mse_loss(box_pred_response[:, :2], box_target_response[:, :2], size_average=False) + F.mse_loss(
        torch.sqrt(box_pred_response[:, 2:4]), torch.sqrt(box_target_response[:, 2:4]), size_average=False)

    # 类别损失
    class_loss = F.mse_loss(class_pred, class_target, size_average=False)

    return (self.l_coord * loc_loss + contain_loss + self.l_noobj * (nooobj_loss + no_contain_loss) + class_loss) / N

解释一下部分代码的含义及其用法:

unsqueeze():增加一个维度

YOLO V1代码_第6张图片

expand_as(tensor):将原本的张量扩充,一般是将通道数扩充,扩充的部分就是将原来的部分复制粘贴。

YOLO V1代码_第7张图片

coo_pred = pred_tensor[coo_mask].view(-1, 30):其中coo_mask必须为一个布尔值张量,它与pred_tensor的size一致,若coo_mask的某一位置为True则提取pred_tensor对应位置的信息。而view就相当于resize用于改变张量的形状。

noo_pred_mask = torch.ByteTensor(noo_pred.size()).bool():创建一个Byte类型的张量,形状与noo_pred一致,然后将所有的参数转化为布尔值。
noo_pred_mask.zero_():将参数全部置0,也就是False
compute_iou :输入两个box,输出IOU。IOU指的时交并比,即(交集/并集)。计算公式:
iou = inter / (area1 + area2 - inter)

YOLO V1代码_第8张图片

上文提到过,bbox的四个信息为(xc,yc,w,h),(xc,yc)是基于7*7网格的,而(w,h)是基于1*1的网格。而我们计算IOU需要知道bbox的四个顶点的坐标(相对于grid ceil的坐标而不是坐标轴),通过以下公式进行换算。参照对应代码。

YOLO V1代码_第9张图片

 计算完IOU之后需要在一个grid ceil中挑选出IOU最大的bbox用于计算置信度损失,另外一个就舍去。然后再去计算坐标损失,置信度损失(含物体grid ceil和不含物体grid ceil)还有类别损失,最后相加即可。

def compute_iou(self, box1, box2):  # box1(2,4)  box2(1,4)
    N = box1.size(0)  # 2
    M = box2.size(0)  # 1

    lt = torch.max(  # 返回张量所有元素的最大值
        # [N,2] -> [N,1,2] -> [N,M,2]
        box1[:, :2].unsqueeze(1).expand(N, M, 2),
        # [M,2] -> [1,M,2] -> [N,M,2]
        box2[:, :2].unsqueeze(0).expand(N, M, 2),
    )

    rb = torch.min(
        # [N,2] -> [N,1,2] -> [N,M,2]
        box1[:, 2:].unsqueeze(1).expand(N, M, 2),
        # [M,2] -> [1,M,2] -> [N,M,2]
        box2[:, 2:].unsqueeze(0).expand(N, M, 2),
    )

    wh = rb - lt  # [N,M,2]
    wh[wh < 0] = 0  # clip at 0
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]  重复面积

    area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])  # [N,]
    area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])  # [M,]
    area1 = area1.unsqueeze(1).expand_as(inter)  # [N,] -> [N,1] -> [N,M]
    area2 = area2.unsqueeze(0).expand_as(inter)  # [M,] -> [1,M] -> [N,M]

    iou = inter / (area1 + area2 - inter)
    return iou  # [2,1]

6、文件

先说明以下代码的结构,一共分为三个类和一个函数。

首先是  Bottleneck 类 :定义一个基本块,会根据参数downsample的取值不同可以变为Conv Block和Identity Block。
class Bottleneck(nn.Module):  # 定义基本块
    def __init__(self, in_channel, out_channel, stride, downsample):
        super(Bottleneck, self).__init__()
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.in_channel = in_channel
        self.out_channel = out_channel
        self.bottleneck = Sequential(

            Conv2d(in_channel, out_channel, kernel_size=1, stride=stride[0], padding=0, bias=False),
            BatchNorm2d(out_channel),
            ReLU(inplace=True),

            Conv2d(out_channel, out_channel, kernel_size=3, stride=stride[1], padding=1, bias=False),
            BatchNorm2d(out_channel),
            ReLU(inplace=True),

            Conv2d(out_channel, out_channel * 4, kernel_size=1, stride=stride[2], padding=0, bias=False),
            BatchNorm2d(out_channel * 4),
        )
        if self.downsample is False:  # 如果 downsample = True则为Conv_Block 为False为Identity_Block
            self.shortcut = Sequential()
        else:
            self.shortcut = Sequential(
                Conv2d(self.in_channel, self.out_channel * 4, kernel_size=1, stride=stride[0], bias=False),
                BatchNorm2d(self.out_channel * 4)
            )

    def forward(self, x):
        out = self.bottleneck(x)
        out += self.shortcut(x)
        out = self.relu(out)
        return out
其次是  output_net 类 :配合后续的代码将网络的输出调整为7*7*30的格式
class output_net(nn.Module):
    # no expansion
    # dilation = 2
    # type B use 1x1 conv
    expansion = 1

    def __init__(self, in_planes, planes, stride=1, block_type='A'):
        super(output_net, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=2, bias=False, dilation=2)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
        self.downsample = nn.Sequential()
        self.relu = nn.ReLU(inplace=True)
        if stride != 1 or in_planes != self.expansion * planes or block_type == 'B':
            self.downsample = nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    self.expansion * planes,
                    kernel_size=1,
                    stride=stride,
                    bias=False),
                nn.BatchNorm2d(self.expansion * planes))

    def forward(self, x):
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out += self.downsample(x)
        out = self.relu(out)
        return out
最后是  ResNet 类 :这个类内置了两个主要的方法分别是_make_layer以及_make_output_layer,前者定义的ResNet50的主要部分,后者将前者输出的张量转化为YOLOV1损失函数需要的张量格式(7*7*30)
class ResNet50(nn.Module):
    def __init__(self, block):
        super(ResNet50, self).__init__()
        self.block = block
        self.layer0 = Sequential(
            Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),
            BatchNorm2d(64),
            ReLU(inplace=True),
            MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self.layer1 = self.make_layer(self.block, channel=[64, 64], stride1=[1, 1, 1], stride2=[1, 1, 1], n_re=3)
        self.layer2 = self.make_layer(self.block, channel=[256, 128], stride1=[2, 1, 1], stride2=[1, 1, 1], n_re=4)
        self.layer3 = self.make_layer(self.block, channel=[512, 256], stride1=[2, 1, 1], stride2=[1, 1, 1], n_re=6)
        self.layer4 = self.make_layer(self.block, channel=[1024, 512], stride1=[2, 1, 1], stride2=[1, 1, 1], n_re=3)
        self.layer5 = self._make_output_layer(in_channels=2048)
        self.avgpool = nn.AvgPool2d(2)  # kernel_size = 2  , stride = 2
        self.conv_end = nn.Conv2d(256, int(CLASS_NUM + 10), kernel_size=3, stride=1, padding=1, bias=False)
        self.bn_end = nn.BatchNorm2d(int(CLASS_NUM + 10))

    def make_layer(self, block, channel, stride1, stride2, n_re):
        layers = []
        for num_layer in range(0, n_re):
            if num_layer == 0:
                layers.append(block(channel[0], channel[1], stride1, downsample=True))
            else:
                layers.append(block(channel[1]*4, channel[1], stride2, downsample=False))
        return Sequential(*layers)

    def _make_output_layer(self, in_channels):
        layers = []
        layers.append(
            output_net(
                in_planes=in_channels,
                planes=256,
                block_type='B'))
        layers.append(
            output_net(
                in_planes=256,
                planes=256,
                block_type='A'))
        layers.append(
            output_net(
                in_planes=256,
                planes=256,
                block_type='A'))
        return nn.Sequential(*layers)

    def forward(self, x):
        # print(x.shape) # 3*448*448
        out = self.layer0(x)
        # print(out.shape) # 64*112*112
        out = self.layer1(out)
        # print(out.shape)  # 256*112*112
        out = self.layer2(out)
        # print(out.shape) # 512*56*56
        out = self.layer3(out)
        # print(out.shape) # 1024*28*28
        out = self.layer4(out)  # 2048*14*14
        out = self.layer5(out)  # batch_size*256*14*14
        out = self.avgpool(out)  # batch_size*256*7*7
        out = self.conv_end(out)  # batch_size*30*7*7
        out = self.bn_end(out)
        out = torch.sigmoid(out)
        out = out.permute(0, 2, 3, 1)  # bitch_size*7*7*30
        return out

最后通过 resnet50 函数 return网络结构。在后续训练的时候只需要调用这个函数即可。

7、文件

1、要将数据集用于训练,首先得将数据集打包为训练集以及测试集,此操作之前已经定义过,为yoloData.py内的yoloDataset类以及DaraLoader类实现。

train_dataset = yoloDataset(img_root=file_root, list_file='voctrain.txt', train=True, transform=[transforms.ToTensor()])
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_dataset = yoloDataset(img_root=file_root, list_file='voctest.txt', train=False, transform=[transforms.ToTensor()])
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=0)

 2、根据定义的网络导入权重参数,因之前定义的网络后面一部分与ResNet50的结构略有差异,所以并不能完全使用使用torchvision的models库中的resnet50导入权重参数。需要对其权重参数进行一定的筛选。

        权重参数导入方法:自己定义的网络以及models库内的网络各自创建一个对象。接着使用state_dict()导入各自的权重参数。网络结构相同的部分将new_state_dict的值赋给op。但是如果自己定义的网络结构的键值与torch自带的库不一直的话,导入权重参数会稍微麻烦一点。这里给出了一种解决办法,参考代码。
 

YOLO V1代码_第10张图片

注意:

state_dict():返回一个类(若直接输出其格式类似与字典)

enumerate: for循环中经常用到,既可以遍历元素又可以遍历索引。

3、定义损失函数优化器并开始训练,网络结构,损失函数以及train_loader返回的图片以及target最后通过torch.save保存模型参数。

criterion = yoloLoss(7, 2, 5, 0.5)
criterion = criterion.to(device)
net.train()  # 训练前需要加入的语句

params = []  # 里面存字典
params_dict = dict(net.named_parameters()) # 返回各层中key(只包含weight and bias) and value
for key, value in params_dict.items():
    params += [{'params': [value], 'lr':learning_rate}]  # value和学习率相加

optimizer = torch.optim.SGD(    # 定义优化器  “随机梯度下降”
    params,   # net.parameters() 为什么不用这个???
    lr=learning_rate,
    momentum=0.9,   # 即更新的时候在一定程度上保留之前更新的方向  可以在一定程度上增加稳定性,从而学习地更快
    weight_decay=5e-4)     # L2正则化理论中出现的概念
# torch.multiprocessing.freeze_support()  # 多进程相关 猜测是使用多显卡训练需要

for epoch in range(num_epochs):
    net.train()
    if epoch == 60:
        learning_rate = 0.0001
    if epoch == 80:
        learning_rate = 0.00001
    for param_group in optimizer.param_groups:   # 其中的元素是2个字典;optimizer.param_groups[0]: 长度为6的字典,包括[‘amsgrad’, ‘params’, ‘lr’, ‘betas’, ‘weight_decay’, ‘eps’]这6个参数;
                                                # optimizer.param_groups[1]: 好像是表示优化器的状态的一个字典;
        param_group['lr'] = learning_rate      # 更改全部的学习率
    print('\n\nStarting epoch %d / %d' % (epoch + 1, num_epochs))
    print('Learning Rate for this epoch: {}'.format(learning_rate))

    total_loss = 0.
    for i, (images, target) in enumerate(train_loader):
        images, target = images.to(device), target.to(device)
        pred = net(images)
        loss = criterion(pred, target)
        total_loss += loss.item()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if (i + 1) % 5 == 0:
            print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f' % (epoch +1, num_epochs,
                                                                                 i + 1, len(train_loader), loss.item(), total_loss / (i + 1)))
    validation_loss = 0.0
    net.eval()
    for i, (images, target) in enumerate(test_loader):  # 导入dataloader 说明开始训练了  enumerate 建立一个迭代序列
        images, target = images.to(device), target.to(device)
        pred = net(images)    # 将图片输入
        loss = criterion(pred, target)
        validation_loss += loss.item()   # 累加loss值  (固定搭配)
    validation_loss /= len(test_loader)  # 计算平均loss

    best_test_loss = validation_loss
    print('get best test loss %.5f' % best_test_loss)
    torch.save(net.state_dict(), 'yolo.pth')

8、文件

 此为预测代码,执行此代码可以通过输入图片得到识别的结果。从上文可知,将图片输入神经网络会输出一个7*7*30的张量,我们需要对其进行分析得到很多个bbox包括坐标信息以及对应的物体类别信息,最终进行NMS非极大值抑制对bbox进行筛选得到最终结果。

以下为NMS的步骤:

(1)对于类别1, 从概率最大的bbox F开始,分别判断A、B、C、D、E与F的IOU是否大于设定的阈值。

(2) 假设B、D与F的重叠度超过阈值,那么就扔掉B、D(将其置信度置0),然后保留F。

(3) 从剩下的矩形框A、C、E中,选择概率最大的E,然后判断A、C与E的重叠度,重叠度大于一定的阈值,那么就扔掉;并标记E是我们保留下来的第二个矩形框。

(4) 重复这个过程,找到此类别所有被保留下来的矩形框。

(5) 对于类别2,类别3等等...都要重复以上4个步骤。

首先需要导入模型以及参数,并且设置好有关NMS的两个参数:置信度以及IOU最大值。然后就可以开始预测了。首先需要通过opencv读取图片并且将其resize为448*448的RGB图像,将其进行均值处理后输入神经网络得到7*7*30的张量。

然后运行 decode  方法:因为一个grid ceil只预测一个物体,而一个grid ceil生成两个bbox。这里对grid ceil进行以下操作。

1、选择置信度较高的bbox。

2、选择20种类别概率中的最大者作为这个grid ceil预测的类别。

3、置信度乘以物体类别概率作为物体最终的概率。

最终输入一个7*7*6的张量,7*7代表grid ceil  。6=bbox的4个坐标信息+类别概率+类别代号

最后运行  NMS 方法对bbox进行筛选:因为bbox的4个坐标信息为(xc,yc,w,h)需要将其转化为(x,y,w,h)后才能进行非极大值抑制处理。

# target 7*7*30  值域为0-1
class Pred():
    def __init__(self, model, img_root):
        self.model = model
        self.img_root = img_root

    def result(self):
        img = cv2.imread(self.img_root)
        h, w, _ = img.shape
        print(h, w)
        image = cv2.resize(img, (448, 448))
        img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mean = (123, 117, 104)  # RGB
        img = img - np.array(mean, dtype=np.float32)
        transform = ToTensor()
        img = transform(img)
        img = img.unsqueeze(0)  # 输入要求是4维的
        Result = self.model(img)   # 1*7*7*30
        bbox = self.Decode(Result)
        bboxes = self.NMS(bbox)    # n*6   bbox坐标是基于7*7网格需要将其转换成448
        if len(bboxes) == 0:
            print("未识别到任何物体")
            print("尝试减小 confident 以及 iou_con")
            print("也可能是由于训练不充分,可在训练时将epoch增大")        
        for i in range(0, len(bboxes)):    # bbox坐标将其转换为原图像的分辨率
            bboxes[i][0] = bboxes[i][0] * 64
            bboxes[i][1] = bboxes[i][1] * 64
            bboxes[i][2] = bboxes[i][2] * 64
            bboxes[i][3] = bboxes[i][3] * 64

            x1 = bboxes[i][0].item()    # 后面加item()是因为画框时输入的数据不可一味tensor类型
            x2 = bboxes[i][1].item()
            y1 = bboxes[i][2].item()
            y2 = bboxes[i][3].item()
            class_name = bboxes[i][5].item()
            print(x1, x2, y1, y2, VOC_CLASSES[int(class_name)])

            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (144, 144, 255))   # 画框

        cv2.imshow('img', image)

        cv2.waitKey(0)

    def Decode(self, result):  # x -> 1**7*30
        result = result.squeeze()   # 7*7*30
        grid_ceil1 = result[:, :, 4].unsqueeze(2)  # 7*7*1
        grid_ceil2 = result[:, :, 9].unsqueeze(2)
        grid_ceil_con = torch.cat((grid_ceil1, grid_ceil2), 2)  # 7*7*2
        grid_ceil_con, grid_ceil_index = grid_ceil_con.max(2)    # 按照第二个维度求最大值  7*7   一个grid ceil两个bbox,两个confidence
        class_p, class_index = result[:, :, 10:].max(2)   # size -> 7*7   找出单个grid ceil预测的物体类别最大者
        class_confidence = class_p * grid_ceil_con   # 7*7   真实的类别概率
        bbox_info = torch.zeros(7, 7, 6)
        for i in range(0, 7):
            for j in range(0, 7):
                bbox_index = grid_ceil_index[i, j]
                bbox_info[i, j, :5] = result[i, j, (bbox_index * 5):(bbox_index+1) * 5]   # 删选bbox 0-5 或者5-10
        bbox_info[:, :, 4] = class_confidence
        bbox_info[:, :, 5] = class_index
        print(bbox_info[1, 5, :])
        return bbox_info  # 7*7*6    6 = bbox4个信息+类别概率+类别代号

    def NMS(self, bbox, iou_con=iou_con):
        for i in range(0, 7):
            for j in range(0, 7):
                # xc = bbox[i, j, 0]        # 目前bbox的四个坐标是以grid ceil的左上角为坐标原点 而且单位不一致
                # yc = bbox[i, j, 1]         # (xc,yc) 单位= 7*7   (w,h) 单位= 1*1
                # w = bbox[i, j, 2] * 7
                # h = bbox[i, j, 3] * 7
                # Xc = i + xc
                # Yc = j + yc
                # xmin = Xc - w/2     # 计算bbox四个顶点的坐标(以整张图片的左上角为坐标原点)单位7*7
                # xmax = Xc + w/2
                # ymin = Yc - h/2
                # ymax = Yc + h/2     # 更新bbox参数  xmin and ymin的值有可能小于0
                xmin = j + bbox[i, j, 0] - bbox[i, j, 2] * 7 / 2     # xmin
                xmax = j + bbox[i, j, 0] + bbox[i, j, 2] * 7 / 2     # xmax
                ymin = i + bbox[i, j, 1] - bbox[i, j, 3] * 7 / 2     # ymin
                ymax = i + bbox[i, j, 1] + bbox[i, j, 3] * 7 / 2     # ymax

                bbox[i, j, 0] = xmin
                bbox[i, j, 1] = xmax
                bbox[i, j, 2] = ymin
                bbox[i, j, 3] = ymax

        bbox = bbox.view(-1, 6)   # 49*6
        bboxes = []
        ori_class_index = bbox[:, 5]
        class_index, class_order = ori_class_index.sort(dim=0, descending=False)
        class_index = class_index.tolist()   # 从0开始排序到7
        bbox = bbox[class_order, :]  # 更改bbox排列顺序
        a = 0
        for i in range(0, CLASS_NUM):
            num = class_index.count(i)
            if num == 0:
                continue
            x = bbox[a:a+num, :]   # 提取同一类别的所有信息
            score = x[:, 4]
            score_index, score_order = score.sort(dim=0, descending=True)
            y = x[score_order, :]   # 同一种类别按照置信度排序
            if y[0, 4] >= confident:    # 物体类别的最大置信度大于给定值才能继续删选bbox,否则丢弃全部bbox
                for k in range(0, num):
                    y_score = y[:, 4]   # 每一次将置信度置零后都重新进行排序,保证排列顺序依照置信度递减
                    _, y_score_order = y_score.sort(dim=0, descending=True)
                    y = y[y_score_order, :]
                    if y[k, 4] > 0:
                        area0 = (y[k, 1] - y[k, 0]) * (y[k, 3] - y[k, 2])
                        for j in range(k+1, num):
                            area1 = (y[j, 1] - y[j, 0]) * (y[j, 3] - y[j, 2])
                            x1 = max(y[k, 0], y[j, 0])
                            x2 = min(y[k, 1], y[j, 1])
                            y1 = max(y[k, 2], y[j, 2])
                            y2 = min(y[k, 3], y[j, 3])
                            w = x2 - x1
                            h = y2 - y1
                            if w < 0 or h < 0:
                                w = 0
                                h = 0
                            inter = w * h
                            iou = inter / (area0 + area1 - inter)
                            # iou大于一定值则认为两个bbox识别了同一物体删除置信度较小的bbox
                            # 同时物体类别概率小于一定值则认为不包含物体
                            if iou >= iou_con or y[j, 4] < confident:
                                y[j, 4] = 0
                for mask in range(0, num):
                    if y[mask, 4] > 0:
                        bboxes.append(y[mask])
            a = num + a
        return bboxes

你可能感兴趣的:(代码理解,目标检测)