# Copyright 2019-present NAVER Corp. # CC BY-NC-SA 3.0 # Available only for non-commercial use import os, pdb from .dataset import Dataset from .pair_dataset import SyntheticPairDataset class ImgFolder (Dataset): """ load all images in a folder (no recursion). """ def __init__(self, root, imgs=None, exts=('.jpg','.png','.ppm')): Dataset.__init__(self) self.root = root self.imgs = imgs or [f for f in os.listdir(root) if f.endswith(exts)] self.nimg = len(self.imgs) def get_key(self, idx): return self.imgs[idx]