2021-11-11yolov5-5.0版本代码详解----augmentations.py的letterbox函数

yolov5-5.0版本代码详解----augmentations.py的letterbox函数

1、作用

将图片缩放调整到指定大小

2、引用位置

在datasets.py多次调用，分别为
1、LoadImages类的__next__函数(242行)；
2、LoadWebcam类的__next__函数(288行)；
3、LoadStreams类的__init__函数(337行)；
4、LoadStreams类的__next__函数(366行)；
5、LoadImagesAndLabels类的__getitem__函数(566行)

3、原理

3.1 传统图片缩放

在常用的目标检测算法中，不同的图片长宽都不相同，因此常用的方式是将原始图片统一缩放到一个标准尺寸，再送入检测网络中。

缺点：缩放填充后，两端的黑边大小都不同，而如果填充的比较多，则存在信息冗余，影响推理速度。

3.2 Yolov5的自适应缩放

图像高度上两端的黑边变少了，在推理时，计算量也会减少，即目标检测速度会得到提升，通过这种简单的改进，推理速度得到了37%的提升，可以说效果很明显。

4、代码

4.1 参数

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114),
              auto=True, scaleFill=False, scaleup=True, stride=32):
    """参数：
            im: 原图 hwc
            new_shape: 缩放后的尺寸
            color: pad的颜色（灰色边框，补齐调整后的区域）
            auto: True 保证缩放后的图片保持原图的比例 即 将原图最长边缩放到指定大小，再将原图较短边按原图比例缩放（不会失真）
                  False 将原图最长边缩放到指定大小，再将原图较短边按原图比例缩放,最后将较短边两边pad操作缩放到最长边大小（不会失真）
            scaleFill: True 简单粗暴的将原图resize到指定的大小，没有pad操作（失真）
            scaleup: True  
                     False 对于大于new_shape的原图进行缩放,小于的不变
            stride: 步长

4.2 计算缩放比例

    # 在满足多个步幅约束的情况下调整图像大小并填充图像

    # 第一层resize后图片大小[h, w]
    shape = im.shape[:2]

    # 如果new_shape参数是整数（这里不是），防止将new_shape = 640
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # 计算比例 (new / old) 选择小的那个
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])

原始缩放尺寸是416×416，都除以原始图像的尺寸后，可以得到0.52，和0.69两个缩放系数，选择小的缩放系数0.52。

4.3 计算缩放后的尺寸

    # 只进行下采样 因为上采样会让图片模糊
    if not scaleup:
        r = min(r, 1.0)

    # 计算pad长宽
    ratio = r, r  # 长宽比例
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))  # wh(512, 343) 保证缩放后图像比例不变

原始图片的长宽都乘以最小的缩放系数0.52，宽变成了416，而高变成了312。

4.4 计算黑边填充数值

    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # 保证原图比例不变，将图像最大边缩放到指定大小
        # 这里的取余操作可以保证padding后的图片是32的整数倍(416x416)，如果是(512x512)可以保证是64的整数倍
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # 简单粗暴的将图片缩放到指定尺寸
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    # 在较小边的两侧进行pad, 而不是在一侧pad
    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # 将原图resize到new_unpad（长边相同，比例相同的新图）
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))   # 计算上下两侧的padding
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))  # 计算左右两侧的padding
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border

将416-312=104，得到原本需要填充的高度。再采用numpy中np.mod取余数的方式，得到8个像素，再除以2，即得到图片高度两端需要填充的数值。

4.5 返回

    return im, ratio, (dw, dh)
    """:img: letterbox后的图片 HWC
        ratio: wh比例
        (dw, dh): w和h的pad
    """

5、例子

import numpy as np
import cv2


#  将图片缩放调整到指定大小
#  在datasets.py多次调用，分别为1、LoadImages类的__next__函数(242行)；2、LoadWebcam类的__next__函数(288行)；
#  3、LoadStreams类的__init__函数(337行)；4、LoadStreams类的__next__函数(366行)；
#  5、LoadImagesAndLabels类的__getitem__函数(566行)
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114),
              auto=True, scaleFill=False, scaleup=True, stride=32):
    """参数：
            im: 原图 hwc
            new_shape: 缩放后的尺寸
            color: pad的颜色（灰色边框，补齐调整后的区域）
            auto: True 保证缩放后的图片保持原图的比例 即 将原图最长边缩放到指定大小，再将原图较短边按原图比例缩放（不会失真）
                  False 将原图最长边缩放到指定大小，再将原图较短边按原图比例缩放,最后将较短边两边pad操作缩放到最长边大小（不会失真）
            scaleFill: True 简单粗暴的将原图resize到指定的大小，没有pad操作（失真）
            scaleup: True  对于小于new_shape的原图进行缩放,大于的不变
                     False 对于大于new_shape的原图进行缩放,小于的不变
            stride: 步长

    :return: img: letterbox后的图片 HWC
             ratio: wh比例
             (dw, dh): w和h的pad
    """

    # 在满足多个步幅约束的情况下调整图像大小并填充图像

    # 第一层resize后图片大小[h, w]
    shape = im.shape[:2]

    # 如果new_shape参数是整数（这里不是），防止将new_shape = 640
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # 计算比例 (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])

    # 只进行下采样 因为上采样会让图片模糊
    if not scaleup:
        r = min(r, 1.0)

    # 计算灰边长宽
    ratio = r, r  # 长宽比例
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))  # wh(512, 343) 保证缩放后图像比例不变
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # 保证原图比例不变，将图像最大边缩放到指定大小
        # 这里的取余操作可以保证padding后的图片是32的整数倍(416x416)，如果是(512x512)可以保证是64的整数倍
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # 简单粗暴的将图片缩放到指定尺寸
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    # 在较小边的两侧进行pad, 而不是在一侧pad
    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # 将原图resize到new_unpad（长边相同，比例相同的新图）
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))   # 计算上下两侧的padding
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))  # 计算左右两侧的padding
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, ratio, (dw, dh)


if __name__ == '__main__':
    img = r'2.jpg'
    im = cv2.imread(img)

    ni, _, _ = letterbox(im, new_shape=320)
    cv2.imwrite('res3.jpg', ni)

注意：
a.训练时没有采用缩减黑边的方式，还是采用传统填充的方式，即缩放到416×416大小。
只是在测试，使用模型推理时，才采用缩减黑边的方式，提高目标检测，推理的速度。
b.为什么np.mod函数的后面用32？因为Yolov5的网络经过5次下采样，而2的5次方，等于32。所以至少要去掉32的倍数，再进行取余。