CristoJV/image-folder

1from dataclasses import dataclass2from pathlib import Path3from typing import Optional4 5import datasets6import pyarrow as pa7from datasets.features import ClassLabel, Image8from datasets.tasks import ImageClassification9 10logger = datasets.utils.logging.get_logger(__name__)11 12 13@dataclass14class ImageFolderConfig(datasets.BuilderConfig):15    """BuilderConfig for ImageFolder."""16 17    features: Optional[datasets.Features] = None18 19    @property20    def schema(self):21        return (22            pa.schema(self.features.type)23            if self.features is not None24            else None25        )26 27 28class ImageFolder(datasets.GeneratorBasedBuilder):29 30    BUILDER_CONFIG_CLASS = ImageFolderConfig31 32    def _info(self):33        filepaths = None34        if isinstance(self.config.data_files, str):35            filepaths = self.config.data_files36        elif isinstance(self.config.data_files, dict):37            filepaths = self.config.data_files.get("train", None)38        if filepaths is None:39            raise RuntimeError("data_files must be specified")40 41        classes = sorted(42            [Path(file_path).parent.name.lower() for file_path in filepaths]43        )44 45        # Remove duplicates46        classes = list(set(classes))47 48        return datasets.DatasetInfo(49            features=datasets.Features(50                {51                    "image_filepath": Image(),52                    "labels": ClassLabel(names=classes),53                }54            ),55            task_templates=[56                ImageClassification(57                    image_column="image_filepath",58                    label_column="labels",59                )60            ],61        )62 63    def _split_generators(self, dl_manager):64        if not self.config.data_files:65            raise ValueError(66                f"At least one data file must be specified, but got data_files={self.config.data_files}"67            )68 69        data_files = self.config.data_files70        if isinstance(data_files, str):71            file_paths = data_files72            return [73                datasets.SplitGenerator(74                    name=datasets.Split.TRAIN,75                    gen_kwargs={"file_paths": file_paths},76                )77            ]78        splits = []79        for split_name, file_paths in data_files.items():80            splits.append(81                datasets.SplitGenerator(82                    name=split_name, gen_kwargs={"file_paths": file_paths}83                )84            )85        return splits86 87    def _generate_examples(self, file_paths):88        logger.info("generating examples from = %s", file_paths)89        extensions = {90            ".jpg",91            ".jpeg",92            ".png",93            ".ppm",94            ".bmp",95            ".pgm",96            ".tif",97            ".tiff",98            ".webp",99        }100        for i, path in enumerate(file_paths):101            path = Path(path)102            if path.suffix in extensions:103                yield i, {104                    "image_filepath": path.as_posix(),105                    "labels": path.parent.name.lower(),106                }107