codecrafters/github-avatars

1"""Custom vs. default GitHub avatars dataset"""2 3from pathlib import Path4from typing import List5 6import datasets7from datasets.tasks import ImageClassification8 9 10logger = datasets.logging.get_logger(__name__)11 12_URL = "https://huggingface.co/datasets/codecrafters/github-avatars/resolve/main/custom-and-default-avatars.zip"13 14_HOMEPAGE = "https://codecrafters.io"15 16_DESCRIPTION = "A dataset of custom vs. default GitHub avatars"17 18 19class CatsVsDogs(datasets.GeneratorBasedBuilder):20    def _info(self):21        return datasets.DatasetInfo(22            description=_DESCRIPTION,23            features=datasets.Features(24                {25                    "image": datasets.Image(),26                    "labels": datasets.features.ClassLabel(names=["custom", "default"]),27                }28            ),29            supervised_keys=("image", "labels"),30            task_templates=[31                ImageClassification(image_column="image", label_column="labels")32            ],33            homepage=_HOMEPAGE,34        )35 36    def _split_generators(37        self, dl_manager: datasets.DownloadManager38    ) -> List[datasets.SplitGenerator]:39        images_path = (40            Path(dl_manager.download_and_extract(_URL)) / "custom-and-default-avatars"41        )42        return [43            datasets.SplitGenerator(44                name=datasets.Split.TRAIN, gen_kwargs={"images_path": images_path}45            ),46        ]47 48    def _generate_examples(self, images_path):49        logger.info("generating examples from = %s", images_path)50        for i, filepath in enumerate(images_path.glob("**/*.png")):51            yield str(i), {52                "image": str(filepath),53                "labels": filepath.parent.name.lower(),54            }55