万字长文解析 YOLOv1-v5 模型（上）

def encoder(self, boxes, labels):
    '''
    boxes (tensor) [[x1,y1,x2,y2],[]] 目标的边界框坐标信息
    labels (tensor) [...] 目标的类别信息
    return 7x7x30
    '''
    grid_num = 7 # 论文中设为7
    target = torch.zeros((grid_num, grid_num, 30))  # 和模型输出张量维尺寸一样都是 14*14*30
    cell_size = 1./grid_num  # 之前已经把目标框的坐标进行了归一化（这里与原论文有区别），故这里用1.作为除数
    # 计算目标框中心点坐标和宽高
    wh = boxes[:, 2:]-boxes[:, :2]  
    cxcy = (boxes[:, 2:]+boxes[:, :2])/2  
    # 1，遍历各个目标框；
    for i in range(cxcy.size()[0]):    # 对应于数据集中的每个框 这里cxcy.size()[0] == num_samples
        # 2，计算第 i 个目标中心点落在哪个 `grid` 上，`target` 相应位置的两个框的置信度值设为 `1`，同时对应类别值也置为 `1`；
        cxcy_sample = cxcy[i]  
        ij = (cxcy_sample/cell_size).ceil()-1 # ij 是一个list, 表示目标中心点cxcy在归一化后的图片中所处的x y 方向的第几个网格
        # [0,1,2,3,4,5,6,7,8,9, 10-19] 对应索引 
        # [x,y,w,h,c,x,y,w,h,c, 20 个类别的 one-hot编码] 与原论文输出张量维度各个索引对应目标有所区别
        target[int(ij[1]), int(ij[0]), 4] = 1  # 第一个框的置信度
        target[int(ij[1]), int(ij[0]), 9] = 1  # 第二个框的置信度
        target[int(ij[1]), int(ij[0]), int(labels[i])+9] = 1 # 第 int(labels[i])+9 个类别为 1
        # 3，计算目标中心所在 `grid`（网格）的左上角相对坐标：`ij*cell_size`，然后目标中心坐标相对于子网格左上角的偏移比例 `delta_xy`；
        xy = ij*cell_size  
        delta_xy = (cxcy_sample -xy)/cell_size  
        # 4，最后将 `target` 对应网格位置的 (x, y, w, h) 分别赋相应 `wh`、`delta_xy` 值。
        target[int(ij[1]), int(ij[0]), 2:4] = wh[i]  # 范围为(0,1)
        target[int(ij[1]), int(ij[0]), :2] = delta_xy
        target[int(ij[1]), int(ij[0]), 7:9] = wh[i]
        target[int(ij[1]), int(ij[0]), 5:7] = delta_xy
    return target
复制代码

# 对于网络输出预测 改为再图片上画出框及score
def decoder(pred):
    """
    pred (tensor)  torch.Size([1, 7, 7, 30])
    return (tensor) box[[x1,y1,x2,y2]] label[...]
    """
    grid_num = 7
    boxes = []
    cls_indexs = []
    probs = []
    cell_size = 1./grid_num
    pred = pred.data  # torch.Size([1, 14, 14, 30])
    pred = pred.squeeze(0)  # torch.Size([14, 14, 30])
    # 0 1      2 3   4    5 6   7 8   9
    # [中心坐标,长宽,置信度,中心坐标,长宽,置信度, 20个类别] x 7x7
    contain1 = pred[:, :, 4].unsqueeze(2)  # torch.Size([14, 14, 1])
    contain2 = pred[:, :, 9].unsqueeze(2)  # torch.Size([14, 14, 1])
    contain = torch.cat((contain1, contain2), 2)    # torch.Size([14, 14, 2])
    mask1 = contain > 0.1  # 大于阈值, torch.Size([14, 14, 2])  content: tensor([False, False])
    mask2 = (contain == contain.max())  # we always select the best contain_prob what ever it>0.9
    mask = (mask1+mask2).gt(0)
    # min_score,min_index = torch.min(contain, 2) # 每个 cell 只选最大概率的那个预测框
    for i in range(grid_num):
        for j in range(grid_num):
            for b in range(2):
                # index = min_index[i,j]
                # mask[i,j,index] = 0
                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(cls_index.item())
                        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)
        # print(type(probs))
        # print(len(probs))
        # print(probs)
        probs = torch.cat(probs, 0)  # (n,)
        # print(probs)
        # print(type(cls_indexs))
        # print(len(cls_indexs))
        # print(cls_indexs)
        cls_indexs = torch.IntTensor(cls_indexs)  # (n,)
    # 去除冗余的候选框，得到最佳检测框（bbox）
    keep = nms(boxes, probs)
    # print("keep:", keep)
    a = boxes[keep]
    b = cls_indexs[keep]
    c = probs[keep]
    return a, b, c
复制代码

def encode(self, boxes, labels, input_size):
    '''Encode target bounding boxes and class labels into YOLOv2 format.
    Args:
        boxes: (tensor) bounding boxes of (xmin,ymin,xmax,ymax) in range [0,1], sized [#obj, 4].
        labels: (tensor) object class labels, sized [#obj,].
        input_size: (int) model input size.
    Returns:
        loc_targets: (tensor) encoded bounding boxes, sized [5,4,fmsize,fmsize].
        cls_targets: (tensor) encoded class labels, sized [5,20,fmsize,fmsize].
        box_targets: (tensor) truth boxes, sized [#obj,4].
    '''
    num_boxes = len(boxes)
    # input_size -> fmsize
    # 320->10, 352->11, 384->12, 416->13, ..., 608->19
    fmsize = (input_size - 320) / 32 + 10
    grid_size = input_size / fmsize
    boxes *= input_size  # scale [0,1] -> [0,input_size]
    bx = (boxes[:,0] + boxes[:,2]) * 0.5 / grid_size  # in [0,fmsize]
    by = (boxes[:,1] + boxes[:,3]) * 0.5 / grid_size  # in [0,fmsize]
    bw = (boxes[:,2] - boxes[:,0]) / grid_size        # in [0,fmsize]
    bh = (boxes[:,3] - boxes[:,1]) / grid_size        # in [0,fmsize]
    tx = bx - bx.floor()
    ty = by - by.floor()
    xy = meshgrid(fmsize, swap_dims=True) + 0.5  # grid center, [fmsize*fmsize,2]
    wh = torch.Tensor(self.anchors)              # [5,2]
    xy = xy.view(fmsize,fmsize,1,2).expand(fmsize,fmsize,5,2)
    wh = wh.view(1,1,5,2).expand(fmsize,fmsize,5,2)
    anchor_boxes = torch.cat([xy-wh/2, xy+wh/2], 3)  # [fmsize,fmsize,5,4]
    ious = box_iou(anchor_boxes.view(-1,4), boxes/grid_size)  # [fmsize*fmsize*5,N]
    ious = ious.view(fmsize,fmsize,5,num_boxes)               # [fmsize,fmsize,5,N]
    loc_targets = torch.zeros(5,4,fmsize,fmsize)  # 5boxes * 4coords
    cls_targets = torch.zeros(5,20,fmsize,fmsize)
    for i in range(num_boxes):
        cx = int(bx[i])
        cy = int(by[i])
        _, max_idx = ious[cy,cx,:,i].max(0)
        j = max_idx[0]
        cls_targets[j,labels[i],cy,cx] = 1
        tw = bw[i] / self.anchors[j][0]
        th = bh[i] / self.anchors[j][1]
        loc_targets[j,:,cy,cx] = torch.Tensor([tx[i], ty[i], tw, th])
    return loc_targets, cls_targets, boxes/grid_size
复制代码

def decode(self, outputs, input_size):
    '''Transform predicted loc/conf back to real bbox locations and class labels.
    Args:
        outputs: (tensor) model outputs, sized [1,125,13,13].
        input_size: (int) model input size.
    Returns:
        boxes: (tensor) bbox locations, sized [#obj, 4].
        labels: (tensor) class labels, sized [#obj,1].
    '''
    fmsize = outputs.size(2)
    outputs = outputs.view(5,25,13,13)
    loc_xy = outputs[:,:2,:,:]   # [5,2,13,13]
    grid_xy = meshgrid(fmsize, swap_dims=True).view(fmsize,fmsize,2).permute(2,0,1)  # [2,13,13]
    box_xy = loc_xy.sigmoid() + grid_xy.expand_as(loc_xy)  # [5,2,13,13]
    loc_wh = outputs[:,2:4,:,:]  # [5,2,13,13]
    anchor_wh = torch.Tensor(self.anchors).view(5,2,1,1).expand_as(loc_wh)  # [5,2,13,13]
    box_wh = anchor_wh * loc_wh.exp()  # [5,2,13,13]
    boxes = torch.cat([box_xy-box_wh/2, box_xy+box_wh/2], 1)  # [5,4,13,13]
    boxes = boxes.permute(0,2,3,1).contiguous().view(-1,4)    # [845,4]
    iou_preds = outputs[:,4,:,:].sigmoid()  # [5,13,13]
    cls_preds = outputs[:,5:,:,:]  # [5,20,13,13]
    cls_preds = cls_preds.permute(0,2,3,1).contiguous().view(-1,20)
    cls_preds = softmax(cls_preds)  # [5*13*13,20]
    score = cls_preds * iou_preds.view(-1).unsqueeze(1).expand_as(cls_preds)  # [5*13*13,20]
    score = score.max(1)[0].view(-1)  # [5*13*13,]
    print(iou_preds.max())
    print(cls_preds.max())
    print(score.max())
    ids = (score>0.5).nonzero().squeeze()
    keep = box_nms(boxes[ids], score[ids])  # NMS 算法去除重复框
    return boxes[ids][keep] / fmsize
复制代码

'''Darknet in PyTorch.'''
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from torch.autograd import Variable
class Darknet(nn.Module):
    # (64,1) means conv kernel size is 1, by default is 3.
    cfg1 = [32, 'M', 64, 'M', 128, (64,1), 128, 'M', 256, (128,1), 256, 'M', 512, (256,1), 512, (256,1), 512]  # conv1 - conv13
    cfg2 = ['M', 1024, (512,1), 1024, (512,1), 1024]  # conv14 - conv18
    def __init__(self):
        super(Darknet, self).__init__()
        self.layer1 = self._make_layers(self.cfg1, in_planes=3)
        self.layer2 = self._make_layers(self.cfg2, in_planes=512)
        #### Add new layers
        self.conv19 = nn.Conv2d(1024, 1024, kernel_size=3, stride=1, padding=1)
        self.bn19 = nn.BatchNorm2d(1024)
        self.conv20 = nn.Conv2d(1024, 1024, kernel_size=3, stride=1, padding=1)
        self.bn20 = nn.BatchNorm2d(1024)
        # Currently I removed the passthrough layer for simplicity
        self.conv21 = nn.Conv2d(1024, 1024, kernel_size=3, stride=1, padding=1)
        self.bn21 = nn.BatchNorm2d(1024)
        # Outputs: 5boxes * (4coordinates + 1confidence + 20classes)
        self.conv22 = nn.Conv2d(1024, 5*(5+20), kernel_size=1, stride=1, padding=0)
    def _make_layers(self, cfg, in_planes):
        layers = []
        for x in cfg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
            else:
                out_planes = x[0] if isinstance(x, tuple) else x
                ksize = x[1] if isinstance(x, tuple) else 3
                layers += [nn.Conv2d(in_planes, out_planes, kernel_size=ksize, padding=(ksize-1)//2),
                           nn.BatchNorm2d(out_planes),
                           nn.LeakyReLU(0.1, True)]
                in_planes = out_planes
        return nn.Sequential(*layers)
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = F.leaky_relu(self.bn19(self.conv19(out)), 0.1)
        out = F.leaky_relu(self.bn20(self.conv20(out)), 0.1)
        out = F.leaky_relu(self.bn21(self.conv21(out)), 0.1)
        out = self.conv22(out)
        return out
def test():
    net = Darknet()
    y = net(Variable(torch.randn(1,3,416,416)))
    print(y.size())  # 模型最后输出张量大小 [1,125,13,13]
if __name__ == "__main__":
    test()
复制代码

from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from utils import box_iou, meshgrid
class YOLOLoss(nn.Module):
    def __init__(self):
        super(YOLOLoss, self).__init__()
    def decode_loc(self, loc_preds):
        '''Recover predicted locations back to box coordinates.
        Args:
          loc_preds: (tensor) predicted locations, sized [N,5,4,fmsize,fmsize].
        Returns:
          box_preds: (tensor) recovered boxes, sized [N,5,4,fmsize,fmsize].
        '''
        anchors = [(1.3221,1.73145),(3.19275,4.00944),(5.05587,8.09892),(9.47112,4.84053),(11.2364,10.0071)]
        N, _, _, fmsize, _ = loc_preds.size()
        loc_xy = loc_preds[:,:,:2,:,:]   # [N,5,2,13,13]
        grid_xy = meshgrid(fmsize, swap_dims=True).view(fmsize,fmsize,2).permute(2,0,1)  # [2,13,13]
        grid_xy = Variable(grid_xy.cuda())
        box_xy = loc_xy.sigmoid() + grid_xy.expand_as(loc_xy)  # [N,5,2,13,13]
        loc_wh = loc_preds[:,:,2:4,:,:]  # [N,5,2,13,13]
        anchor_wh = torch.Tensor(anchors).view(1,5,2,1,1).expand_as(loc_wh)  # [N,5,2,13,13]
        anchor_wh = Variable(anchor_wh.cuda())
        box_wh = anchor_wh * loc_wh.exp()  # [N,5,2,13,13]
        box_preds = torch.cat([box_xy-box_wh/2, box_xy+box_wh/2], 2)  # [N,5,4,13,13]
        return box_preds
    def forward(self, preds, loc_targets, cls_targets, box_targets):
        '''
        Args:
          preds: (tensor) model outputs, sized [batch_size,150,fmsize,fmsize].
          loc_targets: (tensor) loc targets, sized [batch_size,5,4,fmsize,fmsize].
          cls_targets: (tensor) conf targets, sized [batch_size,5,20,fmsize,fmsize].
          box_targets: (list) box targets, each sized [#obj,4].
        Returns:
          (tensor) loss = SmoothL1Loss(loc) + SmoothL1Loss(iou) + SmoothL1Loss(cls)
        '''
        batch_size, _, fmsize, _ = preds.size()
        preds = preds.view(batch_size, 5, 4+1+20, fmsize, fmsize)
        ### loc_loss
        xy = preds[:,:,:2,:,:].sigmoid()   # x->sigmoid(x), y->sigmoid(y)
        wh = preds[:,:,2:4,:,:].exp()
        loc_preds = torch.cat([xy,wh], 2)  # [N,5,4,13,13]
        pos = cls_targets.max(2)[0].squeeze() > 0  # [N,5,13,13]
        num_pos = pos.data.long().sum()
        mask = pos.unsqueeze(2).expand_as(loc_preds)  # [N,5,13,13] -> [N,5,1,13,13] -> [N,5,4,13,13]
        loc_loss = F.smooth_l1_loss(loc_preds[mask], loc_targets[mask], size_average=False)
        ### iou_loss
        iou_preds = preds[:,:,4,:,:].sigmoid()  # [N,5,13,13]
        iou_targets = Variable(torch.zeros(iou_preds.size()).cuda()) # [N,5,13,13]
        box_preds = self.decode_loc(preds[:,:,:4,:,:])  # [N,5,4,13,13]
        box_preds = box_preds.permute(0,1,3,4,2).contiguous().view(batch_size,-1,4)  # [N,5*13*13,4]
        for i in range(batch_size):
            box_pred = box_preds[i]  # [5*13*13,4]
            box_target = box_targets[i]  # [#obj, 4]
            iou_target = box_iou(box_pred, box_target)  # [5*13*13, #obj]
            iou_targets[i] = iou_target.max(1)[0].view(5,fmsize,fmsize)  # [5,13,13]
        mask = Variable(torch.ones(iou_preds.size()).cuda()) * 0.1  # [N,5,13,13]
        mask[pos] = 1
        iou_loss = F.smooth_l1_loss(iou_preds*mask, iou_targets*mask, size_average=False)
        ### cls_loss
        cls_preds = preds[:,:,5:,:,:]  # [N,5,20,13,13]
        cls_preds = cls_preds.permute(0,1,3,4,2).contiguous().view(-1,20)  # [N,5,20,13,13] -> [N,5,13,13,20] -> [N*5*13*13,20]
        cls_preds = F.softmax(cls_preds)  # [N*5*13*13,20]
        cls_preds = cls_preds.view(batch_size,5,fmsize,fmsize,20).permute(0,1,4,2,3)  # [N*5*13*13,20] -> [N,5,20,13,13]
        pos = cls_targets > 0
        cls_loss = F.smooth_l1_loss(cls_preds[pos], cls_targets[pos], size_average=False)
        print('%f %f %f' % (loc_loss.data[0]/num_pos, iou_loss.data[0]/num_pos, cls_loss.data[0]/num_pos), end=' ')
        return (loc_loss + iou_loss + cls_loss) / num_pos
复制代码