| |
|
|
|
|
| import os |
| from torch.utils.data import Dataset |
| from PIL import Image |
|
|
| from utils.data_utils import make_dataset |
|
|
|
|
| class ImagesDataset(Dataset): |
|
|
| def __init__(self, source_root, source_transform=None): |
| self.source_paths = sorted(make_dataset(source_root)) |
| self.source_transform = source_transform |
|
|
| def __len__(self): |
| return len(self.source_paths) |
|
|
| def __getitem__(self, index): |
| fname, from_path = self.source_paths[index] |
| from_im = Image.open(from_path).convert('RGB') |
| |
| if self.source_transform: |
| from_im = self.source_transform(from_im) |
|
|
| return fname, from_im |
|
|
|
|
