MLBench/ReaLens

1"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.2 3It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.4"""5 6import random7import numpy as np8import torch.utils.data as data9from PIL import Image10import torchvision.transforms as transforms11from abc import ABC, abstractmethod12 13 14class BaseDataset(data.Dataset, ABC):15    """This class is an abstract base class (ABC) for datasets.16 17    To create a subclass, you need to implement the following four functions:18    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).19    -- <__len__>:                       return the size of dataset.20    -- <__getitem__>:                   get a data point.21    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.22    """23 24    def __init__(self, opt):25        """Initialize the class; save the options in the class26 27        Parameters:28            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions29        """30        self.opt = opt31        self.root = opt.dataroot32 33    @staticmethod34    def modify_commandline_options(parser, is_train):35        """Add new dataset-specific options, and rewrite default values for existing options.36 37        Parameters:38            parser          -- original option parser39            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.40 41        Returns:42            the modified parser.43        """44        return parser45 46    @abstractmethod47    def __len__(self):48        """Return the total number of images in the dataset."""49        return 050 51    @abstractmethod52    def __getitem__(self, index):53        """Return a data point and its metadata information.54 55        Parameters:56            index - - a random integer for data indexing57 58        Returns:59            a dictionary of data with their names. It ususally contains the data itself and its metadata information.60        """61        pass62 63 64def get_params(opt, size):65    w, h = size66    new_h = h67    new_w = w68    if opt.preprocess == "resize_and_crop":69        new_h = new_w = opt.load_size70    elif opt.preprocess == "scale_width_and_crop":71        new_w = opt.load_size72        new_h = opt.load_size * h // w73 74    x = random.randint(0, np.maximum(0, new_w - opt.crop_size))75    y = random.randint(0, np.maximum(0, new_h - opt.crop_size))76 77    flip = random.random() > 0.578 79    return {"crop_pos": (x, y), "flip": flip}80 81 82def get_transform(opt, params=None, grayscale=False, method=transforms.InterpolationMode.BICUBIC, convert=True):83    transform_list = []84    if grayscale:85        transform_list.append(transforms.Grayscale(1))86    if "resize" in opt.preprocess:87        osize = [opt.load_size, opt.load_size]88        transform_list.append(transforms.Resize(osize, method))89    elif "scale_width" in opt.preprocess:90        transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, opt.crop_size, method)))91 92    if "crop" in opt.preprocess:93        if params is None:94            transform_list.append(transforms.RandomCrop(opt.crop_size))95        else:96            transform_list.append(transforms.Lambda(lambda img: __crop(img, params["crop_pos"], opt.crop_size)))97 98    if opt.preprocess == "none":99        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base=4, method=method)))100 101    if not opt.no_flip:102        if params is None:103            transform_list.append(transforms.RandomHorizontalFlip())104        elif params["flip"]:105            transform_list.append(transforms.Lambda(lambda img: __flip(img, params["flip"])))106 107    if convert:108        transform_list += [transforms.ToTensor()]109        if grayscale:110            transform_list += [transforms.Normalize((0.5,), (0.5,))]111        else:112            transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]113    return transforms.Compose(transform_list)114 115 116def __transforms2pil_resize(method):117    mapper = {118        transforms.InterpolationMode.BILINEAR: Image.BILINEAR,119        transforms.InterpolationMode.BICUBIC: Image.BICUBIC,120        transforms.InterpolationMode.NEAREST: Image.NEAREST,121        transforms.InterpolationMode.LANCZOS: Image.LANCZOS,122    }123    return mapper[method]124 125 126def __make_power_2(img, base, method=transforms.InterpolationMode.BICUBIC):127    method = __transforms2pil_resize(method)128    ow, oh = img.size129    h = int(round(oh / base) * base)130    w = int(round(ow / base) * base)131    if h == oh and w == ow:132        return img133 134    __print_size_warning(ow, oh, w, h)135    return img.resize((w, h), method)136 137 138def __scale_width(img, target_size, crop_size, method=transforms.InterpolationMode.BICUBIC):139    method = __transforms2pil_resize(method)140    ow, oh = img.size141    if ow == target_size and oh >= crop_size:142        return img143    w = target_size144    h = int(max(target_size * oh / ow, crop_size))145    return img.resize((w, h), method)146 147 148def __crop(img, pos, size):149    ow, oh = img.size150    x1, y1 = pos151    tw = th = size152    if ow > tw or oh > th:153        return img.crop((x1, y1, x1 + tw, y1 + th))154    return img155 156 157def __flip(img, flip):158    if flip:159        return img.transpose(Image.FLIP_LEFT_RIGHT)160    return img161 162 163def __print_size_warning(ow, oh, w, h):164    """Print warning information about image size(only print once)"""165    if not hasattr(__print_size_warning, "has_printed"):166        print("The image size needs to be a multiple of 4. " "The loaded image size was (%d, %d), so it was adjusted to " "(%d, %d). This adjustment will be done to all images " "whose sizes are not multiples of 4" % (ow, oh, w, h))167        __print_size_warning.has_printed = True168