CV — 数据增强：yolov5 HSV色调变换

CV — 数据增强：yolov5 HSV色调变换

一、前言

• RGB颜色空间更加适合于显示系统，而 HSV 更加适合于图像处理，大多数做图像处理这一块的都会运用 HSV 颜色空间，因为HSV颜色空间表达起来更加直观。
• RGB 与 HSV 颜色空间区别，可参考资料：https://blog.csdn.net/pentiumCM/article/details/119145452

二、代码实现

• 实现方案：
  参照 yolov5 源码、
  通过在 HSV 色彩空间中，对 h、s、v 三个通道增加扰动，来进行色调变换，达到数据增强的目的。

• 注意点：

  HSV色彩空间中，三通道取值范围为，h：【0 ~ 360】，s：【0 ~ 1】，v：【0 ~ 1】
  在 opencv 中，hsv 取值范围不相同，h：【0 ~ 179】，s：【0 ~ 255】，v：【0 ~ 255】

• 完整代码：

  import cv2
  import numpy as np
  
  def augment_hsv(image, hgain=0.5, sgain=0.5, vgain=0.5):
      """
      HSV color-space augmentation
      :param image:       待增强的图片
      :param hgain:       HSV 中的 h 扰动系数，yolov5：0.015
      :param sgain:       HSV 中的 s 扰动系数，yolov5：0.7
      :param vgain:       HSV 中的 v 扰动系数，yolov5：0.4
      :return:
      """
      if hgain or sgain or vgain:
          # 随机取-1到1三个实数，乘以 hsv 三通道扰动系数
          # r:【1-gain ~ 1+gain】
          r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
  
          image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
          # cv2.split：通道拆分
          # h:[0~180], s:[0~255], v:[0~255]
          hue, sat, val = cv2.split(image_hsv)
          dtype = image.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)
  
          # cv2.LUT：dst(I) = lut(src(I) + d)，d为常数0 / 128
          hue = cv2.LUT(hue, lut_hue)
          sat = cv2.LUT(sat, lut_sat)
          val = cv2.LUT(val, lut_val)
  
          # 通道合并
          image_hsv = cv2.merge((hue, sat, val)).astype(dtype)
  
          # 将 hsv 格式转为 BGR 格式
          image_dst = cv2.cvtColor(image_hsv, cv2.COLOR_HSV2BGR)
          return image_dst
      else:
          return image
      
  if __name__ == '__main__':
      image = cv2.imread('000007.jpg')
  
      rad = np.random.uniform(size=10)
  
      for idx in range(0, 50):
          cv2.imshow('org_img', image)
  
          img_hsv = augment_hsv(image, hgain=0.015, sgain=0.7, vgain=0.4)
          cv2.imshow('img_hsv: %d' % idx, img_hsv)
  
          cv2.waitKey(0)
          cv2.destroyAllWindows()
  

• 效果展示：

  左侧为原图，右侧为色调变化之后的图