custom_pytorch_data_augmentation.py

# https://www.jianshu.com/p/9a9000d226b6
# https://stackoverflow.com/questions/64420379/is-it-possible-to-add-own-function-in-transform-compose-in-pytorch

import cv2
import torchvision.transforms as transforms

class hisEqulColor(object):
    def __call__(self, img):
        img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)   # PyTorch的Transform要求PIL格式的圖片，若要使用OpenCV則需先進行轉換
        ycrcb = cv2.cvtColor(img, cv2.COLOR_BGR2YCR_CB)
        channels = cv2.split(ycrcb)
        cv2.equalizeHist(channels[0], channels[0])
        cv2.merge(channels, ycrcb)
        cv2.cvtColor(ycrcb, cv2.COLOR_YCR_CB2RGB, img)           # OpenCV輸出成PIL格式圖片要轉回RGB
        return Image.fromarray(img)
    def __repr__(self):
        return self.__class__.__name__+'()'

class CLAHE(object):
    def __call__(self, img):
        img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        lab_planes = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=2.0,tileGridSize=(8,8))
        lab_planes[0] = clahe.apply(lab_planes[0])
        lab = cv2.merge(lab_planes)
        img = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
        return Image.fromarray(img)
    def __repr__(self):
        return self.__class__.__name__+'()'

# 中間不相關之程式碼省略
    
transform_test = transforms.Compose([hisEqulColor(),
                                     #CLAHE(),
                                     transforms.Resize([224, 224]),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.RandomVerticalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                                    ])