yolov3代码详解3-预测结果的解码

由第二步可以得到三个特征层，分别为(N,13,13,255)、(N,26,26,255)、(N,52,52,255)，每一个有效特征层将整个图片分成与其长宽对应的网格，如(N,13,13,255)的特征层就是将整个图像分成13x13个网格；

然后从每个网格中心建立多个先验框，这些框是网络预先设定好的框，网络的预测结果会判断这些框内是否包含物体，以及这个物体的种类。由于每一个网格点都具有三个先验框，所以上述的预测结果可以

reshape为：

• (N,13,13,3,85)
• (N,26,26,3,85)
• (N,52,52,3,85)

85可以拆分为3 * （20 + 1 + 4），分别代表二十个类，一个置信度和先验框的调整参数(x,y,w,h)，每一个节点有三个先验框。

yolov3的解码过程可以分为两个步骤：

　　1. 给每个网格点加上它对应的x和y，加完后的结果就是预测框的中心

　　2. 再利用先验框和h、w结合计算出预测框的宽高，就可以得到整个预测框的位置了 

得到最终的预测结果后还要进行得分排序与非极大抑制筛选。

这一部分基本上是所有目标检测通用的部分。其对于每一个类进行判别：
1、取出每一类得分大于self.obj_threshold的框和得分。
2、利用框的位置和得分进行非极大抑制。

from __future__ import division

import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from torchvision.ops import nms

每次只能对一个特征层进行解码，故在调用的时候循环三次，对三个特征层的框进行解码

# 每次只能对一个特征层进行解码
class DecodeBox(nn.Module):
    def __init__(self, anchors, num_classes, img_size):   # 形参分别为先验框的大小、预测的种类、图像的大小
        super(DecodeBox, self).__init__()
        #-----------------------------------------------------------#
        #   13x13的特征层对应的anchor是[116,90],[156,198],[373,326]
        #   26x26的特征层对应的anchor是[30,61],[62,45],[59,119]
        #   52x52的特征层对应的anchor是[10,13],[16,30],[33,23]
        #-----------------------------------------------------------#
        self.anchors = anchors
        self.num_anchors = len(anchors)
        self.num_classes = num_classes
        self.bbox_attrs = 5 + num_classes
        self.img_size = img_size

    def forward(self, input):
        #-----------------------------------------------#
        #   输入的input一共有三个，他们的shape分别是
        #   batch_size, 255, 13, 13
        #   batch_size, 255, 26, 26
        #   batch_size, 255, 52, 52
        # 一共有多少张图片
        batch_size = input.size(0)
        # 以 13 * 13 为例
        input_height = input.size(2)    # 13
        input_width = input.size(3)     # 13

        # 计算步长，每一个特征点对应原来的图片上相应数量的像素点，若特征层为13*13，则每一个特征点对应416/13 = 32个像素点
        stride_h = self.img_size[1] / input_height      # 416/13
        stride_w = self.img_size[0] / input_width       # 416/13

        # 此时获得的scaled_anchors大小是相对于特征层的
        # 把先验框的尺寸调整为特征层的大小，计算出特征层再特征层上对应的宽高
        scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors]

        #   输入的input一共有三个，他们的shape分别是
        #   batch_size, 3, 13, 13, 85
        #   batch_size, 3, 26, 26, 85
        #   batch_size, 3, 52, 52, 85

        # 对预测结果进行reshape通道变换
        # bs,3*(5 + num_classes),13,13 --> bs,3,13,13,(5 + num_classes)
        prediction = input.view(batch_size, self.num_anchors,
                                self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()

        # 先验框的中心位置的调整参数
        '实际预测框的中心和先验框中心的距离'
        x = torch.sigmoid(prediction[..., 0])   # 0 ~ 1之间 使得节点对应右下角的网格，换句话说框由其中心点左上角的网格点负责预测
        y = torch.sigmoid(prediction[..., 1])

        # 先验框的宽高调整参数
        w = prediction[..., 2]
        h = prediction[..., 3]
        # 获得置信度，是否有物体
        conf = torch.sigmoid(prediction[..., 4])
        # 种类置信度
        pred_cls = torch.sigmoid(prediction[..., 5:])

        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
        LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor

        # ======先验框的生成，包含一下两步======

        #  生成网格，先验框中心，网格左上角 ，先验框的中心就是网格的左上角  batch_size，3，13，13
        #  batch_size,3,13,13
        grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_height, 1).repeat(
            batch_size * self.num_anchors, 1, 1).view(x.shape).type(FloatTensor)
        grid_y = torch.linspace(0, input_height - 1, input_height).repeat(input_width, 1).t().repeat(
            batch_size * self.num_anchors, 1, 1).view(y.shape).type(FloatTensor)

        #  按照网格格式生成先验框的宽高
        #   batch_size,3,13,13
        anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
        anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
        anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(w.shape)
        anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(h.shape)

        #   ======利用预测结果对先验框进行调整======
        #   首先调整先验框的中心，从先验框中心向右下角偏移
        #   再调整先验框的宽高。
        #----------------------------------------------------------#
        pred_boxes = FloatTensor(prediction[..., :4].shape)
        pred_boxes[..., 0] = x.data + grid_x
        pred_boxes[..., 1] = y.data + grid_y
        pred_boxes[..., 2] = torch.exp(w.data) * anchor_w   # 对先验框的宽高进行调整   实际坐标系原点再左上角，所以向右下偏移
        pred_boxes[..., 3] = torch.exp(h.data) * anchor_h

        #----------------------------------------------------------#
        #   将输出结果调整成相对于输入图像大小 416 * 416
        #----------------------------------------------------------#
        _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)
        output = torch.cat((pred_boxes.view(batch_size, -1, 4) * _scale,
                            conf.view(batch_size, -1, 1), pred_cls.view(batch_size, -1, self.num_classes)), -1)
        return output.data