naver/PUMP

1# Copyright 2022-present NAVER Corp.2# CC BY-NC-SA 4.03# Available only for non-commercial use4 5from pdb import set_trace as bb6import os7from os.path import *8from PIL import Image9 10 11class ImageSet(object):12    """ Base class for an image dataset.13    """14    def __init__(self, root, imgs):15        self.root = root16        self.imgs = imgs17        assert imgs, f'Empty image set in {root}'18 19    def init_from_folder(self, *args, **kw):20        imset = ImageSet.from_folder(*args, **kw)21        ImageSet.__init__(self, imset.root, imset.imgs)22 23    def __len__(self):24        return len(self.imgs)25 26    def get_image_path(self, idx):27        return os.path.join(self.root, self.imgs[idx])28 29    def get_image(self, idx):30        fname = self.get_image_path(idx)31        try:32            return Image.open(fname).convert('RGB')33        except Exception as e:34            raise IOError("Could not load image %s (reason: %s)" % (fname, str(e)))35 36    __getitem__ = get_image37 38    @staticmethod39    def from_folder(root, exts=('.jpg','.jpeg','.png','.ppm'), recursive=False, listing=False, check_imgs=False):40        """41        recursive: bool or func. If a function, it must evaluate True to the directory name.42        """43        if listing: 44            if listing is True: listing = f"list_imgs{'_recursive' if recursive else ''}.txt"45            flist = join(root, listing)46            try: return ImageSet.from_listing(root,flist)47            except IOError: print(f'>> ImageSet.from_folder(listing=True): entering {root}...')48 49        if check_imgs is True: # default verif function50            check_imgs = verify_img51 52        for _, dirnames, dirfiles in os.walk(root):53            imgs = sorted([f for f in dirfiles if f.lower().endswith(exts)])54            if check_imgs: imgs = [img for img in imgs if check_imgs(join(root,img))]55 56            if recursive:57                for dirname in sorted(dirnames):58                    if callable(recursive) and not recursive(join(root,dirname)): continue59                    imset = ImageSet.from_folder(join(root,dirname), exts=exts, recursive=recursive, listing=listing, check_imgs=check_imgs)60                    imgs += [join(dirname,f) for f in imset.imgs]61            break # recursion is handled internally62 63        if listing: 64            try: open(flist,'w').write('\n'.join(imgs))65            except IOError: pass # write permission denied66        return ImageSet(root, imgs)67 68    @staticmethod69    def from_listing(root, list_path):70        return ImageSet(root, open(list_path).read().splitlines())71 72    def circular_pad(self, min_size):73        assert self.imgs, 'cannot pad an empty image set'74        while len(self.imgs) < min_size: 75            self.imgs += self.imgs # artifically augment size76        self.imgs = self.imgs[:min_size or None]77        return self78 79    def __repr__(self):80        prefix = os.path.commonprefix((self.get_image_path(0),self.get_image_path(len(self)-1)))81        return f'{self.__class__.__name__}({len(self)} images from {prefix}...)'82 83 84 85def verify_img(path, exts=None):86    if exts and not path.lower().endswith(exts): return False87    try: 88        Image.open(path).convert('RGB') # try to open it89        return True90    except: 91        return False92