1aurent/Human-Embryo-Timelapse

1import datasets2import pandas as pd3from pathlib import Path4from PIL import ImageFile5 6ImageFile.LOAD_TRUNCATED_IMAGES = True7 8_URLS = {9    "F-45": "https://zenodo.org/records/7912264/files/embryo_dataset_F-45.tar.gz?download=1",10    "F-30": "https://zenodo.org/records/7912264/files/embryo_dataset_F-30.tar.gz?download=1",11    "F-15": "https://zenodo.org/records/7912264/files/embryo_dataset_F-15.tar.gz?download=1",12    "F0": "https://zenodo.org/records/7912264/files/embryo_dataset.tar.gz?download=1",13    "F+15": "https://zenodo.org/records/7912264/files/embryo_dataset_F15.tar.gz?download=1",14    "F+30": "https://zenodo.org/records/7912264/files/embryo_dataset_F30.tar.gz?download=1",15    "F+45": "https://zenodo.org/records/7912264/files/embryo_dataset_F45.tar.gz?download=1",16    "grades": "https://zenodo.org/records/7912264/files/embryo_dataset_grades.csv?download=1",17    "annotations": "https://zenodo.org/records/7912264/files/embryo_dataset_annotations.tar.gz?download=1",18    "time_elapsed": "https://zenodo.org/records/7912264/files/embryo_dataset_time_elapsed.tar.gz?download=1",19}20 21_EVENT_NAMES = [22    "tPB2", "tPNa", "tPNf", "t2", "t3", "t4", "t5", "t6", "t7", "t8", "t9+", "tM", "tSB", "tB", "tEB", "tHB",23]24 25_GRADES = ["A", "B", "C", "NA"]26 27_DESCRIPTION = """28This dataset is composed of 704 videos, each recorded at 7 focal planes, accompanied by the annotations of 16 cellular events.29"""30 31_VERSION = datasets.Version("0.3.0")32 33_HOMEPAGE = "https://zenodo.org/record/7912264"34 35_LICENSE = "CC BY-NC-SA 4.0"36 37class HumanEmbryoTimelapse(datasets.GeneratorBasedBuilder):38 39    def _info(self):40        return datasets.DatasetInfo(41            description=_DESCRIPTION,42            version=_VERSION,43            homepage=_HOMEPAGE,44            license=_LICENSE,45            features=datasets.Features(46                {47                    "name": datasets.Value("string"),48                    "F-45": datasets.Sequence(datasets.Image()),49                    "F-30": datasets.Sequence(datasets.Image()),50                    "F-15": datasets.Sequence(datasets.Image()),51                    "F0": datasets.Sequence(datasets.Image()),52                    "F+45": datasets.Sequence(datasets.Image()),53                    "F+30": datasets.Sequence(datasets.Image()),54                    "F+15": datasets.Sequence(datasets.Image()),55                    "events": datasets.Sequence(56                        {57                            "name": datasets.ClassLabel(names=_EVENT_NAMES),58                            "frame_index_start": datasets.Value("uint16"),59                            "frame_index_stop": datasets.Value("uint16"),60                        },61                    ),62                    "timeline": {63                        "frame_index": datasets.Sequence(datasets.Value("uint16")),64                        "time": datasets.Sequence(datasets.Value("float32")),65                    },66                    "grades": {67                        "TE": datasets.ClassLabel(names=_GRADES),68                        "ICM": datasets.ClassLabel(names=_GRADES),69                    }70                }71            ),72        )73 74    def _split_generators(self, dl_manager):75        """Generate splits."""76 77        # download and extract all files78        directories = {79            name: Path(dl_manager.download_and_extract(url))80            for name, url in _URLS.items()81        }82 83        # get all subfolders of embryo_names_dir84        embryo_names_dir = directories["F0"] / "embryo_dataset"85        embryo_names = [x.name for x in embryo_names_dir.iterdir() if x.is_dir()]86 87        return [88            datasets.SplitGenerator(89                name=datasets.Split.TRAIN,90                gen_kwargs={91                    "embryo_names": embryo_names,92                    "directories": directories,93                },94            )95        ]96 97    def _generate_examples(self, embryo_names, directories):98        """Generate images and labels for splits."""99        100        # get grades for each embryo (name, TE, ICM)101        pd_grades = pd.read_csv(directories["grades"], keep_default_na=False)102        grades = {103            row["video_name"]: {104                "TE": row["TE"],105                "ICM": row["ICM"],106            }107            for _, row in pd_grades.iterrows()108        }109 110        for index, embryo_name in enumerate(embryo_names):111 112            # get events of the embryo (name, frame_index_start, frame_index_stop)113            pd_events = pd.read_csv(directories["annotations"] / "embryo_dataset_annotations" / f"{embryo_name}_phases.csv", header=None)114            events = [115                {116                    "name": row[0],117                    "frame_index_start": row[1],118                    "frame_index_stop": row[2],119                }120                for _, row in pd_events.iterrows()121            ]122 123            # get frame index and time124            pd_time = pd.read_csv(directories["time_elapsed"] / "embryo_dataset_time_elapsed" / f"{embryo_name}_timeElapsed.csv")125            timeline = {126                "frame_index": pd_time["frame_index"].tolist(),127                "time": pd_time["time"].tolist(),128            }129 130            # get images of the embryo, with focal plane -45131            F_m45 = list(map(132                lambda x: str(x),133                sorted(134                    (directories["F-45"] / "embryo_dataset_F-45" / embryo_name).glob("*.jpeg"),135                    key=lambda x: int(x.stem.split("RUN")[-1]),136                ),137            ))138 139            # get images of the embryo, with focal plane -30140            F_m30 = list(map(141                lambda x: str(x),142                sorted(143                    (directories["F-30"] / "embryo_dataset_F-30" / embryo_name).glob("*.jpeg"),144                    key=lambda x: int(x.stem.split("RUN")[-1]),145                ),146            ))147 148            # get images of the embryo, with focal plane -15149            F_m15 = list(map(150                lambda x: str(x),151                sorted(152                    (directories["F-15"] / "embryo_dataset_F-15" / embryo_name).glob("*.jpeg"),153                    key=lambda x: int(x.stem.split("RUN")[-1]),154                ),155            ))156 157            # get images of the embryo, with focal plane 0158            F_zero = list(map(159                lambda x: str(x),160                sorted(161                    (directories["F0"] / "embryo_dataset" / embryo_name).glob("*.jpeg"),162                    key=lambda x: int(x.stem.split("RUN")[-1]),163                ),164            ))165 166            # get images of the embryo, with focal plane +15167            F_p15 = list(map(168                lambda x: str(x),169                sorted(170                    (directories["F+15"] / "embryo_dataset_F15" / embryo_name).glob("*.jpeg"),171                    key=lambda x: int(x.stem.split("RUN")[-1]),172                ),173            ))174 175            # get images of the embryo, with focal plane +30176            F_p30 = list(map(177                lambda x: str(x),178                sorted(179                    (directories["F+30"] / "embryo_dataset_F30" / embryo_name).glob("*.jpeg"),180                    key=lambda x: int(x.stem.split("RUN")[-1]),181                ),182            ))183 184            # get images of the embryo, with focal plane +45185            F_p45 = list(map(186                lambda x: str(x),187                sorted(188                    (directories["F+45"] / "embryo_dataset_F45" / embryo_name).glob("*.jpeg"),189                    key=lambda x: int(x.stem.split("RUN")[-1]),190                ),191            ))192 193            yield index, {194                "name": embryo_name,195                "F-45": F_m45,196                "F-30": F_m30,197                "F-15": F_m15,198                "F0": F_zero,199                "F+15": F_p15,200                "F+30": F_p30,201                "F+45": F_p45,202                "events": events,203                "grades": grades[embryo_name],204                "timeline": timeline,205            }206