目标检测 YOLOv5 - 推理时的数据增强

目标检测 YOLOv5 - 推理时的数据增强

flyfish

版本 YOLOv5 6.2

参考地址

https://github.com/ultralytics/yolov5/issues/303

在训练时可以使用数据增强，在推理阶段也可以使用数据增强
在测试使用数据增强有个名字叫做Test-Time Augmentation (TTA)

实际使用中使用了大中小三个不同分辨率，中间大小分辨率的图像进行了左右反转
大分辨率
480 * 640 宽度W 高度H 比例为1

中分辨率
416 * 544 宽度W 高度H 比例为0.83

小分辨率
352 * 448 宽度W 高度H 比例为0.67

命令

python detect.py --weights ./yolov5s.pt --source ./data/images/bus.jpg  --imgsz 640 --augment

--augment语法
推理时默认不使用增强

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("-v", "--verbose", help="increase output verbosity",
                    action="store_true")
args = parser.parse_args()
if args.verbose:
    print("verbosity turned on")
else:
    print("verbosity turned off")

假如上段代码是test.py

# python test.py
# 输出     verbosity turned off

# python test.py -v
# 输出 verbosity turned on

验证图像大小是每个维度上的stride的倍数，默认是32的倍数
例如 图像大小是1111 那么就是
--img-size [1111, 1111] 更新为 [1120, 1120]

def check_img_size(imgsz, s=32, floor=0):
    # Verify image size is a multiple of stride s in each dimension
    if isinstance(imgsz, int):  # integer i.e. img_size=640
        new_size = max(make_divisible(imgsz, int(s)), floor)
    else:  # list i.e. img_size=[640, 480]
        imgsz = list(imgsz)  # convert to list if tuple
        new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz]
    if new_size != imgsz:
        LOGGER.warning(f'WARNING: --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}')
    return new_size

推理增强部分

def _forward_augment(self, x):
    img_size = x.shape[-2:]  # height, width
    s = [1, 0.83, 0.67]  # scales
    f = [None, 3, None]  # flips (2-ud, 3-lr)
    y = []  # outputs
    for si, fi in zip(s, f):
        xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
        print("xi.shape[2:]:",xi.shape[2:])
        yi = self._forward_once(xi)[0]  # forward
        print("0 yi:",yi.shape)
        #cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # save
        yi = self._descale_pred(yi, fi, si, img_size)
        print("1 yi.shape:",yi.shape)
        y.append(yi)
    y = self._clip_augmented(y)  # clip augmented tails
    return torch.cat(y, 1), None  # augmented inference, train

def _descale_pred(self, p, flips, scale, img_size):
    # de-scale predictions following augmented inference (inverse operation)
    if self.inplace:
        p[..., :4] /= scale  # de-scale
        if flips == 2:
            p[..., 1] = img_size[0] - p[..., 1]  # de-flip ud
        elif flips == 3:
            p[..., 0] = img_size[1] - p[..., 0]  # de-flip lr
    else:
        x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale  # de-scale
        if flips == 2:
            y = img_size[0] - y  # de-flip ud
        elif flips == 3:
            x = img_size[1] - x  # de-flip lr
        p = torch.cat((x, y, wh, p[..., 4:]), -1)
    return p

def _clip_augmented(self, y):
    # Clip YOLOv5 augmented inference tails
    nl = self.model[-1].nl  # number of detection layers (P3-P5)
    g = sum(4 ** x for x in range(nl))  # grid points
    e = 1  # exclude layer count
    i = (y[0].shape[1] // g) * sum(4 ** x for x in range(e))  # indices
    y[0] = y[0][:, :-i]  # large
    i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e))  # indices
    y[-1] = y[-1][:, i:]  # small
    return y

关于翻转看

if self.inplace:
    p[..., :4] /= scale  # de-scale
    if flips == 2:
        p[..., 1] = img_size[0] - p[..., 1]  # de-flip ud
    elif flips == 3:
        p[..., 0] = img_size[1] - p[..., 0]  # de-flip lr

2表示上下翻转
3表示左右翻转
s = [1, 0.83, 0.67] 是缩放比例，且能被32整除

这里的顺序是HW

xi.shape[2:]: torch.Size([640, 480])
xi.shape[2:]: torch.Size([544, 416])
xi.shape[2:]: torch.Size([448, 352])

yi.shape: torch.Size([1, 18900, 85])
yi.shape: torch.Size([1, 13923, 85])
yi.shape: torch.Size([1, 9702, 85])

合并去冗余之后再进NMS

torch.Size([1, 34233, 85])

原来推理一张图像，增强后是推理3张