Spaces:
Running
Running
| """ | |
| Simply load images from a folder or nested folders (does not have any split). | |
| """ | |
| import logging | |
| from pathlib import Path | |
| import omegaconf | |
| import torch | |
| from ..utils.image import ImagePreprocessor, load_image | |
| from .base_dataset import BaseDataset | |
| class ImageFolder(BaseDataset, torch.utils.data.Dataset): | |
| default_conf = { | |
| "glob": ["*.jpg", "*.png", "*.jpeg", "*.JPG", "*.PNG"], | |
| "images": "???", | |
| "root_folder": "/", | |
| "preprocessing": ImagePreprocessor.default_conf, | |
| } | |
| def _init(self, conf): | |
| self.root = conf.root_folder | |
| if isinstance(conf.images, str): | |
| if not Path(conf.images).is_dir(): | |
| with open(conf.images, "r") as f: | |
| self.images = f.read().rstrip("\n").split("\n") | |
| logging.info(f"Found {len(self.images)} images in list file.") | |
| else: | |
| self.images = [] | |
| glob = [conf.glob] if isinstance(conf.glob, str) else conf.glob | |
| for g in glob: | |
| self.images += list(Path(conf.images).glob("**/" + g)) | |
| if len(self.images) == 0: | |
| raise ValueError( | |
| f"Could not find any image in folder: {conf.images}." | |
| ) | |
| self.images = [i.relative_to(conf.images) for i in self.images] | |
| self.root = conf.images | |
| logging.info(f"Found {len(self.images)} images in folder.") | |
| elif isinstance(conf.images, omegaconf.listconfig.ListConfig): | |
| self.images = conf.images.to_container() | |
| else: | |
| raise ValueError(conf.images) | |
| self.preprocessor = ImagePreprocessor(conf.preprocessing) | |
| def get_dataset(self, split): | |
| return self | |
| def __getitem__(self, idx): | |
| path = self.images[idx] | |
| img = load_image(path) | |
| data = {"name": str(path), **self.preprocessor(img)} | |
| return data | |
| def __len__(self): | |
| return len(self.images) | |