yolov4 预测框解码详解【附代码】

在上一篇文章中yolov4 LOSS代码详解【附代码】 * 详细讲解了训练过程中loss的实现过程。这篇文章将关注预测过程中box的解码问题。

            # ---------------------------------------------------------#
            #   将图像输入网络当中进行预测！
            # ---------------------------------------------------------#
            outputs = self.net(images)
            outputs = self.bbox_util.decode_box(outputs)

上面这段代码中，net是我们定义的yolov4网络，images就是我们输入的图像，outputs就是网络输出。

此刻我的outputs类型是tuple形式，里面有三个output，outputs=(output0,output1,output2)。这三个output的shape分别是(1,255,19,19)、(1,255,38,38)、(1,255,76,76)。 255=3*(5+80)，80是coco类别数量，5表示box的信息和是否有类。

decode_box()函数就是本次的研究内容，它会对网络的输出进行解码，获得新的输出。

---

decode_box详解：

获得输入尺寸

首先定义一个outputs的空列表，用来存储后面我们处理后的输出。

这里的inputs就是 网络的输出，通过enumerate进行枚举，input就是yolo的三个输出特征层。可以获得各个尺寸，以输入大小608为例，第一次遍历时，batch_size=1，input_height=input_height=19。

outputs = []
for i, input in enumerate(inputs):
            #-----------------------------------------------#
            #   输入的input一共有三个，他们的shape分别是
            #   batch_size, 255, 13, 13
            #   batch_size, 255, 26, 26
            #   batch_size, 255, 52, 52
            #-----------------------------------------------#
            batch_size      = input.size(0)
            input_height    = input.size(2)
            input_width     = input.size(3)

 获得步长(缩放比)

 接下来是计算缩放比例，就是计算网络输入缩小多少倍变成特征层大小【有些地方把这个比例也叫步长】。

            #-----------------------------------------------#
            #   输入为416x416时
            #   stride_h = stride_w = 32、16、8
            #-----------------------------------------------#
            stride_h = self.input_shape[0] / input_height
            stride_w = self.input_shape[1] / input_width

由于此时的网络输入大小是608，特征层大小为19，因此部长为32，就是缩小了32倍。

stride_h=stride_w=32 

计算anchor相对特征层缩放尺寸

 接下来是计算anchors的缩放尺寸。也就是将anchors映射到特征层上。

缩放前anchors:

[[ 12.  16.],

[ 19.  36.],

[ 40.  28.],------>对应76 * 76特征层

[ 36.  75.],

[ 76.  55.],

[ 72. 146.],----->对应38 * 38 特征层

[142. 110.],

[192. 243.],

[459. 401.]]---->对应19 * 19特征层

由于我们现在先需要处理的是19 * 19 这个特征层的，通过anchors_mask进行索引获得对应的anchors。此时我们计算得到的缩放比例为32，因此这三个尺寸的anchors也需要缩小32倍。

            #-------------------------------------------------#
            #   此时获得的scaled_anchors大小是相对于特征层的
            #-------------------------------------------------#
            scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors[self.anchors_mask[i]]]

得到缩放后的19 * 19 特征层上3种anchors的尺寸。 

[(4.4375, 3.4375),

(6.0, 7.59375),

(14.34375, 12.53125)] 

获得网络预测信息(box和类) 

我们的input shape为(1,255,19,19)，这样不太直观，也不好处理，因此我们对shape进行一个改变，变为(1,3,85,19,19)，在用permute(0,1,3,4,2)对维度顺序进行改变，此刻变为(1,3,19,19,85)，也就是我们得到的prediction。

这里需要说一下，prediction[...,0] = prediction[:,:,:,:,0]。

取最后一个维度【也就是5+80这个维度上】的前4列可以获得所有cell上预测box的center_x，center_y，w，h，第五列(prediction[...,4])是是否有物体【此刻的conf还是hard形式，所以需要经过sigmoid函数给出概率值】。从prediction[...,5:]就是80个类信息【同样也是hard形式】，进行sigmoid可以获得三种anchor在19 * 19网格上所有类别的预测置信度。

            #-----------------------------------------------#
            #   输入的input一共有三个，他们的shape分别是
            #   batch_size, 3, 13, 13, 85
            #   batch_size, 3, 26, 26, 85
            #   batch_size, 3, 52, 52, 85
            #-----------------------------------------------#
            prediction = input.view(batch_size, len(self.anchors_mask[i]),
                                    self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()

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

比如我这里举个例子。比如我现在要获得所有person这个类在所有cell中内所有anchor的置信度。那么一个有3*19*19*1个值。

tensor([[[[1.5125e-03, 5.7717e-03, 1.2517e-02,  ..., 9.4604e-03,
           6.4645e-03, 3.9380e-03],
          [3.4480e-02, 3.0047e-02, 1.8495e-02,  ..., 1.7670e-02,
           2.7245e-02, 1.4117e-02],
          [1.1721e-01, 8.4902e-02, 2.8659e-02,  ..., 4.6449e-02,
           1.5122e-01, 9.6072e-02],
          ...,
          [8.8295e-01, 9.0356e-01, 3.8897e-01,  ..., 1.9456e-03,
           1.5115e-02, 1.2139e-01],
          [8.5030e-01, 9.4131e-01, 6.9385e-01,  ..., 7.4368e-04,
           4.2023e-03, 1.2149e-01],
          [8.6811e-01, 8.6302e-01, 6.4376e-01,  ..., 2.2633e-03,
           5.0898e-03, 7.8965e-03]],

         [[6.2741e-03, 1.0798e-02, 1.7037e-02,  ..., 9.2335e-03,
           8.9423e-03, 2.2053e-02],
          [3.6071e-02, 2.1278e-02, 1.4221e-02,  ..., 1.4337e-02,
           1.4463e-02, 1.8076e-02],
          [1.4068e-01, 9.8277e-02, 4.1939e-02,  ..., 5.7934e-02,
           1.5278e-01, 1.1965e-01],
          ...,
          [8.6118e-01, 8.5999e-01, 5.3942e-01,  ..., 1.2218e-03,
           1.4375e-02, 2.0362e-01],
          [9.1455e-01, 9.4715e-01, 7.2468e-01,  ..., 8.6567e-04,
           4.2050e-03, 1.5936e-01],
          [8.8682e-01, 9.0414e-01, 7.7474e-01,  ..., 1.7651e-03,
           4.7418e-03, 1.7866e-02]],

         [[1.2892e-02, 1.0655e-02, 1.8664e-02,  ..., 1.4098e-02,
           1.0832e-02, 2.6804e-02],
          [4.9359e-02, 1.9839e-02, 3.1507e-02,  ..., 1.4819e-02,
           9.5383e-03, 2.0111e-02],
          [1.6157e-01, 1.6885e-01, 1.2669e-01,  ..., 9.5604e-02,
           1.5007e-01, 1.0721e-01],
          ...,
          [5.9196e-01, 6.2065e-01, 7.3476e-01,  ..., 2.0107e-03,
           1.2156e-02, 1.9672e-01],
          [8.8735e-01, 9.2701e-01, 8.6674e-01,  ..., 2.9043e-03,
           5.8838e-03, 1.0288e-01],
          [8.3874e-01, 8.6092e-01, 7.4291e-01,  ..., 7.3370e-03,
           8.0380e-03, 1.6741e-02]]]], device='cuda:0')

生成网格

 然后我们就可以划分网格了。可以得到grid_x和grid_y。shape均为【1，3，19，19】.每种anchor都有19*19个网格。

            #----------------------------------------------------------#
            #   生成网格，先验框中心，网格左上角 
            #   batch_size,3,13,13
            #----------------------------------------------------------#
            grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_height, 1).repeat(
                batch_size * len(self.anchors_mask[i]), 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 * len(self.anchors_mask[i]), 1, 1).view(y.shape).type(FloatTensor)

按照网络格式生成先验框的宽和高

下面的代码相当于是给所有的cell都设置好anchor的尺寸。

            #----------------------------------------------------------#
            #   按照网格格式生成先验框的宽高
            #   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)

 利用预测结果对anchors进行调整

接下来的这一步比较关键，因为我们获得分类结果是很好获得的，只需要获得置信度即可，但如果去利用anchor获得目标的尺寸【或者说box呢】。

pred-boxes是用来存储调整后的box，x.data+grid_x这种操作其实就是将预测框的坐标信息映射到特征层(我们的网格)坐标。

            #----------------------------------------------------------#
            #   利用预测结果对先验框进行调整
            #   首先调整先验框的中心，从先验框中心向右下角偏移
            #   再调整先验框的宽高。
            #----------------------------------------------------------#
            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

输出结果归一化

得到的output的shape为【1，3*19*19，85】

            #----------------------------------------------------------#
            #   将输出结果归一化成小数的形式
            #----------------------------------------------------------#
            _scale = torch.Tensor([input_width, input_height, input_width, input_height]).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)
            outputs.append(output.data)

---

上面就是我们box的解码过程了。最后再将输出结果经过NMS进行box过滤即可。

完整的代码：

    def decode_box(self, inputs):
        outputs = []
        for i, input in enumerate(inputs):
            #-----------------------------------------------#
            #   输入的input一共有三个，他们的shape分别是
            #   batch_size, 255, 13, 13
            #   batch_size, 255, 26, 26
            #   batch_size, 255, 52, 52
            #-----------------------------------------------#
            batch_size      = input.size(0)
            input_height    = input.size(2)
            input_width     = input.size(3)

            #-----------------------------------------------#
            #   输入为416x416时
            #   stride_h = stride_w = 32、16、8
            #-----------------------------------------------#
            stride_h = self.input_shape[0] / input_height
            stride_w = self.input_shape[1] / input_width
            #-------------------------------------------------#
            #   此时获得的scaled_anchors大小是相对于特征层的
            #-------------------------------------------------#
            scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors[self.anchors_mask[i]]]

            #-----------------------------------------------#
            #   输入的input一共有三个，他们的shape分别是
            #   batch_size, 3, 13, 13, 85
            #   batch_size, 3, 26, 26, 85
            #   batch_size, 3, 52, 52, 85
            #-----------------------------------------------#
            prediction = input.view(batch_size, len(self.anchors_mask[i]),
                                    self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()

            #-----------------------------------------------#
            #   先验框的中心位置的调整参数
            #-----------------------------------------------#
            x = torch.sigmoid(prediction[..., 0])  
            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
            #----------------------------------------------------------#
            grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_height, 1).repeat(
                batch_size * len(self.anchors_mask[i]), 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 * len(self.anchors_mask[i]), 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

            #----------------------------------------------------------#
            #   将输出结果归一化成小数的形式
            #----------------------------------------------------------#
            _scale = torch.Tensor([input_width, input_height, input_width, input_height]).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)
            outputs.append(output.data)
        return outputs