File size: 974 Bytes
4c62147 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import os
import io
from PIL import Image
from torch.utils.data import Dataset
def add_parser_arguments(parser):
pass
def task_dataset():
return TaskDataset
class TaskDataset(Dataset):
def __init__(self, args=None, is_train=True):
super(TaskDataset, self).__init__()
self.args = args
self.is_train = is_train
self.root_dir = None
self.sample_list = []
self.idxs = []
self.im_loader = ImageLoader()
if is_train:
self.root_dir = list(self.args.trainset.values())[0]
else:
self.root_dir = list(self.args.valset.values())[0]
def __len__(self):
return len(self.sample_list)
def __getitem__(self, idx):
raise NotImplementedError
class ImageLoader:
def __init__(self):
pass
def load(self, name):
image = Image.open(name)
return image
|