import os import random import numpy as np from skimage.filters import gaussian import torch from PIL import Image """ color map 0=background, 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle # 6=bus, 7=car, 8=cat, 9=chair, 10=cow, 11=diningtable, 12=dog, 13=horse, 14=motorbike, 15=person # 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor """ palette = [0, 0, 0, 128, 0, 0, 0, 128, 0, 128, 128, 0, 0, 0, 128, 128, 0, 128, 0, 128, 128, 128, 128, 128, 64, 0, 0, 192, 0, 0, 64, 128, 0, 192, 128, 0, 64, 0, 128, 192, 0, 128, 64, 128, 128, 192, 128, 128, 0, 64, 0, 128, 64, 0, 0, 192, 0, 128, 192, 0, 0, 64, 128] zero_pad = 256 * 3 - len(palette) for i in range(zero_pad): palette.append(0) def colorize_mask(mask): # mask: numpy array of the mask new_mask = Image.fromarray(mask.astype(np.uint8)).convert('P') new_mask.putpalette(palette) return new_mask def make_dataset(mode, root): assert mode in ['train', 'val', 'validate', 'test', 'inference'] items = [] if mode == 'train': img_path = os.path.join(root, 'benchmark_RELEASE', 'dataset', 'img') mask_path = os.path.join(root, 'benchmark_RELEASE', 'dataset', 'cls') data_list = [l.strip('\n') for l in open(os.path.join( root, 'benchmark_RELEASE', 'dataset', 'train.txt')).readlines()] for it in data_list: item = (os.path.join(img_path, it + '.jpg'), os.path.join(mask_path, it + '.mat')) items.append(item) elif mode == 'val' or mode == 'validate': img_path = os.path.join(root, 'VOCdevkit', 'VOC2012', 'JPEGImages') mask_path = os.path.join(root, 'VOCdevkit', 'VOC2012', 'SegmentationClass') data_list = [l.strip('\n') for l in open(os.path.join( root, 'VOCdevkit', 'VOC2012', 'ImageSets', 'Segmentation', 'seg11valid.txt')).readlines()] for it in data_list: item = (os.path.join(img_path, it + '.jpg'), os.path.join(mask_path, it + '.png')) items.append(item) elif mode == 'test' or mode == 'inference': img_path = os.path.join(root, 'VOCdevkit (test)', 'VOC2012', 'JPEGImages') data_list = [l.strip('\n') for l in open(os.path.join( root, 'VOCdevkit (test)', 'VOC2012', 'ImageSets', 'Segmentation', 'test.txt')).readlines()] for it in data_list: items.append((img_path, it)) else: raise Exception("Please choose proper mode for data") return items class RandomVerticalFlip(object): def __call__(self, img): if random.random() < 0.5: return img.transpose(Image.FLIP_TOP_BOTTOM) return img class DeNormalize(object): def __init__(self, mean, std): self.mean = mean self.std = std def __call__(self, tensor): for t, m, s in zip(tensor, self.mean, self.std): t.mul_(s).add_(m) return tensor class MaskToTensor(object): def __call__(self, img): return torch.from_numpy(np.array(img, dtype=np.int32)).long() class FreeScale(object): def __init__(self, size, interpolation=Image.BILINEAR): self.size = tuple(reversed(size)) # size: (h, w) self.interpolation = interpolation def __call__(self, img): return img.resize(self.size, self.interpolation) class FlipChannels(object): def __call__(self, img): img = np.array(img)[:, :, ::-1] return Image.fromarray(img.astype(np.uint8)) class RandomGaussianBlur(object): def __call__(self, img): sigma = 0.15 + random.random() * 1.15 blurred_img = gaussian(np.array(img), sigma=sigma, multichannel=True) blurred_img *= 255 return Image.fromarray(blurred_img.astype(np.uint8))