pytroch 颜色增强ColorJitter，墙裂推荐

目录

函数参数解释：

随机亮度测试，非常方便，墙裂推荐：

单项测试：

举例：

yolov5颜色增强示例，效果差不多，opencv的：

---

函数参数解释：

函数名：
torchvision.transforms.ColorJitter(brightness=0, contrast=0, saturation=0, hue=0)
函数解析：
随机改变一个图像的亮度、对比度、饱和度和色调。如果图像是 tensor，那么它的 shape 为[…，1或3，H，W]，其中…表示 batch。如果图像是PIL图像，那么不支持模式 “1”、“I”、"F "和带有透明度（alpha通道）的模式。

参数：
brightness (类型为 float 或 tuple: float (min, max)) - 亮度的偏移程度。 brightness_factor可以是 [max(0, 1 - brightness), 1 + brightness]，也可以直接给出最大、最小值的范围 [min, max]，然后从中随机采样。brightness_factor 值应该是非负数。

contrast (类型为 float 或 tuple: float (min, max)) - 对比度的偏移程度。 contrast_factor 可以是 [max(0, 1 - contrast), 1 + contrast]，也可以直接给出最大、最小值的范围 [min, max]，然后从中随机采样。contrast_factor 值应该是非负数。

saturation (类型为 float 或 tuple: float (min, max)) - 饱和度的偏移程度。 saturation_factor 可以是 [max(0, 1 - saturation), 1 + saturation]，也可以直接给出最大、最小值的范围 [min, max]，然后从中随机采样。saturation_factor 值应该是非负数。

hue (类型为 float 或 tuple: float (min, max)) - 色调的偏移程度。hue_factor 可以是 [-hue, hue]，也可以直接给出最大、最小值的范围 [min, max]，然后从中随机采样，它的值应当满足 0<= hue <= 0.5 或者 -0.5<= min <= max <= 0.5。为了使色调偏移，输入图像的像素值必须是非负值，以便转换到 HSV 颜色空间。因此，如果将图像归一化到一个有负值的区间，或者在使用这个函数之前使用会产生负值的插值方法，那么它就不会起作用。

随机亮度测试，非常方便，墙裂推荐：

参数是测试过的经验值

import cv2
import numpy as np
import torch
from PIL import Image

from torchvision.transforms import ColorJitter
import random


class CustomColorJitter(ColorJitter):
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        super(CustomColorJitter, self).__init__(brightness=brightness, contrast=contrast, saturation=saturation, hue=hue)

    def get_params(self, brightness, contrast, saturation, hue):
        self.last_brightness =None

        self.bright_param=None
        if brightness is not None:
            self.last_brightness = brightness[0] + random.uniform(0, 1) * (brightness[1] - brightness[0])
            self.bright_param=(self.last_brightness,self.last_brightness)
        self.contrast_param = None
        if contrast is not None:
            self.last_contrast = contrast[0] + random.uniform(0, 1) * (contrast[1] - contrast[0])
            self.contrast_param=(self.last_contrast,self.last_contrast)
        self.saturation_param = None
        if saturation is not None:
            self.last_saturation = saturation[0] + random.uniform(0, 1) * (saturation[1] - saturation[0])
            self.saturation_param=(self.last_saturation,self.last_saturation)
        self.hue_param=None
        if hue is not None:
            self.last_hue = hue[0] + random.uniform(0, 1) * (hue[1] - hue[0])
            self.hue_param=(self.last_hue,self.last_hue)

        return super().get_params(brightness=self.bright_param, contrast=self.contrast_param,
                                      saturation=self.saturation_param, hue=self.hue_param)


img_path = "./aaa.png"

debug=True

if debug:
    transform = CustomColorJitter(brightness=[0.6, 1.3], contrast=[0.5, 1.5], saturation=[0.5, 1.5], hue=[-0.02, 0.02])
    # transform = CustomColorJitter(hue=[-0.02, 0.02])
    # transform = CustomColorJitter(saturation=[0.5, 1.5])
    # transform = CustomColorJitter( contrast=[0.5, 1.5])
    # transform = CustomColorJitter( brightness=[0.7, 1.3])
else:
    transform = ColorJitter(brightness=[0.6, 1.3], contrast=[0.5, 1.5], saturation=[0.5, 1.5], hue=[-0.02, 0.02])


while True:
    img = cv2.imread(img_path)
    pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    image = transform(pil_img)
    if debug:
        if transform.bright_param is not None:
            print("Last brightness value:", transform.bright_param)
        if transform.contrast_param is not None:
            print("Last contrast value:", transform.contrast_param)
        if transform.saturation_param is not None:
            print("Last saturation value:", transform.saturation_param)
        if transform.hue_param is not None:
            print("Last hue value:", transform.hue_param)

    img_cv = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
    cv2.imshow("img_o", img)
    cv2.imshow("img_cv", img_cv)
    cv2.waitKey(0)

单项测试：

import cv2
import numpy as np
import torch
import torchvision.transforms as f
from PIL import Image


from torchvision.transforms import ColorJitter
import random

class CustomColorJitter(ColorJitter):
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        super(CustomColorJitter, self).__init__(brightness=brightness, contrast=contrast, saturation=saturation, hue=hue)

    def get_params(self, brightness, contrast, saturation, hue):
        self.last_brightness = brightness[0] + random.uniform(0, 1) * (brightness[1] - brightness[0])
        return super().get_params(brightness=(self.last_brightness, self.last_brightness), contrast=contrast, saturation=saturation, hue=hue)

img_path = "./aaa.png"

transform = CustomColorJitter(brightness=[0.5, 1.5])

while True:
    img = cv2.imread(img_path)
    pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    image = transform(pil_img)

    print("Last brightness value:", transform.last_brightness)

    img_cv = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
    cv2.imshow("img_o", img)
    cv2.imshow("img_cv", img_cv)
    cv2.waitKey(0)

举例：

以下内容转自：https://blog.csdn.net/lxhRichard/article/details/128083192
1. 以随机亮度为例

import torch
import torchvision.transforms as f
from PIL import Image

img_path = "./1.jpg"
img = Image.open(img_path)
trans = f.ColorJitter(brightness=[0.01,0.05])
image = trans(img)
image.show()

输出对比：

2. 以随机对比度为例

import torch
import torchvision.transforms as f
from PIL import Image

img_path = "./1.jpg"
img = Image.open(img_path)
trans = f.ColorJitter(contrast=[0.3,0.6])
image = trans(img)
image.show()

输出对比：

3. 以随机饱和度为例

import torch
import torchvision.transforms as f
from PIL import Image

img_path = "./1.jpg"
img = Image.open(img_path)
trans = f.ColorJitter(saturation=[0.2,0.5])
image = trans(img)
image.show()

输出对比：

4. 以随机色调为例

import torch
import torchvision.transforms as f
from PIL import Image

img_path = "./1.jpg"
img = Image.open(img_path)
trans = f.ColorJitter(hue=[-0.1,0.2])
image = trans(img)
image.show()

输出对比：

5. 综合调整：

import torch
import torchvision.transforms as f
from PIL import Image

img_path = "./1.jpg"
img = Image.open(img_path)
trans = f.ColorJitter(brightness=0.6, contrast=0.7, saturation=0.5, hue=0.1)
image = trans(img)
image.show()

输出对比：

官方文档链接：https://pytorch.org/vision/stable/generated/torchvision.transforms.ColorJitter.html?highlight=transforms+colorjitter#torchvision.transforms.ColorJitter
 

yolov5颜色增强示例，效果差不多，opencv的：

import cv2
import numpy as np

def augment_hsv(img, h_gain=0.015, s_gain=0.7, v_gain=0.4):
    r = np.random.uniform(-1, 1, 3) * [h_gain, s_gain, v_gain] + 1  # random gains
    hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
    print(r[0], r[1], r[2])
    dtype = img.dtype  # uint8

    x = np.arange(0, 256, dtype=np.int16)
    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)

    img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype)
    cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)  # no return needed

if __name__ == '__main__':

    img_path = "./aaa.png"
    while True:
        img_o = cv2.imread(img_path)
        img=img_o.copy()
        augment_hsv(img)
        cv2.imshow("img_o", img_o)
        cv2.imshow('HSV Augmented Image', img)
        cv2.waitKey(0)