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