Class ImageFolderDataset

datautils.py:185–196

class ImageFolderDataset(data.Dataset):
    def __init__(self, dir, max_size=None, random_seed=0):
        self.image_paths = sorted(glob(os.path.join(self.sequence_dir, "sample_*.png")) + glob(os.path.join(self.sequence_dir, "sample_*.jpg")), key=lambda p: int(re.search(r'\d+', os.path.basename(p)).group()))
        if max_size is not None:
            random.seed(random_seed)  # Set the seed for reproducibility
            self.image_paths = sorted(random.sample(self.image_paths, max_size), key=lambda p: int(re.search(r'\d+', os.path.basename(p)).group()))

    def __len__(self):
        return len(self.image_path)

    def __getitem__(self, index):
        return (pil_to_tensor(PIL.Image.open(self.image_path[index].replace("sample", prefix))) for prefix in ["src", "tgt", "sample"])