背景
上一篇博客介绍了 pytorch
的 transforms
模块,有需要的移步 研究一下 pytorch 的 transforms 模块。现在,我想像 T.Normalise()
一样,对每个图片做一下对比度增强,而其他的转换方法保持不变,依旧采用随机处理。
实现方式是 采用 __call__
机制,具体如下:
-
from PIL import ImageEnhance
-
-
class Contrast(object):
-
def __init__(self,degree):
-
self.degree = degree
-
def __call__(self,img):
-
return contrast(img,self.degree)
-
def contrast(img,degree):
-
enh_contrast = ImageEnhance.Contrast(img)
-
enh_contrast.enhance(degree)
-
return img
上面定义了新的数据转换方法:Contrast,使用方法如下,以开源的代码为例 class ChaojieDataset(Dataset):
-
-
#2.define dataset
-
class ChaojieDataset(Dataset):
-
def __init__(self,label_list,transforms=None,train=True,test=False):
-
self.test = test
-
self.train = train
-
imgs = []
-
if self.test:
-
for index,row in label_list.iterrows():
-
imgs.append((row[ "filename"]))
-
self.imgs = imgs
-
else:
-
for index,row in label_list.iterrows():
-
imgs.append((row[ "filename"],row["label"]))
-
self.imgs = imgs
-
if transforms is None:
-
if self.test or not train:
-
self.transforms = T.Compose([
-
T.Resize((config.img_weight,config.img_height)),
-
T.ToTensor(),
-
T.Normalize(mean = [ 0.485,0.456,0.406],
-
std = [ 0.229,0.224,0.225])])
-
else:
-
self.transforms = T.Compose([
-
T.Resize((config.img_weight,config.img_height)),
-
T.RandomRotation( 30),
-
T.RandomHorizontalFlip(),
-
T.RandomVerticalFlip(),
-
T.RandomAffine( 45),
-
Contrast( 1.8), ## 在此添加 ##
-
T.ToTensor(),
-
T.Normalize(mean = [ 0.485,0.456,0.406],
-
std = [ 0.229,0.224,0.225])])
-
else:
-
self.transforms = transforms
-
def __getitem__(self,index):
-
if self.test:
-
filename = self.imgs[index]
-
img = Image.open(filename)
-
img = self.transforms(img)
-
return img,filename
-
else:
-
filename,label = self.imgs[index]
-
img = Image.open(filename)
-
img = self.transforms(img)
-
return img,label
-
def __len__(self):
-
return len(self.imgs)
这样,每一张图片除了进行归一化和转变成张量外,都做了对比度增强,当然也可以设置一个随机数,每次随机选择要增强的对比度。