import torchvision.transforms as transforms
from antgine.datasets import ImageFolder

_default_train_transform = transforms.Compose([
    transforms.Resize(112),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[127.5/255.] * 3, std=[128.0/255.] * 3),
])

_default_test_transform = transforms.Compose([
    transforms.Resize(112),
    transforms.ToTensor(),
    transforms.Normalize(mean=[127.5/255.] * 3, std=[128.0/255.] * 3),
])



class FacesEmore(ImageFolder):
    """
        FacesEmore dataset class.
    """
    def __init__(self, root: str, batch_size: int,
                 train_transform: transforms.Compose = _default_train_transform,
                 test_transform: transforms.Compose = _default_test_transform,
                 num_workers=8):
        """
        :param str root: Dataset's root directory.
        :param int batch_size: Batch size.
        :param transforms.Compose train_transform: Transform applied to inputs during training.
        :param transforms.Compose test_transform: Transform applied to inputs during testing.
        :param int num_workers: Number of workers launched for loading data.
        """
        super().__init__(root=root, train_path='train', test_path='validation',
                         batch_size=batch_size, train_transform=train_transform,
                         test_transform=test_transform, num_workers=num_workers)
