UniqueData/pose_estimation

1import datasets2import pandas as pd3 4_CITATION = """\5@InProceedings{huggingface:dataset,6title = {pose_estimation},7author = {TrainingDataPro},8year = {2023}9}10"""11 12_DESCRIPTION = """\13The dataset is primarly intended to dentify and predict the positions of major14joints of a human body in an image. It consists of people's photographs with15body part labeled with keypoints.16"""17_NAME = 'pose_estimation'18 19_HOMEPAGE = f"https://huggingface.co/datasets/TrainingDataPro/{_NAME}"20 21_LICENSE = "cc-by-nc-nd-4.0"22 23_DATA = f"https://huggingface.co/datasets/TrainingDataPro/{_NAME}/resolve/main/data/"24 25 26class PoseEstimation(datasets.GeneratorBasedBuilder):27 28    def _info(self):29        return datasets.DatasetInfo(description=_DESCRIPTION,30                                    features=datasets.Features({31                                        'image_id': datasets.Value('uint32'),32                                        'image': datasets.Image(),33                                        'mask': datasets.Image(),34                                        'shapes': datasets.Value('string')35                                    }),36                                    supervised_keys=None,37                                    homepage=_HOMEPAGE,38                                    citation=_CITATION,39                                    license=_LICENSE)40 41    def _split_generators(self, dl_manager):42        images = dl_manager.download(f"{_DATA}images.tar.gz")43        masks = dl_manager.download(f"{_DATA}masks.tar.gz")44        annotations = dl_manager.download(f"{_DATA}{_NAME}.csv")45        images = dl_manager.iter_archive(images)46        masks = dl_manager.iter_archive(masks)47 48        return [49            datasets.SplitGenerator(name=datasets.Split.TRAIN,50                                    gen_kwargs={51                                        "images": images,52                                        "masks": masks,53                                        'annotations': annotations54                                    }),55        ]56 57    def _generate_examples(self, images, masks, annotations):58        annotations_df = pd.read_csv(annotations, sep=',')59        for idx, ((image_path, image),60                  (mask_path, mask)) in enumerate(zip(images, masks)):61            file_name = int(image_path.split('.')[0].split('/')[-1])62            yield idx, {63                'image_id':64                    annotations_df.loc[annotations_df['image_id'] == file_name]65                    ['image_id'].values[0],66                "image": {67                    "path": image_path,68                    "bytes": image.read()69                },70                "mask": {71                    "path": mask_path,72                    "bytes": mask.read()73                },74                'shapes':75                    annotations_df.loc[annotations_df['image_id'] == file_name]76                    ['shapes'].values[0],77            }78