import torch import torchvision.transforms as transforms import torch.utils.data as data from util import task from .image_folder import make_dataset import random import numpy as np import copy import skimage.morphology as sm from PIL import Image, ImageFile, ImageOps ImageFile.LOAD_TRUNCATED_IMAGES = True ###################################################################################### # Create the dataloader ###################################################################################### class CreateDataset(data.Dataset): def __init__(self, opt): self.opt = opt self.img_paths, self.img_size = make_dataset(opt.img_file) if opt.mask_file != 'none': # load the random mask files for training and testing self.mask_paths, self.mask_size = make_dataset(opt.mask_file) self.transform = get_transform(opt, convert=False, augment=False) fixed_opt = copy.deepcopy(opt) fixed_opt.preprocess = 'scale_longside' fixed_opt.load_size = fixed_opt.fixed_size fixed_opt.no_flip = True self.transform_fixed = get_transform(fixed_opt, convert=True, augment=False) def __len__(self): """return the total number of examples in the dataset""" return self.img_size def __getitem__(self, item): """return a data point and its metadata information""" # load the image and conditional input img_org, img, img_path = self._load_img(item) if self.opt.batch_size > 1: # padding the image to the same size for batch training img_org = transforms.functional.pad(img_org, (0, 0, self.opt.fine_size-self.img_h, self.opt.fine_size-self.img_w)) img = transforms.functional.pad(img, (0, 0, self.opt.fixed_size - img.size(-1), self.opt.fixed_size - img.size(-2))) pad_mask = torch.zeros_like(img_org) pad_mask[:, :self.img_w, :self.img_h] = 1 # load the mask mask, mask_type = self._load_mask(item, img_org) if self.opt.reverse_mask: if self.opt.isTrain: mask = 1 - mask if random.random() > 0.8 else mask else: mask = 1 - mask return {'img_org': img_org, 'img': img, 'img_path': img_path, 'mask': mask, 'pad_mask': pad_mask} def name(self): return "" def _load_img(self, item): """load the original image and preprocess image""" img_path = self.img_paths[item % self.img_size] img_pil = Image.open(img_path).convert('RGB') img_org = self.transform(img_pil) img = self.transform_fixed(img_org) img_org = transforms.ToTensor()(img_org) img_pil.close() self.img_c, self.img_w, self.img_h = img_org.size() return img_org, img, img_path def _mask_dilation(self, mask): """mask erosion for different region""" mask = np.array(mask) pixel = np.random.randint(3, 25) mask = sm.erosion(mask, sm.square(pixel)).astype(np.uint8) return mask def _load_mask(self, item, img): """load the mask for image completion task""" c, h, w = img.size() if isinstance(self.opt.mask_type, list): mask_type_index = random.randint(0, len(self.opt.mask_type) - 1) mask_type = self.opt.mask_type[mask_type_index] else: mask_type = self.opt.mask_type if mask_type == 0: # center mask if random.random() > 0.3 and self.opt.isTrain: return task.random_regular_mask(img), mask_type # random regular mask return task.center_mask(img), mask_type elif mask_type == 1: # random regular mask return task.random_regular_mask(img), mask_type elif mask_type == 2: # random irregular mask return task.random_irregular_mask(img), mask_type elif mask_type == 3: # external mask from "Image Inpainting for Irregular Holes Using Partial Convolutions (ECCV18)" if self.opt.isTrain: mask_index = random.randint(0, self.mask_size-1) mask_transform = transforms.Compose( [ transforms.RandomHorizontalFlip(), transforms.RandomRotation(10), transforms.RandomCrop([self.opt.fine_size + 64, self.opt.fine_size + 64]), transforms.Resize([h, w]) ] ) else: mask_index = item mask_transform = transforms.Compose( [ transforms.Resize([h, w]) ] ) mask_pil = Image.open(self.mask_paths[mask_index]).convert('L') mask = mask_transform(mask_pil) mask_pil.close() if self.opt.isTrain: mask = self._mask_dilation(mask) else: mask = np.array(mask) < 128 mask = torch.tensor(mask).view(1, h, w).float() return mask, mask_type else: raise NotImplementedError('mask type [%s] is not implemented' % str(mask_type)) def dataloader(opt): datasets = CreateDataset(opt) dataset = data.DataLoader(datasets, batch_size=opt.batch_size, shuffle=not opt.no_shuffle, num_workers=int(opt.nThreads), drop_last=True) return dataset ###################################################################################### # Basic image preprocess function ###################################################################################### def _make_power_2(img, power, method=Image.BICUBIC): """resize the image to the size of log2(base) times""" ow, oh = img.size base = 2 ** power nw, nh = int(max(1, round(ow / base)) * base), int(max(1, round(oh / base)) * base) if nw == ow and nh == oh: return img return img.resize((nw, nh), method) def _random_zoom(img, target_width, method=Image.BICUBIC): """random resize the image scale""" zoom_level = np.random.uniform(0.8, 1.0, size=[2]) ow, oh = img.size nw, nh = int(round(max(target_width, ow * zoom_level[0]))), int(round(max(target_width, oh * zoom_level[1]))) return img.resize((nw, nh), method) def _scale_shortside(img, target_width, method=Image.BICUBIC): """resize the short side to the target width""" ow, oh = img.size shortsize = min(ow, oh) scale = target_width / shortsize return img.resize((round(ow * scale), round(oh * scale)), method) def _scale_longside(img, target_width, method=Image.BICUBIC): """resize the long side to the target width""" ow, oh = img.size longsize = max(ow, oh) scale = target_width / longsize return img.resize((round(ow * scale), round(oh * scale)), method) def _scale_randomside(img, target_width, method=Image.BICUBIC): """resize the side to the target width with random side""" if random.random() > 0.5: return _scale_shortside(img, target_width, method) else: return _scale_longside(img, target_width, method) def _crop(img, pos=None, size=None): """crop the image based on the given pos and size""" ow, oh = img.size if size is None: return img nw = min(ow, size) nh = min(oh, size) if (ow > nw or oh > nh): if pos is None: x1 = np.random.randint(0, int(ow-nw)+1) y1 = np.random.randint(0, int(oh-nh)+1) else: x1, y1 = pos return img.crop((x1, y1, x1 + nw, y1 + nh)) return img def _pad(img): """expand the image to the square size""" ow, oh = img.size size = max(ow, oh) return ImageOps.pad(img, (size, size), centering=(0, 0)) def _flip(img, flip): if flip: return img.transpose(Image.FLIP_LEFT_RIGHT) return img def get_transform(opt, params=None, method=Image.BICUBIC, convert=True, augment=False): """get the transform functions""" transforms_list = [] if 'resize' in opt.preprocess: osize = [opt.load_size, opt.load_size] transforms_list.append(transforms.Resize(osize)) elif 'scale_shortside' in opt.preprocess: transforms_list.append(transforms.Lambda(lambda img: _scale_shortside(img, opt.load_size, method))) elif 'scale_longside' in opt.preprocess: transforms_list.append(transforms.Lambda(lambda img: _scale_longside(img, opt.load_size, method))) elif "scale_randomside" in opt.preprocess: transforms_list.append(transforms.Lambda(lambda img: _scale_randomside(img, opt.load_size, method))) if 'zoom' in opt.preprocess: transforms_list.append(transforms.Lambda(lambda img: _random_zoom(img, opt.load_size, method))) if 'crop' in opt.preprocess and opt.isTrain: transforms_list.append(transforms.Lambda(lambda img: _crop(img, size=opt.fine_size))) if 'pad' in opt.preprocess: transforms_list.append(transforms.Lambda(lambda img: _pad(img))) # padding image to square transforms_list.append(transforms.Lambda(lambda img: _make_power_2(img, opt.data_powers, method))) if not opt.no_flip and opt.isTrain: transforms_list.append(transforms.RandomHorizontalFlip()) if augment and opt.isTrain: transforms_list.append(transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2)) if convert: transforms_list.append(transforms.ToTensor()) return transforms.Compose(transforms_list)