File size: 571 Bytes
404d2af
 
 
 
 
 
 
 
 
 
8b973ee
 
 
 
404d2af
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# 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]