[Feat] Add layout detection dataset class (#2064)

Author: felixdittrich92

docs/source/modules/datasets.rst

@@ -50,6 +50,8 @@ Custom dataset loader
 
 .. autoclass:: DetectionDataset
 
+.. autoclass:: LayoutDataset
+
 .. autoclass:: RecognitionDataset
 
 .. autoclass:: OCRDataset

doctr/datasets/__init__.py

@@ -2,6 +2,7 @@
 from .coco_text import *
 from .cord import *
 from .detection import *
+from .layout import *
 from .doc_artefacts import *
 from .funsd import *
 from .ic03 import *

doctr/datasets/layout.py

@@ -0,0 +1,122 @@
+# Copyright (C) 2021-2026, 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 typing import Any
+
+import numpy as np
+
+from .datasets import AbstractDataset
+from .utils import pre_transform_multiclass
+
+__all__ = ["LayoutDataset"]
+
+
+class LayoutDataset(AbstractDataset):
+    """Implements a document layout detection dataset.
+
+    >>> from doctr.datasets import LayoutDataset
+    >>> train_set = LayoutDataset(
+    >>>     img_folder="/path/to/images",
+    >>>     label_path="/path/to/labels.json",
+    >>> )
+    >>> img, target = train_set[0]
+
+    Args:
+        img_folder: folder containing the dataset images
+        label_path: path to the labels.json file
+        use_polygons: whether to keep polygons instead of converting to straight boxes
+        **kwargs: keyword arguments from `AbstractDataset`
+    """
+
+    def __init__(
+        self,
+        img_folder: str,
+        label_path: str,
+        use_polygons: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            img_folder,
+            pre_transforms=pre_transform_multiclass,
+            **kwargs,
+        )
+        self._class_names: list[str] = []
+        if not os.path.exists(label_path):
+            raise FileNotFoundError(f"unable to locate {label_path}")
+        with open(label_path, "rb") as f:
+            labels = json.load(f)
+
+        self.data: list[tuple[str, tuple[np.ndarray, list[str]]]] = []
+        np_dtype = np.float32
+
+        for img_name, label in labels.items():
+            img_path = os.path.join(self.root, img_name)
+
+            # File existence check
+            if not os.path.exists(img_path):
+                raise FileNotFoundError(f"unable to locate {img_path}")
+
+            polygons = label.get("polygons")
+            class_names = label.get("class_names")
+
+            if polygons is None:
+                raise KeyError(f"missing 'polygons' for image: {img_name}")
+            if class_names is None:
+                raise KeyError(f"missing 'class_names' for image: {img_name}")
+
+            if len(polygons) != len(class_names):
+                raise ValueError(
+                    f"number of polygons ({len(polygons)}) does not match "
+                    f"number of class_names ({len(class_names)}) for image: {img_name}"
+                )
+
+            geoms, polygon_classes = self.format_polygons(
+                polygons=polygons,
+                class_names=class_names,
+                use_polygons=use_polygons,
+                np_dtype=np_dtype,
+            )
+
+            self.data.append((
+                img_name,
+                (
+                    np.asarray(geoms, dtype=np_dtype),
+                    polygon_classes,
+                ),
+            ))
+
+    def format_polygons(
+        self,
+        polygons: list,
+        class_names: list[str],
+        use_polygons: bool,
+        np_dtype: type,
+    ) -> tuple[np.ndarray, list[str]]:
+        """Format polygons into an array.
+
+        Args:
+            polygons: list of polygons
+            class_names: list of class names corresponding to polygons
+            use_polygons: whether polygons should be preserved
+            np_dtype: numpy dtype
+
+        Returns:
+            geoms: polygons or straight boxes
+            polygon_classes: list of classes
+        """
+        self._class_names += class_names
+
+        _polygons: np.ndarray = np.asarray(polygons, dtype=np_dtype)
+        if _polygons.ndim != 3 or _polygons.shape[1:] != (4, 2):
+            raise ValueError(f"polygons are expected to have shape (N, 4, 2), got {_polygons.shape}")
+        geoms = _polygons if use_polygons else np.concatenate((_polygons.min(axis=1), _polygons.max(axis=1)), axis=1)
+
+        return geoms, class_names
+
+    @property
+    def class_names(self) -> list[str]:
+        return sorted(set(self._class_names))

tests/common/test_cli.py

@@ -61,16 +61,16 @@ def test_parse_args_custom_values():
     assert args.output == "output.json"
 
 
-def test_main_with_image(mock_image_path):
-    output_path = "results.json"
-    cli.main(["--input_path", mock_image_path, "--output", output_path])
+def test_main_with_image(mock_image_path, tmp_path):
+    output_path = tmp_path / "results.json"
+    cli.main(["--input_path", mock_image_path, "--output", str(output_path)])
 
     assert Path(output_path).exists()
 
 
-def test_main_with_pdf(mock_pdf):
-    output_path = "results.json"
-    cli.main(["--input_path", mock_pdf, "--output", output_path])
+def test_main_with_pdf(mock_pdf, tmp_path):
+    output_path = tmp_path / "results.json"
+    cli.main(["--input_path", mock_pdf, "--output", str(output_path)])
 
     assert Path(output_path).exists()
 

tests/conftest.py

@@ -132,6 +132,37 @@ def mock_detection_label(tmpdir_factory):
     return str(labels_path)
 
 
+@pytest.fixture(scope="session")
+def mock_layout_label(tmpdir_factory):
+    folder = tmpdir_factory.mktemp("labels")
+
+    labels = {}
+    for idx in range(5):
+        labels[f"mock_image_file_{idx}.jpeg"] = {
+            "img_dimensions": [800, 600],
+            "img_hash": "dummy_hash",
+            "polygons": [
+                [[1, 2], [1, 3], [2, 1], [2, 3]],
+                [[10, 20], [10, 30], [20, 10], [20, 30]],
+                [[3, 2], [3, 3], [4, 1], [4, 3]],
+                [[30, 20], [30, 30], [40, 10], [40, 30]],
+            ],
+            "class_names": [
+                "Table",
+                "Header",
+                "Footer",
+                "Text",
+            ],
+        }
+
+    labels_path = folder.join("labels.json")
+
+    with open(labels_path, "w") as f:
+        json.dump(labels, f)
+
+    return str(labels_path)
+
+
 @pytest.fixture(scope="session")
 def mock_recognition_label(tmpdir_factory):
     label_file = tmpdir_factory.mktemp("labels").join("labels.json")

tests/pytorch/test_datasets_pt.py

@@ -1,3 +1,4 @@
+import json
 import os
 from shutil import move
 
@@ -177,6 +178,119 @@ def test_detection_dataset(mock_image_folder, mock_detection_label):
     move(os.path.join(ds.root, "tmp_file"), os.path.join(ds.root, img_name))
 
 
+@pytest.mark.parametrize(
+    "use_polygons",
+    [False, True],
+)
+def test_layout_dataset(mock_image_folder, mock_layout_label, use_polygons):
+    input_size = (1024, 1024)
+
+    ds = datasets.LayoutDataset(
+        img_folder=mock_image_folder,
+        label_path=mock_layout_label,
+        img_transforms=Resize(input_size),
+        use_polygons=use_polygons,
+    )
+
+    assert len(ds) == 5
+    img, target_dict = ds[0]
+    assert isinstance(img, torch.Tensor)
+    assert img.dtype == torch.float32
+    assert img.shape[-2:] == input_size
+    assert isinstance(target_dict, dict)
+    expected_classes = {"Table", "Header", "Footer", "Text"}
+    assert set(target_dict.keys()) == expected_classes
+    for class_name, target in target_dict.items():
+        assert isinstance(target, np.ndarray)
+        assert target.dtype == np.float32
+        if use_polygons:
+            assert target.ndim == 3
+            assert target.shape[1:] == (4, 2)
+        else:
+            assert target.ndim == 2
+            assert target.shape[1:] == (4,)
+        assert np.all(np.logical_and(target >= 0, target <= 1))
+    assert ds.class_names == sorted(expected_classes)
+    loader = DataLoader(ds, batch_size=2, collate_fn=ds.collate_fn)
+    images, targets = next(iter(loader))
+    assert isinstance(images, torch.Tensor)
+    assert images.shape == (2, 3, *input_size)
+    assert isinstance(targets, list)
+    assert all(isinstance(target, dict) for target in targets)
+    for target in targets:
+        assert set(target.keys()) == expected_classes
+        assert all(isinstance(v, np.ndarray) for v in target.values())
+
+    # File existence check
+    img_name, _ = ds.data[0]
+    move(os.path.join(ds.root, img_name), os.path.join(ds.root, "tmp_file"))
+    with pytest.raises(FileNotFoundError):
+        datasets.LayoutDataset(
+            img_folder=mock_image_folder,
+            label_path=mock_layout_label,
+        )
+    move(os.path.join(ds.root, "tmp_file"), os.path.join(ds.root, img_name))
+
+    with pytest.raises(FileNotFoundError):
+        datasets.LayoutDataset(
+            img_folder=mock_image_folder,
+            label_path="/tmp/does_not_exist.json",
+        )
+
+    with open(mock_layout_label) as f:
+        original_labels = json.load(f)
+    first_key = next(iter(original_labels))
+
+    test_cases = [
+        (
+            {"class_names": ["Text"]},
+            KeyError,
+            "missing 'polygons'",
+        ),
+        (
+            {"polygons": [[[0, 0], [1, 0], [1, 1], [0, 1]]]},
+            KeyError,
+            "missing 'class_names'",
+        ),
+        (
+            {
+                "polygons": [
+                    [[0, 0], [1, 0], [1, 1], [0, 1]],
+                    [[0, 0], [1, 0], [1, 1], [0, 1]],
+                ],
+                "class_names": ["Text"],
+            },
+            ValueError,
+            "number of polygons",
+        ),
+        (
+            {
+                "polygons": [[[0, 0], [1, 0], [1, 1]]],  # only 3 points
+                "class_names": ["Text"],
+            },
+            ValueError,
+            "polygons are expected to have shape",
+        ),
+    ]
+
+    for sample, exc_type, match in test_cases:
+        broken_labels = dict(original_labels)
+        broken_labels[first_key] = sample
+
+        with open(mock_layout_label, "w") as f:
+            json.dump(broken_labels, f)
+
+        with pytest.raises(exc_type, match=match):
+            datasets.LayoutDataset(
+                img_folder=mock_image_folder,
+                label_path=mock_layout_label,
+            )
+
+    # Restore original labels
+    with open(mock_layout_label, "w") as f:
+        json.dump(original_labels, f)
+
+
 def test_recognition_dataset(mock_image_folder, mock_recognition_label):
     input_size = (32, 128)
     ds = datasets.RecognitionDataset(