训练数据不够怎么造？yolo5 最有用的trick 之数据增强详解

在深度学习领域，模型参数量往往数以百万计，因此要训练一个鲁棒的模型往往需要大量的数据，与此同时，我们还可以利用各种数据增强的手段来增加模型训练数据的多样性，令模型尽可能适应不同的场景，防止过拟合的发生。

yolo5是Ultralytics公司于2020年开源的目标检测算法https://github.com/ultralytics/yolov5，算法在backbone，bbox asign，loss function，数据增强等方面，集成了当时最有用的trick，使模型在速度与精度上都要优于当时同期的其他模型，这也使得yolo5在工业界、学术界及各路竞赛中成为了研究者优先尝试的检测模型。

yolo5应用的trick繁多，此本文将针对yolo5的数据增强方面进行详细解读。

一、YOLO5的数据增强

 NO.1 

Albumentations增强

yolo5使用了Albumentations开源模块中实现的诸个数据增强方法，包括了Blur、MedianBlur、ToGray、CLAHE、RandomBrightnessContrast、RandomGamma、ImageCompression，这些数据增强方法的解释、训练中使用的概率如下表所示：

（使用概率是指在使用了Albumentations增强的情况下；上面的各种增强手段的概率值为yolo5作者网格搜索得出，可能随着yolo5项目的更新而发生变化。）

● 代码详解

在yolo5中，Albumentations增强的核心代码可以简化为：

（简化版旨在讲解逻辑，原始版本请访问原代码[1]）

import albumentations as A

# 1. 将所有的数据增强对象装填至列表T中
T = [
    A.Blur(p=0.01),               # Blur数据增强，使用概率为0.01               
    A.MedianBlur(p=0.01),         # MedianBlur数据增强，使用概率为0.01
    A.ToGray(p=0.01),             # 转换为灰度图，使用概率为0.01
    A.CLAHE(p=0.01),              # CLAHE直方图均衡，使用概率为0.01
    A.RandomBrightnessContrast(p=0.0),   # 随机改变亮度与对比度，不使用
    A.RandomGamma(p=0.0),         # 随机灰度系数，不使用
    A.ImageCompression(quality_lower=75, p=0.0)]  # 对图像进行压缩，最低质量为75，不使用
# 2. 调用Compose类，将列表T中的所有增强对象合成一个对象
transform = A.Compose(T, bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))  # 将列表`T`转换为一个组合数据增强对象`transform`

def augment(im, labels, p=1.0):
    if random.random() < p:
    # 3. 调用transform增强，使用的概率为 p
        new = transform(image=im, bboxes=labels[:, 1:], class_labels=labels[:, 0])   # 调用`transform`对图像进行变换，实现数据增强
        im, labels = new['image'], np.array([[c, *b] for c, b in zip(new['class_labels'], new['bboxes'])])
    return im, labels

（左右滑动查看）

上述代码的功能为：

1. 创建数据增强对象列表T，将上述提及的数据增强方法装填至列表T当中；

2. 调用A.Compose方法，将列表T转换为一个组合数据增强对象transform；

3. 在augment函数中，调用transform对图像进行变换，实现数据增强。

● 可视化

本节将对上述提及的数据增强方法进行可视化，并提供demo代码：

import albumentations as A
from PIL import Image
import math
import numpy as np
import matplotlib.pyplot as plt

T = [
    A.Blur(p=1),
    A.MedianBlur(p=1),
    A.ToGray(p=1),
    A.CLAHE(p=1),
    A.RandomBrightnessContrast(p=1),
    A.RandomGamma(p=1),
    A.ImageCompression(quality_lower=75, p=1)]
img = Image.open('lena.jpg')

plt.figure(figsize=(64, 32))
plt.subplot(2, math.ceil(len(T) / 2), 1)
plt.title("original image", fontsize=50)
plt.imshow(img)
for i in range(len(T)):
    title = T[i].__class__.__name__
    plt.subplot(2, math.ceil(len(T) / 2), i + 2)
    plt.title(title, fontsize=50)
    new_img = Image.fromarray(T[i](image=np.array(img))['image'])
    plt.imshow(new_img)
plt.show()

（左右滑动查看）

执行结果：

● 小节

在yolo5中，实际上只使用了Albumentations中的Blur，MedianBlur，CLAHE，ToGray这四种数据增强，并且使用的概率都是0.01。这些数据增强有各自的应用场景：

1. Blur与MedianBlur都是对图像进行模糊处理，但二者有细微的差别，Blur使用了均值模糊，而MedianBlur使用了中值模糊，后者相对于前者更不易受到局部极值点的影响，对异常值的敏感性更低。二者都可以模拟图像模糊的场景，有利于提升模型对于模糊图像的识别效果；

2. ToGray增强的功能为将图像从RGB转换为灰度图，这种数据增强的方法旨在令模型能够不依赖颜色信息就对图像做出预测，使用该数据增强可以促进模型从图像的内容的形状信息而非色彩信息提取信息，可以提升模型的对于色度等信息的泛化能力；

3. CLAHE的作用为对图像进行CLAHE直方图均衡，该操作可以令图像的灰度图的分布更加均衡，令对比度较低的图像得到增强，细节信息更加明显。使用该数据增强方法的作用即为对图像本身进行增强，令其内容更加明显。如果数据集中存在图像灰度普遍偏大或普遍偏小的情况，可以考虑使用该数据增强对图像进行预处理。在这里yolo5使用CLAHE的原因仅仅是为了令数据更加多变，提升模型的泛化能力。

4. 上述的几种数据增强的手段的功能互相垂直，可以叠加使用，且顺序也可以任意调整。

 NO.2 

augment hsv

yolo5实现了自己的HSV增强，首先将原图从RGB空间转换到HSV空间，然后再分别再H（hue，色相），S（saturation，饱和度），V（value，亮度）三个通道分别增强或减弱随机的值，从而简便地模拟出不同场景、光照的情况，提升模型的泛化能力。

yolo5的实现方式很简单，代码为：

def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5):
    # HSV color-space augmentation
    if hgain or sgain or vgain:
        r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
        hue, sat, val = cv2.split(cv2.cvtColor(im, cv2.COLOR_BGR2HSV))
        dtype = im.dtype  # uint8

        x = np.arange(0, 256, dtype=r.dtype)
        lut_hue = ((x * r[0]) % 180).astype(dtype)
        lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
        lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

        im_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
        cv2.cvtColor(im_hsv, cv2.COLOR_HSV2BGR, dst=im)  # no return needed

（左右滑动查看）

下面本文将对一些主要步骤进行详解。

● 代码详解

1.函数参数解释：

- im：np.array格式的图像，通道为BGR

- hgain：H通道的增益，以0.5举例，原图的H通道的值的变换范围会是原来的50%~150%

- sgain：S通道的增益

- vgain：V通道的增益

2. 函数解释：

a. 首先函数通过下面的代码随机地对HSV三个通道的增益值进行采样，H通道的采样的范围为 。 S，V通道同理。

np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains

（左右滑动查看）

b. 通过下面的代码将原图从BGR格式转换为HSV空间的图像。并将图像的HSV通道分别赋值给hue，sat，val三个变量。

sat, val = cv2.split(cv2.cvtColor(im, cv2.COLOR_BGR2HSV))

（左右滑动查看）

c. 根据采样的增益值对HSV 3个通道增加/减小相应的值。yolo5的实现方式是创建一个查找表，将通过查找表将原值映射为新值，创建3个通道的查找表的代码如下：

x = np.arange(0, 256, dtype=r.dtype)
lut_hue = ((x * r[0]) % 180).astype(dtype)
lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

（左右滑动查看）

d. 上述代码分别将创建了H，S，V三个通道将原值映射至增加了增量后的值的查找表，映射关系分别为：

H： 

S：

V：

e. 利用下面的代码，将原值映射为新值，并将hue，sat，val三个通道合并为HSV空间的图像：

sv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))

（左右滑动查看）

f. 将HSV空间的图像转回BGR格式的图像：

cvtColor(im_hsv, cv2.COLOR_HSV2BGR, dst=im)  # no return needed

（左右滑动查看）

● 可视化

下面将通过代码对HSV增强的结果进行可视化：

import numpy as np
import matplotlib.pyplot as plt
import cv2


def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5):

    r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
    hue, sat, val = cv2.split(cv2.cvtColor(im, cv2.COLOR_RGB2HSV))
    dtype = im.dtype  # uint8

    x = np.arange(0, 256, dtype=r.dtype)
    lut_hue = ((x * r[0]) % 180).astype(dtype)
    lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
    lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

    im_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
    result = cv2.cvtColor(im_hsv, cv2.COLOR_HSV2RGB)  # no return needed
    return result


img = cv2.imread("./lena.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(100, 32))
plt.subplot(1, 4, 1)
plt.title("original image", fontsize=80)
plt.imshow(img)

img_h_aug = augment_hsv(img, 0.5, 0, 0)
plt.subplot(1, 4, 2)
plt.title("hue augment", fontsize=80)
plt.imshow(img_h_aug)

img_s_aug = augment_hsv(img, 0, 0.5, 0)
plt.subplot(1, 4, 3)
plt.title("saturation augment", fontsize=80)
plt.imshow(img_s_aug)

img_v_aug = augment_hsv(img, 0, 0, 0.5)
plt.subplot(1, 4, 4)
plt.title("value augment", fontsize=80)
plt.imshow(img_v_aug)

plt.show()

（左右滑动查看）

执行结果：

● 小节

Yolo 5使用hsv增强的目的是令模型在训练过程中看到的数据更加的多样，而通过HSV增强获得的”多样性“也可以从3个角度来说：

1. 色调（Hue）多样：通过随机地调整色调可以模拟不同颜色风格的输入图像，比如不同滤镜，不同颜色光照等场景下的图像，从而提升模型在这些场景下的泛化能力；

2. 饱和度（Saturation）多样：通过随机调整饱和度可以提升模型对不同鲜艳程度的目标的识别的泛化能力；

3. 亮度（Value）多样：通过随机调整亮度可以提升模型应对不同光亮场景下的输入图像。

HSV增强在目标检测模型的训练中是非常常用的方法，它在不破坏图像中关键信息的前提下提高了数据集的丰富程度，且计算成本很低，是很实用的数据增强方法。

 NO.3 

hist equalize

yolo5还使用了直方图均衡对原图进行数据增强，对图像的亮度进行了直方图均衡操作，可选择的均衡方法有原始的直方图均衡，与CLAHE直方图均衡。

通过该数据增强，图像中的明暗分布会更加均衡，具体的代码解析如下：

def hist_equalize(im, clahe=True, bgr=False):
    # Equalize histogram on BGR image 'im' with im.shape(n,m,3) and range 0-255
    yuv = cv2.cvtColor(im, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)
    if clahe:
        c = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        yuv[:, :, 0] = c.apply(yuv[:, :, 0])
    else:
        yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0])  # equalize Y channel histogram
    return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR if bgr else cv2.COLOR_YUV2RGB)  # convert YUV image to RGB

（左右滑动查看）

● 参数介绍

- im：np.array格式的图像，通道为BGR或RGB

- clahe：是否使用CLAHE直方图均衡

- bgr：输入图像im是BGR还是RGB格式，如果为BGR格式则bgr=True

● 函数介绍

1. 首先通过下面的代码将原图从RGB/BGR格式转换到YUV空间

yuv = cv2.cvtColor(im, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)

2. 如果clahe为True，则使用下面的代码将图像的Y通道进行CLAHE直方图均衡

cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
yuv[:, :, 0] = c.apply(yuv[:, :, 0])

3. 如果clahe为False，则使用常规的直方图均衡

yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0])  # equalize Y channel histogram

4. 最终将直方图均衡后的图像转回RGB/BGR通道并返回

return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR if bgr else cv2.COLOR_YUV2RGB)  # convert YUV image to RGB

（左右滑动查看）

● 可视化

下面将对hist equalize进行可视化展示其效果，代码为：

import matplotlib.pyplot as plt
import cv2


def hist_equalize(im, clahe=True, bgr=False):
    # Equalize histogram on BGR image 'im' with im.shape(n,m,3) and range 0-255
    yuv = cv2.cvtColor(im, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)
    if clahe:
        c = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        yuv[:, :, 0] = c.apply(yuv[:, :, 0])
    else:
        yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0])  # equalize Y channel histogram
    return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR if bgr else cv2.COLOR_YUV2RGB)  # convert YUV image to RGB


img = cv2.imread("./lena.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(70, 32))
plt.subplot(1, 3, 1)
plt.title("original image", fontsize=80)
plt.imshow(img)

img_aug1 = hist_equalize(img, clahe=False)
plt.subplot(1, 3, 2)
plt.title("Hist Equalize", fontsize=80)
plt.imshow(img_aug1)

img_s_aug = hist_equalize(img, clahe=True)
plt.subplot(1, 3, 3)
plt.title("CLAHE Hist Equalize", fontsize=80)
plt.imshow(img_s_aug)

plt.show()

（左右滑动查看）

可视化结果为：

● 小节

直方图均衡操作的目的是令图像的明暗分布更加合理，从而提升那些过明或过暗的图像的对比度，解决图像光照不均匀的问题，所以当我们的数据集存在光亮过明或过暗的情况下，可以尝试使用直方图均衡操作，在输入模型之前进行预处理。

但在yolo 5中，模型训练过程中仅有1%的概率会进行直方图均衡，所以这里的目的仅仅只是提升数据集的丰富度，提高模型的泛化能力。

二、总结

本文主要简单地介绍了yolo5的3个数据增强方法，并对他们进行了代码的讲解与可视化的展示，希望可以帮助读者对这些数据增强方法有更加直观的认识。yolo5的其他数据增强方法，比如cutout，mixup，mosaic等都会在后续的文章中讲解。

参考文献

[1]https://github.com/ultralytics/yolov5/blob/master/utils/augmentations.py

[2]Reza A M. Realization of the contrast limited adaptive histogram equalization (CLAHE) for real-time image enhancement[J]. Journal of VLSI signal processing systems for signal, image and video technology, 2004, 38(1): 35-44.

作者丨Hanx

在AI路上不断摸爬滚打的调包侠

以上就是本次分享，获取海量数据集资源，请访问OpenDataLab官网；获取更多开源工具及项目，请访问OpenDataLab Github空间。另外还有哪些想看的内容，快来告诉小助手吧。更多数据集上架动态、更全面的数据集内容解读、最牛大佬在线答疑、最活跃的同行圈子……欢迎添加微信opendatalab_yunying加入OpenDataLab官方交流群。