|
|
|
|
|
|
|
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 |
|
|
|
|
