|
"""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.current_epoch = 0
|
|
|
|
@staticmethod
|
|
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
|
|
|
|
@abstractmethod
|
|
def __len__(self):
|
|
"""Return the total number of images in the dataset."""
|
|
return 0
|
|
|
|
@abstractmethod
|
|
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
|
|
|
