YOLOV1详解——Pytorch版

山居秋暝LS

已于 2022-10-16 10:23:36 修改

阅读量4.3k

点赞数 9

分类专栏：计算机视觉 Pytorch 文章标签： pytorch 计算机视觉深度学习

于 2022-08-07 09:51:26 首次发布

本文链接：https://blog.csdn.net/qq_35732321/article/details/126206643

版权

计算机视觉同时被 2 个专栏收录

45 篇文章 3 订阅

订阅专栏

Pytorch

6 篇文章 0 订阅

订阅专栏

YOLOV1详解——Pytorch版

1 YOLOV1

1 YOLOV1

YOLOV1的算法流程。YOLOV1把图像从448x448x3下采样6次得到特征图7x7x30，特征图上的每个网格生成两个框预测物体和物体存在的概率以及20个物体的类别，30=Bx(loc+conf)+cls = 2x(4+1)+20。B的数量为2，是预测框的数量。loc的数量为4，是预测框中心店的偏移和宽高。conf的数量为2，代表每个框与物体的概率。cls代表预测物体的数量。如果真实框落在某个网格内，网格内置信度高得一个框预测真实框。
YOLOV1的优缺点。YOLOV1是end-to-end的模型，模型简单；速度快，预测精度逊色与Faster_RCNN，但是预测速度大幅提升；在不同的数据集上验证证明模型泛化能力强。不需要滑动窗口得到区域框，预测精度相对于其他的模型有所提高。

由于YOLOV1是end没有设置先验框，因此预测精度受的影响且迁移能力差；只有一个特征层，对小物体和群体的小物体预测能力差，这样也意味着有更多的修改空间。

参考代码：https://github.com/abeardear/pytorch-YOLO-v1

'''
数据集：VOC2007
        VOC2007/Annotations  
                00001.xml
                00002.xml
                .....
        VOC2007/JPEGImages   
                00001.jpg
                00002.jpg
                .....
        VOCdevkit/VOC2007/ImageSets/Main
                train.txt
                test.txt
                val.txt
代码流程：
    1. 数据处理
        1.1 data_split_11.py 把数据集划分成训练集train.txt、验证集、测试集
        1.2 data_jpgxlm_12.py 根据生成的数据集取出数据集每个样本的的图片地址和框、类别信息
        1.3 data_label_13.py 数据增强和编码生成标签
    2.训练
        2.1 主干网络、迁移学习 resnet_yolo.py
        2.2 损失函数 yoloLoss.py
        2.3 优化器   
        2.4 导入数据，生成器
        2.5 训练
    3.预测
        3.1 图片输入模型得到预测结果。result = predict_gpu(model,image_name)
        3.2 画图
        
 
        
1.1 data_split_11.py
    功能：获取数据集索引。
    inputs  : xmlpath = r'./VOCdevkit/VOC2007/Annotations'   # 所有数据集的.xml.
    outputs : ./VOCdevkit/VOC2007/ImageSets/Main/train.txt   # 数据集索引
              ./VOCdevkit/VOC2007/ImageSets/Main/test.txt
              ./VOCdevkit/VOC2007/ImageSets/Main/val.txt
              ./VOCdevkit/VOC2007/ImageSets/Main/trainval.txt
    ./VOCdevkit/VOC2007/Annotations:
                00001.xml
                00002.xml
                .....
    ./VOCdevkit/VOC2007/ImageSets/Main
                train.txt
                test.txt
                trainval.txt
                val.txt
    train.txt:
            000025
            000018
            000022
            .....
    1.1.1 获取'./VOCdevkit/VOC2007/Annotations' 下以'.xml'结尾的文件名。
    1.1.2 从总的样本中按照数据集比例抽取样本，得到每个数据集的索引
    1.1.3 不同数据集的存储地址
    1.1.4 遍历样本，根据抽取的样本索引放入不同的数据集中
    1.1.5 关闭文件
    
1.2 data_jpgxlm_12.py 根据生成的数据集取出数据集每个样本的的图片地址和框、类别信息
    inputs  : sets、classes
    outputs : 2007_train.txt、2007_test.txt、2007_val.txt
              2007_train.txt:
                    VOCdevkit/VOC2007/JPEGImages/000022.jpg 68,103,368,283,12 186,44,255,230,14
                    VOCdevkit/VOC2007/JPEGImages/000028.jpg 63,18,374,500,7
                    ......
    1.2.1 遍历每个数据集获取数据集图片名。生成解析文件，用来存储图片地址和框的信息。
    1.2.2 遍历每张图片，解析文件写入图片地址，打开xml文件，读入类别信息和框的信息。
          (1) getroot()-->'object'-->'difficult','name','bndbox'-->cls_id + box
    1.2.3 关闭解析文件

1.3 data_label_13.py  
    inputs  : list_f = r'2007_train.txt'
    outputs :
    1.3.1 读入文件信息，存储图片、框、类别fenb放入fnames、boxes、labels中。
    1.3.2 根据索引遍历每张图片，如果是train模式，进行数据增强。
          （1）随机左右反转。先翻转图片，计算框与左右边的距离，框左右边数值调换位置。
          （2）随机变换尺度。选择随机数放缩宽，框的x轴数值按同比例放缩。
          （3）随机模糊。
          （4）随机改变色度、饱和度、明暗。bgr->hsv->放缩->bgr
          （5）随机平移。随机选平移数值，平移图片。判断框框平移后是否在图片上，如果框平移后存在，平移框。筛选平移后留下的目标。
          （6）随机剪切。生成随机剪切起点和宽高，选取剪切图片上的框和目标，剪切图片。
    1.3.3 图像归一化
    1.3.4 编码
          （1）计算框的中心点和宽高
          （2）计算框再特征图上的位置，计算中心点偏移
          （3）再特征图上真实框对应的位置写入框、置信度、类别信息
    1.3.5 根据需要转换通道

2.1 主干网络、迁移学习 resnet_yolo.py
    inputs  : torch.Size([2, 3, 224, 224])
    outputs : torch.Size([2, 7, 7, 30])
    2.1.1 input --> cnnv(s=2)+bn+re+maxpool(s=2) --> layers*5((s=2) * 3) --> cnnv+bn+sigmoid


2.2 损失函数 yoloLoss.py
    inputs  :
        pred.shape:  [n, 7, 7, 30]
        target.shape:[n, 7, 7, 30]
    outputs :scalar
    2.2.1 有物体的掩码coo_mask[n, 7, 7, 30]，没有物体的掩码noo_mask[n, 7, 7, 30]。
    2.2.2 根据noo_mask得到没有物体的预测负样本noo_pred = pred_tensor[noo_mask].view(-1,30)和物体的真实负样本noo_target = target_tensor[noo_mask].view(-1,30)，提取置信度,计算负样本的损失.
          noo_mask.shape: torch.Size([2, 14, 14, 30])
          pred_tensor.shape: torch.Size([2, 14, 14, 30])
          noo_pred.shape: torch.Size([0, 30])
          noo_target.shape: torch.Size([0, 30])
    2.2.3 根据coo_mask提取预测框、预测类别、真实框、真实类别。根据预测类别和真实类别计算类别损失。
    2.2.4 遍历真实框，根据IoU标记与真实框最匹配的预测框和最优的IoU.提取最匹配的预测框，分别计算框和正样本的损失。
    2.2.5 对各个损失加权，计算总损失。

2.3 训练
    2.3.1 加载模型，设置损失函数、优化器、加载数据、前向传播、向后传播跟新参数。
   
   
3.1 图片输入模型得到预测结果
    inputs  : img
    outputs : boxes,cls_indexs,probs 
    3.1.1 数据预处理+前向传 --> 预测结果pred7x7x30
    3.1.2 预测结果解码 + 筛选
            (1) 先根据置信度筛选满足条件的mask.
            (2) 遍历特征图上每一个网格点和mask，对满足条件的框解码。
            (3) 根据置信度和预测物体的条件概率得到预测概率，根据预测概率对框筛选。
            (4) 对上一步筛选的框进行NMS。
    3.1.3 把预测框映射到原图
    
3.2 画框
    3.2.1 画框
    3.2.2 在框上标记预测物体类别和概率
    
'''

1 数据处理

数据处理主要包括数据集划分、读入xml文件、数据增强三大部分。

首先把数据划分为训练集、验证集和测试集，每种类型的数据集里存储的是图片的名称，比如2002001.jpg,2002002.jpg,2002003.jpg图片，在数据集中是2002001，2002002、2002003。第二步根据数据集读入图片的地址和和真实框的信息。第三部就是根据图片和真实框的信息进行数据增强及编码得到标签。

三部分的代码依次存放在 data_split_11.py 、data_jpgxlm_12.py、 data_label_13.py，点击运行就会得到对应结果。然后点击 train.py 就可以训练模型，根据训练模型的参数就可以点击 predict.py 预测了。

1.1 数据集划分

数据集划分在data_split_11.py文件中。

已知信息：在VOCdevkit/VOC2007/Annotations存储的是每个图片对应的xml信息、训练集、验证集、测试集比例。
输出的是各个数据集的图片名称。

数据集划分流程：

获取’./VOCdevkit/VOC2007/Annotations’ 下以’.xml’结尾的文件名。
从总的样本中按照数据集比例抽取样本，得到每个数据集的索引。
不同数据集的存储地址。
遍历样本，根据抽取的样本索引放入不同的数据集中。
关闭文件。

'''
dataSplit
'''
import os
import random

xml_path = r'./VOCdevkit/VOC2007/Annotations'
base_path = r'./VOCdevkit/VOC2007/ImageSets/Main'

# 1 样本名字
tmp = []
img_names = os.listdir(xml_path)
for i in img_names:
    if i.endswith('.xml'):
        tmp.append(i[:-4])

# 2 数据集划分
trainval_ratio = 0.9
train_ratio = 0.9
N = len(tmp)
trainval_num = int(trainval_ratio*N)
train_num = int(train_ratio*trainval_num)
trainval_idx = random.sample(range(N),trainval_num)
train_idx = random.sample(trainval_idx,train_num)

# 3 数据集的存储地址
ftrainval = open(os.path.join(base_path,'LS_trainval.txt'),'w')
ftrain = open(os.path.join(base_path,'LS_train.txt'),'w')
fval = open(os.path.join(base_path,'LS_val.txt'),'w')
ftest = open(os.path.join(base_path,'LS_test.txt'),'w')

# 4 写入数据
for i in range(N):
    name = tmp[i]+'\n'
    if i in trainval_idx:
        ftrainval.write(name)
        if i in train_idx:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)
        
# 5 关闭文件
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

1.2 读入xml文件

data_jpgxlm_12.py把图片信息和xml信息放在一起。

遍历每个数据集获取数据集图片名。生成解析文件，用来存储图片地址和框的信息。
遍历每张图片，解析文件写入图片地址，打开xml文件，读入类别信息和框的信息。 getroot()–>‘object’–>‘difficult’,‘name’,‘bndbox’–>cls_id + box
关闭解析文件

import xml.etree.ElementTree as ET

sets = [('2007','train'),('2007','val'),('2007','test')]
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle",
           "bus", "car", "cat", "chair", "cow", "diningtable",
           "dog", "horse", "motorbike", "person", "pottedplant",
           "sheep", "sofa", "train", "tvmonitor"]       # len(classes) = 20

def convert_annotation(year,img_id,list_file):
    list_file.write("VOCdevkit/VOC%s/JPEGImages/%s.jpg"%(year,img_id))
    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year,img_id))
    root = ET.parse(in_file).getroot()

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult)==1:
            continue
        cls_id = classes.index(cls)
        xml_box = obj.find('bndbox')
        b = (int(xml_box.find('xmin').text),
             int(xml_box.find('ymin').text),
             int(xml_box.find('xmax').text),
             int(xml_box.find('ymax').text))
        list_file.write(" "+','.join([str(i) for i in b])+','+str(cls_id))

for year ,img_set in sets:
    img_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year ,img_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year,img_set),'w')
    for img_id in img_ids:
        convert_annotation(year,img_id,list_file)
        list_file.write('\n')
    list_file.close()

1.3 数据增强

data_label_13.py对图片数据增强，增加样本，提高模型的泛化能力。

首先把图片名称、框、类别信息分别存储。
生成迭代器，对每个图片进行数据增强。翻转、缩放、模糊、随机变换亮度、随机变换色度、随机变换饱和度、随机平移、随机剪切。
编码。对于增强后的图片，根据图片宽高，获取框在图片的相对位置、去均值、统一图片尺寸、编码。编码的关键在于找到真实框在特征图上的相对位置。先把真实框左上角和右下角坐标转换为中心点和宽高，中心点坐标x特征图宽高后，向下取整在-1就得到真实框在特征图上的位置ij。中心点坐标x特征图宽高-ij得到真实偏移。根据ij输入偏移、宽高、类别信息。

（1）数据增强代码


''' 数据增强 '''

import os
import sys
import cv2
import torch
import random
import os.path
import numpy as np
import matplotlib.pyplot as plt
import torch.utils.data as data
import torchvision.transforms as transforms

class yoloDataset(data.Dataset):
    image_size = 448
    def __init__(self,root,list_file,train,transform):
        print('data init')
        self.root = root    ## img.jpg_path
        self.train = train  ## bool:True，如果是训练模式，就进行数据增强。
        self.transform = transform  # 转置
        self.fnames = []   # 存储图片名  eg: 00014.jpg
        self.boxes = []    # 存放真实框信息
        self.labels = []   # 存放标签
        self.mean = (123,117,104)   # 图片各个通道均值，用于归一化数据，加速训练。


        with open(list_file) as f:
            lines = f.readlines()   # 读取数据 

        # line： /Users/ls/PycharmProjects/YOLOV1_LS/VOCdevkit/VOC2007/JPEGImages/000022.jpg 68,103,368,283,12 186,44,255,230,14
        for line in lines:   
            splited = line.strip().split()
            self.fnames.append(splited[0])  # 保存图片名
            num_boxes = (len(splited)-1)    # 一张图片中真实框的个数
            box = []     # 存储框
            label = []   # 存储标签
            for i in range(1,num_boxes+1): # 遍历每个框
                tmp = [float(j) for j in splited[i].split(',')] # 把真实框油字符变为float类型，并用‘,’隔开。
                box.append(tmp[:4]) 
                label.append(int(tmp[4])+1)
            self.boxes.append(torch.Tensor(box))
            self.labels.append(torch.LongTensor(label))
        self.num_samples = len(self.boxes)

    def __getitem__(self,idx):
        fname = self.fnames[idx] # 用迭代器遍历每张图片
        img = cv2.imread(fname)  # 读取图片 cv2.imread(os.path.join(self.root+fname))
        boxes = self.boxes[idx].clone()
        labels = self.labels[idx].clone()

        if self.train:
            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)
        img = self.subMean(img,self.mean)
        img = cv2.resize(img,(self.image_size,self.image_size))
        target = self.encoder(boxes,labels)
        for t in self.transform:
            img = t(img)
        return  img,target

    def __len__(self):
        return self.num_samples

    def encoder(self,boxes,labels):
        grid_num = 7
        target = torch.zeros((grid_num,grid_num,30))
        wh = boxes[:,2:] - boxes[:,:2]
        cxcy = (boxes[:,2:] + boxes[:,:2])/2
        for i in range(cxcy.size()[0]):
            cxcy_sample = cxcy[i]
            ij = (cxcy_sample*grid_num).ceil()-1
            dxy = cxcy_sample*grid_num-ij
            target[int(ij[1]),int(ij[0]),:2] = target[int(ij[1]),int(ij[0]),5:7] = dxy
            target[int(ij[1]),int(ij[0]),2:4] = target[int(ij[1]),int(ij[0]),7:9] = wh[i]
            target[int(ij[1]),int(ij[0]),4] = target[int(ij[1]),int(ij[0]),9] = 1
            target[int(ij[1]),int(ij[0]),int(labels[i])+9] = 1
        return target

    def BGR2RGB(self,img):
        return cv2.cvtColor(img,cv2.COLOR_BGR2RGB)

    def BGR2HSV(self,img):
        return cv2.cvtColor(img,cv2.COLOR_BGR2HSV)

    def HVS2BGR(self,img):
        return cv2.cvtColor(img,cv2.COLOR_HSV2BGR)

    def RandomBrightness(self,bgr):
        if random.random()<0.5:
            hsv = self.BGR2HSV(bgr)
            h,s,v = cv2.split(hsv)
            v = v.astype(float)
            v *= random.choice([0.5,1.5])
            v = np.clip(v,0,255).astype(hsv.dtype)
            hsv = cv2.merge((h,s,v))
            bgr = self.HVS2BGR(hsv)
        return  bgr

    def RandomSaturation(self,bgr):
        if random.random()<0.5:
            hsv = self.BGR2HSV(bgr)
            h,s,v = cv2.split(hsv)
            s = s.astype(float)
            s *= random.choice([0.5,1.5])
            s = np.clip(s,0,255).astype(hsv.dtype)
            hsv = cv2.merge((h,s,v))
            bgr = self.HVS2BGR(hsv)
        return bgr

    def RandomHue(self,bgr):
        if random.random() < 0.5:
            hsv = self.BGR2HSV(bgr)
            h,s,v = cv2.split(hsv)
            h = h.astype(float)
            h *= random.choice([0.5,1.5])
            h = np.clip(h,0,255).astype(hsv.dtype)
            hsv=cv2.merge((h,s,v))
            bgr = self.HVS2BGR(hsv)
        return bgr

    def randomBlur(self,bgr):
        if random.random() < 0.5:
            bgr = cv2.blur(bgr,(5,5))
        return bgr

    def randomShift(self,bgr,boxes,labels):
        center = (boxes[:,2:]+boxes[:,:2])/2
        if random.random()<0.5:
            height,width,c = bgr.shape
            after_shift_imge = np.zeros((height,width,c),dtype=bgr.dtype)
            after_shift_imge[:,:,:] = (104,117,123)
            shift_x = random.uniform(-width*0.2,width*0.2)
            shift_y = random.uniform(-height*0.2,height*0.2)
            if shift_x>=0 and shift_y>=0:
                after_shift_imge[int(shift_y):,int(shift_x):,:] = bgr[:height-int(shift_y),:width-int(shift_x),:]
            elif shift_x>=0 and shift_y<0:
                after_shift_imge[:height+int(shift_y),int(shift_x):,:] = bgr[-int(shift_y):,:width-int(shift_x),:]
            elif shift_x <0 and shift_y >=0:
                after_shift_imge[int(shift_y):,:width+int(shift_x),:] = bgr[:height-int(shift_y),-int(shift_x):,:]
            elif shift_x<0 and shift_y<0:
                after_shift_imge[:height+int(shift_y),:width+int(shift_x),:] = bgr[-int(shift_y):,-int(shift_x):,:]

            shift_xy = torch.FloatTensor([[int(shift_x),int(shift_y)]]).expand_as(center)
            center = center + shift_xy
            mask1 = (center[:,0]>0)& (center[:,0]>height)
            mask2 = (center[:,1]>0)& (center[:,1]>width)
            mask = (mask1 & mask2).view(-1,1)
            boxes_in = boxes[mask.expand_as(boxes)].view(-1,4)
            if len(boxes_in) == 0:
                return bgr,boxes,labels
            box_shift = torch.FloatTensor([[int(shift_x),int(shift_y),int(shift_x),int(shift_y)]]).expand_as(boxes_in)
            boxes_in = boxes_in+box_shift
            labels_in = labels[mask.view(-1)]
            return after_shift_imge,boxes_in,labels_in
        return bgr,boxes,labels

    def randomScale(self,bgr,boxes):
        if random.random() < 0.5:
            scale = random.uniform(0.8,1.2)
            h,w,c = bgr.shape
            bgr = cv2.resize(bgr,(int(w*scale),h))
            scale_tensor = torch.FloatTensor([[scale,1,scale,1]]).expand_as(boxes)
            boxes = boxes*scale_tensor
            return bgr,boxes
        return bgr,boxes

    def randomCrop(self,bgr,boxes,labels):
        if random.random() < 0.5:
            center = (boxes[:,:2]+boxes[:,2:])/2
            height,width,c = bgr.shape
            h = random.uniform(0.6*height,height)
            w = random.uniform(0.6*width,width)
            x = random.uniform(0,width-w)
            y = random.uniform(0,height-h)
            x,y,h,w = int(x),int(y),int(h),int(w)

            center = center - torch.FloatTensor([[x,y]]).expand_as(center)
            mask1 = (center[:,0]>0) & (center[:,0]<w)
            mask2 = (center[:,1]>0) & (center[:,0]<h)
            mask = (mask1 & mask2).view(-1,1)

            boxes_in = boxes[mask.expand_as(boxes)].view(-1,4)
            if(len(boxes_in)==0):
                return bgr,boxes,labels

            box_shift = torch.FloatTensor([[x,y,x,y]]).expand_as(boxes_in)
            boxes_in = boxes_in-box_shift
            boxes_in[:,0]=boxes_in[:,0].clamp(0,w)
            boxes_in[:,2]=boxes_in[:,2].clamp(0,w)
            boxes_in[:,1]=boxes_in[:,1].clamp(0,h)
            boxes_in[:,3]=boxes_in[:,3].clamp(0,h)

            labels_in = labels[mask.view(-1)]
            img_croped = bgr[y:y+h,x:x+w,:]
            return img_croped,boxes_in,labels_in
        return bgr,boxes,labels

    def subMean(self,bgr,mean):
        mean = np.array(mean,dtype=np.float32)
        bgr = bgr - mean
        return bgr

    def random_flip(self,im,boxes):
        if random.random() < 0.5:
            im_lr = np.fliplr(im).copy()
            h,w,_ = im.shape
            xmin = w - boxes[:,2]
            xmax = w - boxes[:,0]
            boxes[:,0] = xmin
            boxes[:,2] = xmax
            return im_lr,boxes
        return im,boxes

    def random_bright(self,im,delta=16):
        alpha = random.random()
        if alpha > 0.3:
            im = im * alpha + random.randrange(-delta,delta)
            im = im.clip(min=0,max=255).astype(np.uint8)
        return im

if __name__=='__main__':
    from torch.utils.data import DataLoader
    import torchvision.transforms as transforms
    file_root ='/Users/ls/PycharmProjects/YOLOV1_LS'  ## xx.jpg
    list_f = r'/Users/ls/PycharmProjects/YOLOV1_LS/2007_train.txt'
    train_dataset = yoloDataset(root=file_root,list_file=list_f,train=True,transform = [transforms.ToTensor()])
    train_loader = DataLoader(train_dataset,batch_size=1,shuffle=False,num_workers=0)
    train_iter = iter(train_loader)
    # for i in range(5):
    #     img,target = next(train_iter)
    #     print(target[target[...,0]>0])
    for i,(images,target) in enumerate(train_loader):
        print(1111111111111111111111)
        print(target)
        print(images)

（2）以下是数据增强中主要函数的讲解。数据增强中，在对图片操作的同时，关键是也要对真实框作相应的操作。例如对图像进行翻转，目标物体位置发生变化，真实框也要进行同样的翻转，利于保障训练预测结果的准确性。

编码要点

I.编码中输入的框是框在图像中的相对位置，即真实框的左上角和右下角坐标除以图像宽高后的值。因为要在找到特征图上真实框的位置，虽然原图和特征图的尺寸不一样，但是真实框在原图和特征图的相对位置一样，通过相对位置把真实框映射在特征图上。

img.shape:[400,500,3]         # 原图 h:400;w:500
box:[100,120,200,250]         # 真实框的坐标 [x1,y1,x2,y2]
box_img: [100/500，120/400,200/500,250/400]            # 真实框在原图上的相对位置 [x1/w,y1/h,x2/w,y2/h]
feature.shape:[7,7,30]        # 特征图
box_feature: [100/500*7，120/400*7,200/500*7,250/400*7] # 真实框在特征图上的位置

II. dxy

cxcy_sample = cxcy[i]          # 真实框相对原图的中心点
ij = (cxcy_sample*grid_num).ceil()-1  # 真实框在特征图上的中心点，为了防止中心点越界，因此中心点坐标向下取整并减一。
dxy = cxcy_sample*grid_num-ij  # 偏移

III. target.shape[7,7,30]，前10个数据是两个框和置信度数据，后面20个是类别的one_hot形式。

dim = 2
label = 5
0   1  2  3    4    5   6   7  8    9     10  11  12  13  14 ...29
dx  dy w  h  conf   dx  dy  w  h  conf     0   0   0   0   1 ... 0

编码流程：

I. 真实框左上角、右下角[x1,y1,x2,y2]坐标转换为中心点和宽高[cx,cy,w,h]。

II. 遍历图像中真实框，计算真实框在特征图上的中心点坐标、中心点坐标偏移。把框、置信度、类别信息填写在标签target的对应位置上。


def encoder(self,boxes,labels):
        grid_num = 7     # 特征图边长，feature.shape:[7,7,30]
        target = torch.zeros((grid_num,grid_num,30))   # 标签
        wh = boxes[:,2:] - boxes[:,:2]         # 真实框相对原图的宽高
        cxcy = (boxes[:,2:] + boxes[:,:2])/2   # 真实框相对原图的中心点
        for i in range(cxcy.size()[0]):  # 遍历每个框
            cxcy_sample = cxcy[i]         
            ij = (cxcy_sample*grid_num).ceil()-1  # 真实框在特征图上的中心点
            dxy = cxcy_sample*grid_num-ij     # 中心偏移
            target[int(ij[1]),int(ij[0]),:2] = target[int(ij[1]),int(ij[0]),5:7] = dxy   
            target[int(ij[1]),int(ij[0]),2:4] = target[int(ij[1]),int(ij[0]),7:9] = wh[i]
            target[int(ij[1]),int(ij[0]),4] = target[int(ij[1]),int(ij[0]),9] = 1  # 置信度
            target[int(ij[1]),int(ij[0]),int(labels[i])+9] = 1   # 类别
        return target

随机调节亮度

def RandomBrightness(self,bgr):
    if random.random()<0.5:        # 随机值
        hsv = self.BGR2HSV(bgr)    # bgr-->hsv
        h,s,v = cv2.split(hsv)     # 通道分割
        v = v.astype(float)        # 取出亮度通道，转换数据类型
        v *= random.choice([0.5,1.5])  # 随机变换亮度值
        v = np.clip(v,0,255).astype(hsv.dtype)  # 限制数据范围，转换数据类型
        hsv = cv2.merge((h,s,v))   # 通道合并
        bgr = self.HVS2BGR(hsv)    # hsv-->bgr
    return  bgr

随机平移

I. 找到平移距离shift_x，shift_y，平移图片。

II. 判断框平移后是否合理。

III.对框平移。

bgr.shape[100,100,3]
1: shift_x= 20，shift_y= 30,after_shift_imge[30:,20:,:]=bgr[:70,:80,:]向右下角移动
2: shift_x= 20，shift_y=-30,after_shift_imge[:70,20:,:]=bgr[30:,:80,:]向上平移
3: shift_x=-20，shift_y= 30,after_shift_imge[30:,:80,:]=bgr[:70,20:,:]向左平移
4: shift_x=-20，shift_y=-30,after_shift_imge[:70,:80,:]=bgr[30:,20:,:]向左上角移动

def randomShift(self,bgr,boxes,labels):
        center = (boxes[:,2:]+boxes[:,:2])/2
        if random.random()<0.5:
            height,width,c = bgr.shape
            after_shift_imge = np.zeros((height,width,c),dtype=bgr.dtype)
            after_shift_imge[:,:,:] = (104,117,123)
            shift_x = random.uniform(-width*0.2,width*0.2)
            shift_y = random.uniform(-height*0.2,height*0.2)
            if shift_x>=0 and shift_y>=0:
                after_shift_imge[int(shift_y):,int(shift_x):,:] = bgr[:height-int(shift_y),:width-int(shift_x),:]
            elif shift_x>=0 and shift_y<0:
                after_shift_imge[:height+int(shift_y),int(shift_x):,:] = bgr[-int(shift_y):,:width-int(shift_x),:]
            elif shift_x <0 and shift_y >=0:
                after_shift_imge[int(shift_y):,:width+int(shift_x),:] = bgr[:height-int(shift_y),-int(shift_x):,:]
            elif shift_x<0 and shift_y<0:
                after_shift_imge[:height+int(shift_y),:width+int(shift_x),:] = bgr[-int(shift_y):,-int(shift_x):,:]

            shift_xy = torch.FloatTensor([[int(shift_x),int(shift_y)]]).expand_as(center)
            center = center + shift_xy # 框中心点移动后的位置
            mask1 = (center[:,0]>0)& (center[:,0]>height)  # 
            mask2 = (center[:,1]>0)& (center[:,1]>width)
            mask = (mask1 & mask2).view(-1,1) 
            boxes_in = boxes[mask.expand_as(boxes)].view(-1,4)  # 筛选移动后仍存在的框
            if len(boxes_in) == 0:
                return bgr,boxes,labels
            box_shift = torch.FloatTensor([[int(shift_x),int(shift_y),int(shift_x),int(shift_y)]]).expand_as(boxes_in)
            boxes_in = boxes_in+box_shift
            labels_in = labels[mask.view(-1)]   # 筛选移动后仍存在的标签
            return after_shift_imge,boxes_in,labels_in
        return bgr,boxes,labels

随机缩放

def randomScale(self,bgr,boxes):
    if random.random() < 0.5:
        scale = random.uniform(0.8,1.2)  # 缩放因子
        h,w,c = bgr.shape
        bgr = cv2.resize(bgr,(int(w*scale),h)) # 缩放图片宽，高度不变
        scale_tensor = torch.FloatTensor([[scale,1,scale,1]]).expand_as(boxes)
        boxes = boxes*scale_tensor    # 缩放框的宽度
        return bgr,boxes
    return bgr,boxes

随机剪切

I.确定随机剪切的宽高hw，根据宽高确定剪切起始点xy.

II.根据框的中心点判断框是否在剪切区间。

III.对在剪切区间的框移动。返回剪切后的图片、框、标签。

def randomCrop(self,bgr,boxes,labels):
    if random.random() < 0.5:
        center = (boxes[:,:2]+boxes[:,2:])/2
        height,width,c = bgr.shape
        h = random.uniform(0.6*height,height)
        w = random.uniform(0.6*width,width)
        x = random.uniform(0,width-w)
        y = random.uniform(0,height-h)
        x,y,h,w = int(x),int(y),int(h),int(w)

        center = center - torch.FloatTensor([[x,y]]).expand_as(center)
        mask1 = (center[:,0]>0) & (center[:,0]<w)
        mask2 = (center[:,1]>0) & (center[:,0]<h)
        mask = (mask1 & mask2).view(-1,1)

        boxes_in = boxes[mask.expand_as(boxes)].view(-1,4)
        if(len(boxes_in)==0):
            return bgr,boxes,labels

        box_shift = torch.FloatTensor([[x,y,x,y]]).expand_as(boxes_in)
        boxes_in = boxes_in-box_shift
        boxes_in[:,0]=boxes_in[:,0].clamp(0,w)
        boxes_in[:,2]=boxes_in[:,2].clamp(0,w)
        boxes_in[:,1]=boxes_in[:,1].clamp(0,h)
        boxes_in[:,3]=boxes_in[:,3].clamp(0,h)

        labels_in = labels[mask.view(-1)]
        img_croped = bgr[y:y+h,x:x+w,:]
        return img_croped,boxes_in,labels_in
    return bgr,boxes,labels

随机左右翻转

图片左右翻转后，框距左右边界翻转。

def random_flip(self,im,boxes):
    if random.random() < 0.5:
        im_lr = np.fliplr(im).copy()
        h,w,_ = im.shape
        xmin = w - boxes[:,2]   # 框距右边边界的距离
        xmax = w - boxes[:,0]   # 框距左边边界的距离
        # 调换左右边界距离就是翻转后的框
        boxes[:,0] = xmin       
        boxes[:,2] = xmax
        return im_lr,boxes
    return im,boxes

2 训练

论文中的主干模型是由24层卷积和两个全联接组成。代码中
训练主要包括backbone(ResNet、VGG)、LOSS、代入数据训练模型。作者在ImageNet数据集上训练模型，达到top5上达到 88%的准确率。

2.1 Backbone

（1）ResNet
ResNet主要分为三部分。首先通过卷积和池化进行两次步长为2的下采样，然后通过残差模块layer1～layer5扩展通道数和三次步长为2的下采样，最后一次卷积、批归一化、激活得到特征图[7,7,30]。

resnet_yolo.py

import math
import torch.nn as nn
import torch.utils.model_zoo as model_zoo
import torch.nn.functional as F

__all__ = ['ResNet','resnet18','resnet34','resnet50','resnet101','resnet152']

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}

def conv3x3(in_planes,out_planes,stride=1):
    return nn.Conv2d(in_planes,out_planes,kernel_size=3,stride=stride,padding=1,bias=False)

class BasicBlock(nn.Module):
    expansion = 1
    def __init__(self,inplanes,planes,stride=1,downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self,x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.downsample is not None:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out

class Bottleneck(nn.Module):
    expansion = 4
    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes,planes,kernel_size=1,bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes,planes,kernel_size=3,stride=stride,padding=1,bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(inplanes, planes * self.expansion, kernel_size=1,bias=False)
        self.bn3 = nn.BatchNorm2d(planes* self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self,x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.downsample is not None:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out

class detnet_bottleneck(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(detnet_bottleneck, 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()
        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 = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out += self.downsample(x)
        out = F.relu(out)
        return out

class ResNet(nn.Module): # model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
    def __init__(self, block, layers, num_classes=1470):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.conv1 = nn.Conv2d(3,64,kernel_size=7,stride=2,padding=3,bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3,stride=2,padding=1)  # torch.Size([2, 64, 112, 112])
        self.layer1 = self._make_layer(block,64,layers[0])
        self.layer2 = self._make_layer(block,128,layers[1],stride=2)
        self.layer3 = self._make_layer(block,256,layers[2],stride=2)
        self.layer4 = self._make_layer(block,512,layers[3],stride=2)
        self.layer5 = self._make_detnet_layer(in_channels=512)   # (in_channels=2048)
        self.conv_end = nn.Conv2d(256,30,kernel_size=3,stride=1,padding=1,bias=False)
        self.bn_end = nn.BatchNorm2d(30)
        for m in self.modules():
            if isinstance(m,nn.Conv2d):
                n = m.kernel_size[0]*m.kernel_size[1]*m.out_channels
                m.weight.data.normal_(0,math.sqrt(2./n))
            elif isinstance(m,nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self,block, planes, blocks, stride=1): # 64,3
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes,planes*block.expansion,kernel_size=1,stride=stride,bias=False),
                nn.BatchNorm2d(planes*block.expansion),
            )
        layers = []
        layers.append(block(self.inplanes,planes,stride,downsample))
        self.inplanes = planes*block.expansion
        for i in range(1,blocks):
            layers.append(block(self.inplanes,planes))
        return nn.Sequential(*layers)

    def _make_detnet_layer(self,in_channels):
        layers = []
        layers.append(detnet_bottleneck(in_planes=in_channels, planes=256, block_type='B'))
        layers.append(detnet_bottleneck(in_planes=256, planes=256, block_type='A'))
        layers.append(detnet_bottleneck(in_planes=256, planes=256, block_type='A'))
        return nn.Sequential(*layers)
    def forward(self,x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        # print(x.shape)
        x = self.layer5(x)
        # print(x.shape)

        x = self.conv_end(x)
        x = self.bn_end(x)
        x = F.sigmoid(x)

        x = x.permute(0,2,3,1)
        return x

def resnet18(pretrained=False,**kwargs):
    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
    return model

def resnet34(pretrained=False,**kwargs):
    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
    return model

def resnet50(pretrained=False,**kwargs):
    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
    return model

def resnet101(pretrained=False,**kwargs):
    model = ResNet(BasicBlock, [3, 4, 23, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
    return model

def resnet152(pretrained=False,**kwargs):
    model = ResNet(BasicBlock, [3, 8, 36, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
    return model


def test():
    import torch
    from torch.autograd import Variable
    model = resnet18()
    x = torch.rand(2, 64, 112, 112)
    x = Variable(x)
    out = model(x)
    print(out.shape)

if __name__ == '__main__':
    test()

（2）VGG
VGG 模型分为两大部分，一部分用卷积、池化提取特征，然后通过两次全连接得到特征。
net.py

#encoding:utf-8
import math
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.model_zoo as model_zoo

__all__ = ['VGG','vgg11','vgg11_bn',
           'vgg13','vgg13_bn', 'vgg16',
           'vgg16_bn','vgg19','vgg19_bn']

model_urls = {
    'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
    'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
    'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
    'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
    'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
    'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
    'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
    'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
}

class VGG(nn.Module):
    def __init__(self,features,num_classes=1000,image_size=448):
        super(VGG,self).__init__()
        self.features = features
        self.image_size = image_size
        self.classifier = nn.Sequential(
            nn.Linear(512*7*7,4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096,1470))
        self._initialize_weights()

    def forward(self,x):
        x = self.features(x)
        x = x.view(x.size(0),-1)
        x = self.classifier(x)
        x = F.sigmoid(x)
        x = x.view(-1,7,7,30)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m,nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0,math.sqrt(2./n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

def make_layers(cfg,batch_norm=False):
    layers = []
    in_channels = 3
    first_flag = True
    for v in cfg:
        s = 1
        if (v == 64 and first_flag):
            s = 2
            first_flag = False
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2,stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels,v,kernel_size=3,stride=s,padding=1)
            if batch_norm:
                layers += [conv2d,nn.BatchNorm2d(v),nn.ReLU(inplace=True)]
            else:
                layers += [conv2d,nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

cfg = { 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
        'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
        'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
        'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']}

def vgg11(pretrained=False,**kwargs):
    model = VGG(make_layers(cfg['A']),**kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg11']))
    return model

def vgg11_bn(pretrained=False,**kwargs):
    model  = VGG(make_layers(cfg['A'],batch_norm=True),**kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url([vgg11_bn]))
    return model

def vgg13(pretrained=False,**kwargs):
    model = VGG(make_layers(cfg['B']),**kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg13']))
    return model

def vgg13_bn(pretrained=False, **kwargs):
    model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg13_bn']))
    return model

def vgg16(pretrained=False, **kwargs):
    model = VGG(make_layers(cfg['D']), **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
    return model

def vgg16_bn(pretrained=False, **kwargs):
    model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn']))
    return model


def vgg19(pretrained=False, **kwargs):
    model = VGG(make_layers(cfg['E']), **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg19']))
    return model

def vgg19_bn(pretrained=False, **kwargs):
    model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg19_bn']))
    return model

def test():
    # import torch
    # from torch.autograd import Variable
    # model = vgg16()
    # img = torch.rand(2,3,448,448)
    # img = Variable(img)
    # output = model(img)
    # print(output.size())
    import torch
    from torch.autograd import Variable
    model = vgg16()
    img = torch.rand(2,3,448,448)
    img = Variable(img)
    output = model(img)
    print(output.size())

if __name__ == '__main__':
    test()

2.2 Loss

YOLOV1计算损失的特殊性：用MSE计算损失，对框回归的宽高先开方在进行MSE，解决大小物体损失差异过大的问题。对回归和前景分类赋予不同的权重，解决正负样本不均衡问题。
yoloLoss.py
计算损失时，输入的是真实框target_tensor、和解码后的预测框pred_tensor[batch_size,7,7,30].

（1）计算损失流程：

根据真实框的置信度对target_tensor和pred_tensor取出没有真实框的样本sample_nobj[-1,30],在取出样本的第5列和第10列，用mse计算负样本的损失noobj_loss。
根据真实框的置信度对target_tensor和pred_tensor取出没有真实框的样本sample_obj[-1,30]。对取出的样本分别在提取target_tensor和pred_tensor的框和物体类别，计算类别损失。
根据sample_obj，计算预测框和真实框的IuO,根据IuO选出与真实框匹配的样本，计算框的回归损失和正样本的损失。预测正样本的真实值用IuO计算。
对各种损失加权以平衡正负样本不平衡。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

class yoloLoss(nn.Module):
    def __init__(self,S,B,l_coord,l_noobj):
        super(yoloLoss, self).__init__()
        self.S = S
        self.B = B
        self.l_coord = l_coord
        self.l_noobj = l_noobj

    def compute_iou(self,box1,box2):
        '''
        Args:
            box1[N,4],box2[M,4]
        Return:
            iou, sized [N,M].
        '''
        N = box1.size()[0]
        M = box2.size()[0]

        lt = torch.max(
            box1[:,:2].unsqueeze(1).expand(N,M,2),
            box2[:,:2].unsqueeze(0).expand(N,M,2)
                       )
        rd = torch.min(
            box1[:,2:].unsqueeze(1).expand(N,M,2),
            box2[:,2:].unsqueeze(0).expand(N,M,2)
        )

        wh = rd-lt
        wh[wh<0] = 0
        inter = wh[...,0] * wh[...,1]

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

        iou = inter/(area1+area2-inter)
        return iou

    def forward(self,pred_tensor,target_tensor):
        ''' pred_tensor[b,S,S,B*5+20] ; target_tensor[b,S,S,30]'''
        # 1 mask_obj_nobj
        N = pred_tensor.size(0)
        coo_mask = target_tensor[...,4] > 0   # 存在物体的mask [batch_size,7,7]
        noo_mask = target_tensor[...,4] == 0  # 不存在物体的mask
        coo_mask = coo_mask.unsqueeze(-1).expand_as(target_tensor) # [b,7,7,30]
        noo_mask = noo_mask.unsqueeze(-1).expand_as(target_tensor)

        # 2 nobj loss
        noo_pred = pred_tensor[noo_mask].view(-1,30) # 没有物体的预测值
        # print('noo_mask.shape:',noo_mask.shape)
        # print('pred_tensor.shape:',pred_tensor.shape)
        # print('noo_pred.shape:',noo_pred.shape)
        noo_target = target_tensor[noo_mask].view(-1,30) # 存在物体的预测值
        noo_pred_c = noo_pred[:,[4,9]].flatten() # 取出预测值中的负样本的置信度
        noo_target_c = noo_target[:,[4,9]].flatten()  # 取出标签中负样本的置信度
        noobj_loss = F.mse_loss(noo_pred_c,noo_target_c,size_average=False)  # 计算负样本损失

        # 3  obj: box , class
        coo_pred = pred_tensor[coo_mask].view(-1,30)  # 存在物体的预测值
        box_pred = coo_pred[:,:10].contiguous().view(-1,5) # 预测框
        class_pred = coo_pred[:,10:]  # 预测类别
        coo_target = target_tensor[coo_mask].view(-1,30)  # 存在物体的标签
        box_target = coo_target[:,:10].contiguous().view(-1,5)  # 真实框
        class_target = coo_target[:,10:] # 真实类别

        # 3.1  class loss
        class_loss = F.mse_loss(class_pred,class_target,size_average=False) # 类别损失
        # 4  obj_iou（每个网格上有两个预测框，根据IoU选出与真实框最匹配的预测框计算回归损失和正样本损失）
        coo_response_mask = torch.ByteTensor(box_target.size()).zero_()
        # coo_response_mask = torch.tensor(coo_response_mask,dtype=torch.bool)
        box_target_iou = torch.zeros(box_target.size())
        for i in range(0,box_target.size(0),2):  # 遍历存在物体的框
            box1 = box_pred[i:i+2]   # 存在物体的两个预测框
            box1_xy = Variable(torch.FloatTensor(box1.size()))
            box1_xy[:,:2] = box1[:,:2] / 14. - 0.5*box1[:,2:4]
            box1_xy[:,2:4] = box1[:,:2] / 14. + 0.5*box1[:,2:4]
            box2 = box_target[i].view(-1,5)  # 存在物体的一个真实框
            box2_xy = Variable(torch.FloatTensor(box2.size()))
            box2_xy[:,:2] = box2[:,:2] / 14. - 0.5*box2[:,2:4]
            box2_xy[:,2:4] = box2[:,:2] / 14. + 0.5*box2[:,2:4]
            iou = self.compute_iou(box1_xy[:,:4],box2_xy[:,:4])
            max_iou,max_index = iou.max(0)  # 计算预测框和真实框的IoU，并返回最有的IoU和预测框的下标
            coo_response_mask[i+max_index] = 1
            box_target_iou[i+max_index,4] = max_iou
        box_target_iou = Variable(box_target_iou)

        # 4.1 obj_loss
        box_pred_response = box_pred[coo_response_mask].view(-1,5) # 与真实框最匹配的预测框
        box_target_response = box_target[coo_response_mask].view(-1,5)  # 真是框，这一步多余。
        box_target_response_iou = box_target_iou[coo_response_mask].view(-1,5)  # 正样本的概率
        # 4.1.1 contain_loss
        contain_loss = F.mse_loss(box_pred_response[:,4],box_target_response_iou[:,4],size_average = False)   # 正样本损失
        # 4.1.2 loc_loss
        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:]),torch.sqrt(box_target_response[:,2:]),size_average = False)  # 框的回归损失

        return (self.l_noobj*noobj_loss + class_loss + 2*contain_loss + self.l_coord*loc_loss)/N   加权平均损失

if __name__ == '__main__':
    pred_tensor = torch.randn(2,14,14,30)
    target_tensor =  pred_tensor+0.01
    yolo_loss = yoloLoss(14,8,5,0.5)
    loss = yolo_loss(pred_tensor,target_tensor)
    print(loss)

（2）IoU
I. 计算框相交部分的左上角和右下角坐标lt,rd。
II. 计算交集面积inter和相交框的各自面积area1、area2。
III.根据以上步骤计算交并比iou。

def compute_iou(self,box1,box2):
        '''
        Args:
            box1[N,4],box2[M,4]
        Return:
            iou, sized [N,M].
        '''
        N = box1.size()[0]
        M = box2.size()[0]

        lt = torch.max(
            box1[:,:2].unsqueeze(1).expand(N,M,2),  # box1.shape[N,4]-->box1[:,:2].shape[N,2]-->box1[:,:2].unsqueeze(1).shape[N,1,2]-->lt.shape[N,M,2]
            box2[:,:2].unsqueeze(0).expand(N,M,2)
        )
        rd = torch.min(
            box1[:,2:].unsqueeze(1).expand(N,M,2),
            box2[:,2:].unsqueeze(0).expand(N,M,2)
        )

        wh = rd-lt     # wh.shape(N,M,2)
        wh[wh<0] = 0
        inter = wh[...,0] * wh[...,1]  # [N,M]

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

        iou = inter/(area1+area2-inter)   # iou.shape[N,M]
        return iou

2.3 train

train.py
训练流程:

导入库
设置超参数
模型
导入模型参数
损失函数
设置优化器
导入数据
训练

# 1 导入库
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision import models
from torch.autograd import Variable

from net import vgg16_bn
from resnet_yolo import resnet50
from yoloLoss import yoloLoss
from data_label_13 import yoloDataset
from visualize import Visualizer
import numpy as np

# 2 设置参数
# use_gpu = torch.cuda.is_available()
file_root = r'/Users/liushuang/PycharmProjects/YOLOV1_LS'
learning_rate = 0.001
num_epochs = 2
batch_size = 1
use_resnet = False

# 3 backbone
if use_resnet:
    net = resnet50()
else:
    net = vgg16_bn()
# print(net)

# 3.1 导入预训练参数
if use_resnet:
    resnet = models.resnet50(pretrained=False)  # True
    new_state_dict = resnet.state_dict()
    dd = net.state_dict()
    for k in new_state_dict.keys():
        if k in dd.keys() and not k.startswith('fc'):
            dd[k] = new_state_dict[k]
    net.load_state_dict(dd)
else:
    vgg = models.vgg16_bn(pretrained=False)
    new_state_dict = vgg.state_dict()
    dd = net.state_dict()
    for k in new_state_dict.keys():
        if k in dd.keys() and k.startswith('features'):
            dd[k] = new_state_dict[k]
    net.load_state_dict(dd)
if False:
    net.load_state_dict(torch.load('best.pth'))

# 4 Loss
criterion = yoloLoss(7,2,5,0.5)

# if use_gpu:
#     net.cuda()
# 模型训练
net.train()

# 5 参数
params = []
params_dict = dict(net.named_parameters())

for k,v in params_dict.items():
    if k.startswith('features'):
        params += [{'params':[v],'lr':learning_rate*1}]
    else:
        params += [{'params':[v],'lr':learning_rate*1}]

# 6 优化器
optimizer = torch.optim.SGD(params,lr=learning_rate,momentum=0.9,weight_decay=5e-4)

# 7 导入数据

train_dataset = yoloDataset(root=file_root,list_file='2007_train.txt',train=True,transform = [transforms.ToTensor()] )
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=0)

test_dataset = yoloDataset(root=file_root,list_file='2007_test.txt',train=False,transform = [transforms.ToTensor()] )
test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False,num_workers=4)

print('the dataset has %d images' % (len(train_dataset)))
print('the batch_size is %d' % (batch_size))
logfile = open('log.txt', 'w')

num_iter = 0
vis = Visualizer(env='LS')
best_test_loss = np.inf

# 8 训练
for epoch in range(num_epochs):
    net.train()
    if epoch == 30:
        learning_rate = 0.0001
    if epoch == 40:
        learning_rate = 0.00001
    for params_group in optimizer.param_groups:
        params_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 = Variable(images)
        target = Variable(target)
        # if use_gpu:
        #     images,target = images.cuda(),target.cuda()

        pred = net(images)
        # print('pred.shape:',pred.shape)
        # print('target.shape:',target.shape)
        loss = criterion(pred,target)
        total_loss += loss.data.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.data.item(), total_loss / (i+1)))
            num_iter += 1
            vis.plot_train_val(loss_train = total_loss/(i+1))

    validation_loss = 0.0
    net.eval()
    for i,(images,target) in enumerate(test_loader):
        images = Variable(images,volatile=True)
        target = Variable(target,volatile=True)
        if use_gpu:
            images,target = images.cuda(),target.cuda()
        pred = net(images)
        loss = criterion(pred,target)
        validation_loss += loss.data[0]
    validation_loss /= len(test_loader)
    vis.plot_train_val(loss_val=validation_loss)

    if best_test_loss > validation_loss:
        best_test_loss = validation_loss
        print('get best test loss %.5f' % best_test_loss)
        torch.save(net.state_dict(),'best.pth')
    logfile.writelines(str(epoch) + '\t' + str(validation_loss) + '\n')
    logfile.flush()
    torch.save(net.state_dict(),'yolo.pth')

3 预测

3.1 图片输入模型得到预测结果

图片输入模型得到预测结果

（1）预测
predict.py

import torch
from torch.autograd import Variable
from resnet_yolo import resnet50
import torchvision.transforms as transforms
import cv2
import numpy as np

VOC_CLASSES = (    # always index 0
    'aeroplane', 'bicycle', 'bird', 'boat',
    'bottle', 'bus', 'car', 'cat', 'chair',
    'cow', 'diningtable', 'dog', 'horse',
    'motorbike', 'person', 'pottedplant',
    'sheep', 'sofa', 'train', 'tvmonitor')

Color = [[0, 0, 0],[128, 0, 0],[0, 128, 0],[128, 128, 0],[0, 0, 128],
         [128, 0, 128],[0, 128, 128],[128, 128, 128],[64, 0, 0],[192, 0, 0],
         [64, 128, 0],[192, 128, 0],[64, 0, 128],[192, 0,128],[64, 128, 128],
         [192, 128, 128],[0, 64, 0],[128, 64, 0],[0, 192, 0],[128, 192, 0],[0, 64, 128]]

def nms(bboxes,scores,threshold=0.5):
    x1 = bboxes[:,0]
    y1 = bboxes[:,1]
    x2 = bboxes[:,2]
    y2 = bboxes[:,3]
    areas = (x2-x1)*(y2-y1)

    _,order = scores.sort(0,descending=True)
    keep = []
    while order.numel() > 0:
        if order.numel()>1:
            i = order[0]
        else:
            i = order
        keep.append(i)

        if order.numel() == 1:
            break

        xx1 = x1[order[1:]].clamp(min=x1[i])
        yy1 = y1[order[1:]].clamp(min=y1[i])
        xx2 = x2[order[1:]].clamp(max=x1[i])
        yy2 = y2[order[1:]].clamp(max=y1[i])
        w = (xx2-xx1).clamp(min=0)
        h = (yy2-yy1).clamp(min=0)
        inter = w*h

        ove = inter/(areas[i]+areas[order[1:]]-inter)
        #ids = (ove <= threshold).nonzero().squeeze()
        ids = torch.nonzero(ove <= threshold).squeeze()
        if ids.numel() == 0:
            break
        order = order[ids+1]
    return torch.LongTensor(keep)

def decoder(pred):
    grid_num = 7
    boxes = []
    cls_indexs = []
    probs = []
    cell_size = 1./grid_num
    pred = pred.data
    pred = pred.squeeze(0)  # 7x7x30
    contain1 = pred[:,:,4].unsqueeze(2)  # [7, 7, 1]
    contain2 = pred[:,:,9].unsqueeze(2)  # [7, 7, 1]
    contain = torch.cat((contain1,contain2),2)  # [7, 7, 2]
    mask1 = contain > 0.1  # [7, 7, 2]
    mask2 = (contain==contain.max())  # [7, 7, 2]
    mask = (mask1+mask2).gt(0)  # [7, 7, 2]
    for i in range(grid_num):
        for j in range(grid_num):
            for b in range(2):
                if mask[i,j,b] == 1:
                    box = pred[i,j,b*5:b*5+4]
                    contain_prob = torch.FloatTensor([pred[i,j,b*5+4]])
                    xy = torch.FloatTensor([j,i])*cell_size #cell左上角  up left of cell
                    box[:2] = box[:2]*cell_size + xy # return cxcy relative to image
                    box_xy = torch.FloatTensor(box.size())#转换成xy形式    convert[cx,cy,w,h] to [x1,y1,x2,y2]
                    box_xy[:2] = box[:2] - 0.5*box[2:]
                    box_xy[2:] = box[:2] + 0.5*box[2:]
                    max_prob,cls_index = torch.max(pred[i,j,10:],0)
                    if float((contain_prob*max_prob)[0]) > 0.1:
                        boxes.append(box_xy.view(1,4))
                        cls_indexs.append(torch.LongTensor(cls_index,0))
                        probs.append(contain_prob*max_prob)
    if len(boxes) == 0:
        boxes = torch.zeros((1,4))
        probs = torch.zeros(1)
        cls_indexs = torch.zeros(1)
    else:
        boxes = torch.cat(boxes,0) #(n,4)
        probs = torch.cat(probs,0) #(n,)
        cls_indexs = torch.cat(cls_indexs,0) #(n,)
    keep = nms(boxes,probs)
    return boxes[keep],cls_indexs[keep],probs[keep]

def predict_gpu(model,image_name,root_path='/Users/ls/PycharmProjects/YOLOV1_LS/VOCdevkit/VOC2007/JPEGImages/'):

    result = []
    image = cv2.imread(root_path+image_name)
    # 1 图片预处理
    h,w,_ = image.shape
    img = cv2.resize(image,(448,448))  # 统一输入模型的图片尺寸
    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)  # 色彩空间转换
    mean = (123,117,104)  #RGB
    img = img - np.array(mean,dtype=np.float32)  # 去均值

    transform = transforms.Compose([transforms.ToTensor(),])
    img = transform(img)  # 转置
    img = Variable(img[None,:,:,:],volatile=True)
    # img = img.cuda()

    # 2 预测
    pred = model(img) #1x7x7x30   
    pred = pred.cpu()
    # 3 解码
    boxes,cls_indexs,probs =  decoder(pred)

    for i,box in enumerate(boxes):
        x1 = int(box[0]*w)
        x2 = int(box[2]*w)
        y1 = int(box[1]*h)
        y2 = int(box[3]*h)
        cls_index = cls_indexs[i]
        if cls_index.numel()==0:return
        cls_index = int(cls_index) # convert LongTensor to int
        prob = probs[i]
        prob = float(prob)
        result.append([(x1,y1),(x2,y2),VOC_CLASSES[cls_index],image_name,prob])
    return result

if __name__ == '__main__':
    model = resnet50()
    print('load model...')
    # model.load_state_dict(torch.load('best.pth'))
    model.eval()
    #model.cuda()
    image_name = '000015.jpg'
    image = cv2.imread(image_name)
    print('predicting...')
    result = predict_gpu(model,image_name)
    #  4 画框
    for left_up,right_bottom,class_name, _ ,prob in result:
        color = Color[VOC_CLASSES.index(class_name)]
        cv2.rectangle(image,left_up,right_bottom,color,2)
        label = class_name+str(round(prob,2))
        text_size, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)
        p1 = (left_up[0], left_up[1]- text_size[1])
        cv2.rectangle(image, (p1[0] - 2//2, p1[1] - 2 - baseline), (p1[0] + text_size[0], p1[1] + text_size[1]), color, -1)
        cv2.putText(image, label, (p1[0], p1[1] + baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255,255,255), 1, 8)

    cv2.imwrite('result.jpg',image)

（2）NMS

NMS的目的是根据IoU删除重复的预测框，原理是同一物体的预测框IoU 大，不同物体的预测框IoU小。

I. 计算预测框面积。

II. 对置信度降序排列，返回排序下标。把排序好的第一个概率大的预测框取出，计算剩余框与第一个框的IoU。

III. 根据阈值筛选保留的框，并对保留的框做同样的操作，直到剩余一个框，则停止操作。


def nms(bboxes,scores,threshold=0.5):
    x1 = bboxes[:,0]
    y1 = bboxes[:,1]
    x2 = bboxes[:,2]
    y2 = bboxes[:,3]
    # 1 计算所有预测框的面积
    areas = (x2-x1)*(y2-y1)
    # 2 按照预测概率排序
    _,order = scores.sort(0,descending=True)
    keep = []
    while order.numel() > 0:
        # 3 取出第一个框
        if order.numel()>1:
            i = order[0]
        else:
            i = order
        keep.append(i)

        if order.numel() == 1:
            break
        # 4 计算剩余框与第一个框的IoU
        xx1 = x1[order[1:]].clamp(min=x1[i])
        yy1 = y1[order[1:]].clamp(min=y1[i])
        xx2 = x2[order[1:]].clamp(max=x1[i])
        yy2 = y2[order[1:]].clamp(max=y1[i])
        w = (xx2-xx1).clamp(min=0)
        h = (yy2-yy1).clamp(min=0)
        inter = w*h

        ove = inter/(areas[i]+areas[order[1:]]-inter)
        #ids = (ove <= threshold).nonzero().squeeze()
        # 5 根据IoU剔除重合的框
        ids = torch.nonzero(ove <= threshold).squeeze()
        if ids.numel() == 0:
            break
        # 6 取出与第一个框重叠小或不重叠的框作为下一轮筛选的对象
        order = order[ids+1]
    return torch.LongTensor(keep)

（3）解码

I. 取出预测值中的置信度，根据置信度阈值初筛预测框。

II.遍历输出特征图的行、列、每个网格的框，取出对应的预测框、类别概率。根据预测偏移计算预测框的中心点。预测类别概率与置信度乘积作为最终的预测概率，再根据最终的预测概率设置阈值帅选一遍框。

III.根据预测框和物体类别概率进行非极大值抑制，输出符合条件的预测值。

def decoder(pred):
    grid_num = 7
    boxes = []
    cls_indexs = []
    probs = []
    cell_size = 1./grid_num
    pred = pred.data
    pred = pred.squeeze(0)  # 7x7x30
    contain1 = pred[:,:,4].unsqueeze(2)  # [7, 7, 1]
    contain2 = pred[:,:,9].unsqueeze(2)  # [7, 7, 1]
    # 1 根据置信度筛选框
    contain = torch.cat((contain1,contain2),2)  # [7, 7, 2]
    mask1 = contain > 0.1  # [7, 7, 2]
    mask2 = (contain==contain.max())  # [7, 7, 2]
    mask = (mask1+mask2).gt(0)  # [7, 7, 2]
    for i in range(grid_num):
        for j in range(grid_num):
            for b in range(2):
                if mask[i,j,b] == 1:
                    box = pred[i,j,b*5:b*5+4]
                    contain_prob = torch.FloatTensor([pred[i,j,b*5+4]])
                    xy = torch.FloatTensor([j,i])*cell_size #cell左上角  up left of cell
                    # 2 解码
                    box[:2] = box[:2]*cell_size + xy # return cxcy relative to image
                    box_xy = torch.FloatTensor(box.size())#转换成xy形式    convert[cx,cy,w,h] to [x1,y1,x2,y2]
                    box_xy[:2] = box[:2] - 0.5*box[2:]
                    box_xy[2:] = box[:2] + 0.5*box[2:]
                    max_prob,cls_index = torch.max(pred[i,j,10:],0)
                    # 3 根据最终预测概率筛选框
                    if float((contain_prob*max_prob)[0]) > 0.1:
                        boxes.append(box_xy.view(1,4))
                        cls_indexs.append(torch.LongTensor(cls_index,0))
                        probs.append(contain_prob*max_prob)
    if len(boxes) == 0:
        boxes = torch.zeros((1,4))
        probs = torch.zeros(1)
        cls_indexs = torch.zeros(1)
    else:
        boxes = torch.cat(boxes,0) #(n,4)
        probs = torch.cat(probs,0) #(n,)
        cls_indexs = torch.cat(cls_indexs,0) #(n,)
    # 4 非极大值抑制
    keep = nms(boxes,probs)
    return boxes[keep],cls_indexs[keep],probs[keep]