Class GetDataLoader

mlh/data_preprocessing/data_loader.py:36–221

class GetDataLoader(object):
    def __init__(self, args):
        self.args = args
        self.data_path = args.data_path
        self.input_shape = args.input_shape

    def parse_dataset(self, dataset, train_transform, test_transform):

        if dataset in configs.SUPPORTED_IMAGE_DATASETS:
            _loader = getattr(datasets, dataset)
            if dataset != "EMNIST":
                train_dataset = _loader(root=self.data_path,
                                        train=True,
                                        transform=train_transform,
                                        download=True)
                test_dataset = _loader(root=self.data_path,
                                       train=False,
                                       transform=test_transform,
                                       download=True)
            else:
                train_dataset = _loader(root=self.data_path,
                                        train=True,
                                        split="byclass",
                                        transform=train_transform,
                                        download=True)
                test_dataset = _loader(root=self.data_path,
                                       train=False,
                                       split="byclass",
                                       transform=test_transform,
                                       download=True)
            dataset = train_dataset + test_dataset

        elif dataset in configs.SUPPORTED_IMAGE_DATASETS_ATTRIBUTE_INFERENCE:
            from mlh.data_preprocessing.attribute_data_parser import CelebA
            _loader = CelebA
            train_dataset = _loader(root=self.data_path,
                                    train=True,
                                    transform=train_transform,
                                    download=True)
            test_dataset = _loader(root=self.data_path,
                                   train=False,
                                   transform=test_transform,
                                   download=True)
            dataset = train_dataset + test_dataset

        else:
            raise ValueError("Dataset Not Supported: ", dataset)
        return dataset

    def get_data_transform(self, dataset, use_transform="simple"):
        transform_list = [transforms.Resize(
            (self.input_shape[0], self.input_shape[0])), ]

        if use_transform == "simple":
            transform_list += [transforms.RandomCrop(
                32, padding=4), transforms.RandomHorizontalFlip(), ]

            print("add simple data augmentation!")