Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
# 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])
])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment