YOLO v3中的矩形推理pytorch代码详解
- 矩形推理简介
- 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