bald_classification.py

import datasets
import pandas as pd

_CITATION = """\
@InProceedings{huggingface:dataset,
title = {bald_classification},
author = {TrainingDataPro},
year = {2023}
}
"""

_DESCRIPTION = """\
Dataset consists of 5000 photos of people with 7 stages of hairloss according
to the Norwood scale. Dataset is useful for training neural networks for the
recommendation systems, optimizing the work processes of trichologists and
applications in the Med / Beauty spheres. 
"""
_NAME = 'bald_classification'

_HOMEPAGE = f"https://huggingface.co/datasets/TrainingDataPro/{_NAME}"

_LICENSE = ""

_DATA = f"https://huggingface.co/datasets/TrainingDataPro/{_NAME}/resolve/main/data/"


class BaldClassification(datasets.GeneratorBasedBuilder):
    """Small sample of image-text pairs"""

    def _info(self):
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=datasets.Features({
                'image_id': datasets.Value('int32'),
                'image': datasets.Image(),
                'annotations': datasets.Value('string')
            }),
            supervised_keys=None,
            homepage=_HOMEPAGE,
            citation=_CITATION,
        )

    def _split_generators(self, dl_manager):
        images = dl_manager.download(f"{_DATA}images.tar.gz")
        annotations = dl_manager.download(f"{_DATA}{_NAME}.csv")
        images = dl_manager.iter_archive(images)
        return [
            datasets.SplitGenerator(name=datasets.Split.TRAIN,
                                    gen_kwargs={
                                        "images": images,
                                        'annotations': annotations
                                    }),
        ]

    def _generate_examples(self, images, annotations):
        annotations_df = pd.read_csv(annotations)

        for idx, (image_path, image) in enumerate(images):
            yield idx, {
                'image_id':
                    annotations_df.loc[
                        annotations_df['image_name'] == image_path]
                    ['image_id'].values[0],
                "image": {
                    "path": image_path,
                    "bytes": image.read()
                },
                'annotations':
                    annotations_df.loc[
                        annotations_df['image_name'] == image_path]
                    ['annotations'].values[0]
            }