src/python-doctr/doctr/datasets/funsd.py

# Copyright (C) 2021-2025, Mindee.

# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.

import json
import os
from pathlib import Path
from typing import Any

import numpy as np
from tqdm import tqdm

from .datasets import VisionDataset
from .utils import convert_target_to_relative, crop_bboxes_from_image

__all__ = ["FUNSD"]


class FUNSD(VisionDataset):
    """FUNSD dataset from `"FUNSD: A Dataset for Form Understanding in Noisy Scanned Documents"
    <https://arxiv.org/pdf/1905.13538.pdf>`_.

    .. image:: https://doctr-static.mindee.com/models?id=v0.5.0/funsd-grid.png&src=0
        :align: center

    >>> from doctr.datasets import FUNSD
    >>> train_set = FUNSD(train=True, download=True)
    >>> img, target = train_set[0]

    Args:
        train: whether the subset should be the training one
        use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)
        recognition_task: whether the dataset should be used for recognition task
        detection_task: whether the dataset should be used for detection task
        **kwargs: keyword arguments from `VisionDataset`.
    """

    URL = "https://guillaumejaume.github.io/FUNSD/dataset.zip"
    SHA256 = "c31735649e4f441bcbb4fd0f379574f7520b42286e80b01d80b445649d54761f"
    FILE_NAME = "funsd.zip"

    def __init__(
        self,
        train: bool = True,
        use_polygons: bool = False,
        recognition_task: bool = False,
        detection_task: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            self.URL,
            self.FILE_NAME,
            self.SHA256,
            True,
            pre_transforms=convert_target_to_relative if not recognition_task else None,
            **kwargs,
        )
        if recognition_task and detection_task:
            raise ValueError(
                "`recognition_task` and `detection_task` cannot be set to True simultaneously. "
                + "To get the whole dataset with boxes and labels leave both parameters to False."
            )

        self.train = train
        np_dtype = np.float32

        # Use the subset
        subfolder = os.path.join("dataset", "training_data" if train else "testing_data")

        # # list images
        tmp_root = os.path.join(self.root, subfolder, "images")
        self.data: list[tuple[str | np.ndarray, str | dict[str, Any] | np.ndarray]] = []
        for img_path in tqdm(
            iterable=os.listdir(tmp_root), desc="Preparing and Loading FUNSD", total=len(os.listdir(tmp_root))
        ):
            # File existence check
            if not os.path.exists(os.path.join(tmp_root, img_path)):
                raise FileNotFoundError(f"unable to locate {os.path.join(tmp_root, img_path)}")

            stem = Path(img_path).stem
            with open(os.path.join(self.root, subfolder, "annotations", f"{stem}.json"), "rb") as f:
                data = json.load(f)

            _targets = [
                (word["text"], word["box"])
                for block in data["form"]
                for word in block["words"]
                if len(word["text"]) > 0
            ]
            text_targets, box_targets = zip(*_targets)
            if use_polygons:
                # xmin, ymin, xmax, ymax -> (x, y) coordinates of top left, top right, bottom right, bottom left corners
                box_targets = [  # type: ignore[assignment]
                    [
                        [box[0], box[1]],
                        [box[2], box[1]],
                        [box[2], box[3]],
                        [box[0], box[3]],
                    ]
                    for box in box_targets
                ]

            if recognition_task:
                crops = crop_bboxes_from_image(
                    img_path=os.path.join(tmp_root, img_path), geoms=np.asarray(box_targets, dtype=np_dtype)
                )
                for crop, label in zip(crops, list(text_targets)):
                    # filter labels with unknown characters
                    if not any(char in label for char in ["☑", "☐", "\u03bf", "\uf703", "\uf702", " "]):
                        self.data.append((crop, label.replace("–", "-")))
            elif detection_task:
                self.data.append((img_path, np.asarray(box_targets, dtype=np_dtype)))
            else:
                self.data.append((
                    img_path,
                    dict(boxes=np.asarray(box_targets, dtype=np_dtype), labels=list(text_targets)),
                ))

        self.root = tmp_root

    def extra_repr(self) -> str:
        return f"train={self.train}"