DL-Art-School/codes/data/torch_dataset.py

63 lines
2.0 KiB
Python
Raw Normal View History

2020-12-03 22:32:21 +00:00
import torch
from torch.utils.data import Dataset
import torchvision.transforms as T
from torchvision import datasets
# Wrapper for basic pytorch datasets which re-wraps them into a format usable by ExtensibleTrainer.
from utils.util import opt_get
2020-12-03 22:32:21 +00:00
class TorchDataset(Dataset):
def __init__(self, opt):
DATASET_MAP = {
"mnist": datasets.MNIST,
"fmnist": datasets.FashionMNIST,
"cifar10": datasets.CIFAR10,
2021-06-04 23:29:07 +00:00
"cifar100": datasets.CIFAR100,
2021-01-12 03:09:16 +00:00
"imagenet": datasets.ImageNet,
"imagefolder": datasets.ImageFolder
2020-12-03 22:32:21 +00:00
}
2021-01-12 03:09:16 +00:00
normalize = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
if opt_get(opt, ['random_crop'], False):
2021-01-12 03:09:16 +00:00
transforms = [
T.RandomResizedCrop(opt['image_size']),
T.RandomHorizontalFlip(),
T.ToTensor(),
normalize,
]
else:
transforms = [
T.Resize(opt['image_size']),
2021-01-12 03:09:16 +00:00
T.CenterCrop(opt['image_size']),
T.RandomHorizontalFlip(),
2021-01-12 03:09:16 +00:00
T.ToTensor(),
normalize,
]
2020-12-03 22:32:21 +00:00
transforms = T.Compose(transforms)
2021-01-12 03:09:16 +00:00
self.dataset = DATASET_MAP[opt['dataset']](transform=transforms, **opt['kwargs'])
2021-06-05 03:21:04 +00:00
self.len = opt_get(opt, ['fixed_len'], len(self.dataset))
self.offset = opt_get(opt, ['offset'], 0)
2020-12-03 22:32:21 +00:00
def __getitem__(self, item):
2021-06-05 03:21:04 +00:00
underlying_item, lbl = self.dataset[item+self.offset]
2021-01-12 03:09:16 +00:00
return {'lq': underlying_item, 'hq': underlying_item, 'labels': lbl,
2020-12-03 22:32:21 +00:00
'LQ_path': str(item), 'GT_path': str(item)}
def __len__(self):
2021-06-05 03:21:04 +00:00
return self.len-self.offset
2021-01-12 03:09:16 +00:00
if __name__ == '__main__':
opt = {
'flip': True,
'crop_sz': None,
2021-06-04 23:29:07 +00:00
'dataset': 'cifar100',
'image_size': 32,
2021-01-12 03:09:16 +00:00
'normalize': True,
'kwargs': {
2021-06-04 23:29:07 +00:00
'root': 'E:\\4k6k\\datasets\\images\\cifar100',
'download': True
2021-01-12 03:09:16 +00:00
}
}
set = TorchDataset(opt)
j = set[0]