docling-project/DocLayNet

1"""2Inspired from3https://huggingface.co/datasets/ydshieh/coco_dataset_script/blob/main/coco_dataset_script.py4"""5 6import json7import os8import datasets9import collections10 11 12class COCOBuilderConfig(datasets.BuilderConfig):13    def __init__(self, name, splits, **kwargs):14        super().__init__(name, **kwargs)15        self.splits = splits16 17 18# Add BibTeX citation19# Find for instance the citation on arxiv or on the dataset repo/website20_CITATION = """\21@article{doclaynet2022,22  title = {DocLayNet: A Large Human-Annotated Dataset for Document-Layout Analysis},  23  doi = {10.1145/3534678.353904},24  url = {https://arxiv.org/abs/2206.01062},25  author = {Pfitzmann, Birgit and Auer, Christoph and Dolfi, Michele and Nassar, Ahmed S and Staar, Peter W J},26  year = {2022}27}28"""29 30# Add description of the dataset here31# You can copy an official description32_DESCRIPTION = """\33DocLayNet is a human-annotated document layout segmentation dataset from a broad variety of document sources.34"""35 36# Add a link to an official homepage for the dataset here37_HOMEPAGE = "https://developer.ibm.com/exchanges/data/all/doclaynet/"38 39# Add the licence for the dataset here if you can find it40_LICENSE = "CDLA-Permissive-1.0"41 42# Add link to the official dataset URLs here43# The HuggingFace dataset library don't host the datasets but only point to the original files44# This can be an arbitrary nested dict/list of URLs (see below in `_split_generators` method)45 46_URLs = {47    "core": "https://codait-cos-dax.s3.us.cloud-object-storage.appdomain.cloud/dax-doclaynet/1.0.0/DocLayNet_core.zip",48}49 50# Name of the dataset usually match the script name with CamelCase instead of snake_case51class COCODataset(datasets.GeneratorBasedBuilder):52    """An example dataset script to work with the local (downloaded) COCO dataset"""53 54    VERSION = datasets.Version("1.0.0")55 56    BUILDER_CONFIG_CLASS = COCOBuilderConfig57    BUILDER_CONFIGS = [58        COCOBuilderConfig(name="2022.08", splits=["train", "val", "test"]),59    ]60    DEFAULT_CONFIG_NAME = "2022.08"61 62    def _info(self):63        features = datasets.Features(64            {65                "image_id": datasets.Value("int64"),66                "image": datasets.Image(),67                "width": datasets.Value("int32"),68                "height": datasets.Value("int32"),69                # Custom fields70                "doc_category": datasets.Value(71                    "string"72                ),  # high-level document category73                "collection": datasets.Value("string"),  # sub-collection name74                "doc_name": datasets.Value("string"),  # original document filename75                "page_no": datasets.Value("int64"),  # page number in original document76            }77        )78        object_dict = {79            "category_id": datasets.ClassLabel(80                names=[81                    "Caption",82                    "Footnote",83                    "Formula",84                    "List-item",85                    "Page-footer",86                    "Page-header",87                    "Picture",88                    "Section-header",89                    "Table",90                    "Text",91                    "Title",92                ]93            ),94            "image_id": datasets.Value("string"),95            "id": datasets.Value("int64"),96            "area": datasets.Value("int64"),97            "bbox": datasets.Sequence(datasets.Value("float32"), length=4),98            "segmentation": [[datasets.Value("float32")]],99            "iscrowd": datasets.Value("bool"),100            "precedence": datasets.Value("int32"),101        }102        features["objects"] = [object_dict]103 104        return datasets.DatasetInfo(105            # This is the description that will appear on the datasets page.106            description=_DESCRIPTION,107            # This defines the different columns of the dataset and their types108            features=features,  # Here we define them above because they are different between the two configurations109            # If there's a common (input, target) tuple from the features,110            # specify them here. They'll be used if as_supervised=True in111            # builder.as_dataset.112            supervised_keys=None,113            # Homepage of the dataset for documentation114            homepage=_HOMEPAGE,115            # License for the dataset if available116            license=_LICENSE,117            # Citation for the dataset118            citation=_CITATION,119        )120 121    def _split_generators(self, dl_manager):122        """Returns SplitGenerators."""123        archive_path = dl_manager.download_and_extract(_URLs)124        splits = []125        for split in self.config.splits:126            if split == "train":127                dataset = datasets.SplitGenerator(128                    name=datasets.Split.TRAIN,129                    # These kwargs will be passed to _generate_examples130                    gen_kwargs={131                        "json_path": os.path.join(132                            archive_path["core"], "COCO", "train.json"133                        ),134                        "image_dir": os.path.join(archive_path["core"], "PNG"),135                        "split": "train",136                    },137                )138            elif split in ["val", "valid", "validation", "dev"]:139                dataset = datasets.SplitGenerator(140                    name=datasets.Split.VALIDATION,141                    # These kwargs will be passed to _generate_examples142                    gen_kwargs={143                        "json_path": os.path.join(144                            archive_path["core"], "COCO", "val.json"145                        ),146                        "image_dir": os.path.join(archive_path["core"], "PNG"),147                        "split": "val",148                    },149                )150            elif split == "test":151                dataset = datasets.SplitGenerator(152                    name=datasets.Split.TEST,153                    # These kwargs will be passed to _generate_examples154                    gen_kwargs={155                        "json_path": os.path.join(156                            archive_path["core"], "COCO", "test.json"157                        ),158                        "image_dir": os.path.join(archive_path["core"], "PNG"),159                        "split": "test",160                    },161                )162            else:163                continue164 165            splits.append(dataset)166        return splits167 168    def _generate_examples(169        # method parameters are unpacked from `gen_kwargs` as given in `_split_generators`170        self,171        json_path,172        image_dir,173        split,174    ):175        """Yields examples as (key, example) tuples."""176        # This method handles input defined in _split_generators to yield (key, example) tuples from the dataset.177        # The `key` is here for legacy reason (tfds) and is not important in itself.178        def _image_info_to_example(image_info, image_dir):179            image = image_info["file_name"]180            return {181                "image_id": image_info["id"],182                "image": os.path.join(image_dir, image),183                "width": image_info["width"],184                "height": image_info["height"],185                "doc_category": image_info["doc_category"],186                "collection": image_info["collection"],187                "doc_name": image_info["doc_name"],188                "page_no": image_info["page_no"],189            }190 191        with open(json_path, encoding="utf8") as f:192            annotation_data = json.load(f)193            images = annotation_data["images"]194            annotations = annotation_data["annotations"]195            image_id_to_annotations = collections.defaultdict(list)196            for annotation in annotations:197                image_id_to_annotations[annotation["image_id"]].append(annotation)198 199        for idx, image_info in enumerate(images):200            example = _image_info_to_example(image_info, image_dir)201            annotations = image_id_to_annotations[image_info["id"]]202            objects = []203            for annotation in annotations:204                category_id = annotation["category_id"]  # Zero based counting205                if category_id != -1:206                    category_id = category_id - 1207                annotation["category_id"] = category_id208                objects.append(annotation)209            example["objects"] = objects210            yield idx, example211