Spaces:
Running
on
A10G
Running
on
A10G
| """This module implements an abstract base class (ABC) 'BaseDataset' for datasets. | |
| It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses. | |
| """ | |
| import random | |
| import numpy as np | |
| import torch.utils.data as data | |
| from PIL import Image | |
| import torchvision.transforms as transforms | |
| from abc import ABC, abstractmethod | |
| class BaseDataset(data.Dataset, ABC): | |
| """This class is an abstract base class (ABC) for datasets. | |
| To create a subclass, you need to implement the following four functions: | |
| -- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt). | |
| -- <__len__>: return the size of dataset. | |
| -- <__getitem__>: get a data point. | |
| -- <modify_commandline_options>: (optionally) add dataset-specific options and set default options. | |
| """ | |
| def __init__(self, opt): | |
| """Initialize the class; save the options in the class | |
| Parameters: | |
| opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions | |
| """ | |
| self.opt = opt | |
| # self.root = opt.dataroot | |
| self.current_epoch = 0 | |
| def modify_commandline_options(parser, is_train): | |
| """Add new dataset-specific options, and rewrite default values for existing options. | |
| Parameters: | |
| parser -- original option parser | |
| is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options. | |
| Returns: | |
| the modified parser. | |
| """ | |
| return parser | |
| def __len__(self): | |
| """Return the total number of images in the dataset.""" | |
| return 0 | |
| def __getitem__(self, index): | |
| """Return a data point and its metadata information. | |
| Parameters: | |
| index - - a random integer for data indexing | |
| Returns: | |
| a dictionary of data with their names. It ususally contains the data itself and its metadata information. | |
| """ | |
| pass | |
| def get_transform(grayscale=False): | |
| transform_list = [] | |
| if grayscale: | |
| transform_list.append(transforms.Grayscale(1)) | |
| transform_list += [transforms.ToTensor()] | |
| return transforms.Compose(transform_list) | |
| def get_affine_mat(opt, size): | |
| shift_x, shift_y, scale, rot_angle, flip = 0., 0., 1., 0., False | |
| w, h = size | |
| if 'shift' in opt.preprocess: | |
| shift_pixs = int(opt.shift_pixs) | |
| shift_x = random.randint(-shift_pixs, shift_pixs) | |
| shift_y = random.randint(-shift_pixs, shift_pixs) | |
| if 'scale' in opt.preprocess: | |
| scale = 1 + opt.scale_delta * (2 * random.random() - 1) | |
| if 'rot' in opt.preprocess: | |
| rot_angle = opt.rot_angle * (2 * random.random() - 1) | |
| rot_rad = -rot_angle * np.pi/180 | |
| if 'flip' in opt.preprocess: | |
| flip = random.random() > 0.5 | |
| shift_to_origin = np.array([1, 0, -w//2, 0, 1, -h//2, 0, 0, 1]).reshape([3, 3]) | |
| flip_mat = np.array([-1 if flip else 1, 0, 0, 0, 1, 0, 0, 0, 1]).reshape([3, 3]) | |
| shift_mat = np.array([1, 0, shift_x, 0, 1, shift_y, 0, 0, 1]).reshape([3, 3]) | |
| rot_mat = np.array([np.cos(rot_rad), np.sin(rot_rad), 0, -np.sin(rot_rad), np.cos(rot_rad), 0, 0, 0, 1]).reshape([3, 3]) | |
| scale_mat = np.array([scale, 0, 0, 0, scale, 0, 0, 0, 1]).reshape([3, 3]) | |
| shift_to_center = np.array([1, 0, w//2, 0, 1, h//2, 0, 0, 1]).reshape([3, 3]) | |
| affine = shift_to_center @ scale_mat @ rot_mat @ shift_mat @ flip_mat @ shift_to_origin | |
| affine_inv = np.linalg.inv(affine) | |
| return affine, affine_inv, flip | |
| def apply_img_affine(img, affine_inv, method=Image.BICUBIC): | |
| return img.transform(img.size, Image.AFFINE, data=affine_inv.flatten()[:6], resample=Image.BICUBIC) | |
| def apply_lm_affine(landmark, affine, flip, size): | |
| _, h = size | |
| lm = landmark.copy() | |
| lm[:, 1] = h - 1 - lm[:, 1] | |
| lm = np.concatenate((lm, np.ones([lm.shape[0], 1])), -1) | |
| lm = lm @ np.transpose(affine) | |
| lm[:, :2] = lm[:, :2] / lm[:, 2:] | |
| lm = lm[:, :2] | |
| lm[:, 1] = h - 1 - lm[:, 1] | |
| if flip: | |
| lm_ = lm.copy() | |
| lm_[:17] = lm[16::-1] | |
| lm_[17:22] = lm[26:21:-1] | |
| lm_[22:27] = lm[21:16:-1] | |
| lm_[31:36] = lm[35:30:-1] | |
| lm_[36:40] = lm[45:41:-1] | |
| lm_[40:42] = lm[47:45:-1] | |
| lm_[42:46] = lm[39:35:-1] | |
| lm_[46:48] = lm[41:39:-1] | |
| lm_[48:55] = lm[54:47:-1] | |
| lm_[55:60] = lm[59:54:-1] | |
| lm_[60:65] = lm[64:59:-1] | |
| lm_[65:68] = lm[67:64:-1] | |
| lm = lm_ | |
| return lm | |