YOLO v3中的矩形推理pytorch代码详解

矩形推理简介

矩形推理的作用是减少推理时间，在YOLOv3中为32倍降采样，因此输入网络的长宽需要是32的倍数，最常用的分辨率就是416了。可是如果直接将图片填充为416*416的画，可能会产生大量的冗余，如下图1。而矩形推理则是先将较长的边先调整成416大小，然后较短的边按照原来图像的宽高比缩放，然后再填充为416大小，从而达到32的倍数，如图2

图1

图2

pytorch代码实现

• 首先计算每个图片的高宽比，再将其升序排序得到排序后的索引值，再根据索引值将图像和标签重新排序
• 计算每个batch采用的统一尺度，如果高/宽小于1(w > h)，将w设为img_size， 如果高/宽大于1(w < h)，将h设置为img_size
• 最后计算每个batch输入网络的shape值(向上设置为32的整数倍)
• 例如下图，三个图像的高宽比都小于1，最大的为0.75，所以这个批次的图像的shape就设置为[0.75,1],然后根据这个比例，向上取32的整数倍计算出shape值，得到图像最终的宽高为[312,416]
  

        if self.rect:
            s = self.shapes  # wh
            # 计算每个图片的高/宽比
            ar = s[:, 1] / s[:, 0]  
            # argsort函数返回的是数组值从小到大的索引值
            # 按照高宽比例进行排序，这样后面划分的每个batch中的图像就拥有类似的高宽比
            irect = ar.argsort()
            # 根据排序后的顺序重新设置图像顺序、标签顺序以及shape顺序
            self.img_files = [self.img_files[i] for i in irect]
            self.label_files = [self.label_files[i] for i in irect]
            self.shapes = s[irect]  # wh
            ar = ar[irect]

            # 计算每个batch采用的统一尺度
            shapes = [[1, 1]] * nb  # 创建nb个shape用来存储每个batch的尺度
            for i in range(nb):
                ari = ar[bi == i]  # bi是batch的蒙版，bi = np.floor(np.arange(n) / batch_size).astype(np.int)
                # 获取第i个batch中，最小和最大高宽比
                mini, maxi = ari.min(), ari.max()

                # 如果高/宽小于1(w > h)，将w设为img_size
                if maxi < 1:
                    shapes[i] = [maxi, 1]
                # 如果高/宽大于1(w < h)，将h设置为img_size
                elif mini > 1:
                    shapes[i] = [1, 1 / mini]
            # 计算每个batch输入网络的shape值(向上设置为32的整数倍)
            self.batch_shapes = np.ceil(np.array(shapes) * img_size / 32. + pad).astype(np.int) * 32