object_removal/TFill/dataloader/data_loader.py

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)