inference/src/ax_doclayout/postprocess.py

# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Map letterbox-space detector outputs onto canonical ``PageStructure`` elements.

Class id 0 is reserved and never returned. Names that exist only on the
DocLayout-YOLO incumbent (``figure_caption``, ``table_caption``, and similar)
are not invented here.
"""

from __future__ import annotations

from collections.abc import Sequence

import numpy as np

from ax_doclayout.contract import (
    CLASS_NAMES,
    DEFAULT_THRESHOLD_PROFILE,
    BoundingBox,
    Element,
    InvalidPageStructureError,
    ModelIdentity,
    PageImage,
    PageStructure,
)
from ax_doclayout.preprocess import Letterbox

DEFAULT_SCORE_THRESHOLD = 0.30


def elements_from_letterbox_detections(
    *,
    labels: object,
    boxes: object,
    scores: object,
    page: PageImage,
    letterbox: Letterbox,
    threshold: float = DEFAULT_SCORE_THRESHOLD,
) -> tuple[Element, ...]:
    """Keep finite, in-range detections and invert them onto the original raster."""

    try:
        threshold_value = float(threshold)
    except (TypeError, ValueError) as exc:
        raise InvalidPageStructureError("threshold must be within [0, 1]") from exc
    if not 0.0 <= threshold_value <= 1.0:
        raise InvalidPageStructureError("threshold must be within [0, 1]")
    try:
        label_array = np.asarray(labels)
        box_array = np.asarray(boxes, dtype=np.float64)
        score_array = np.asarray(scores, dtype=np.float64)
    except (TypeError, ValueError) as exc:
        raise InvalidPageStructureError(
            "detection arrays must be rectangular numeric arrays"
        ) from exc
    if label_array.ndim == 0 or box_array.ndim == 0 or score_array.ndim == 0:
        raise InvalidPageStructureError("detection arrays must share a length")
    if label_array.ndim == 2 and label_array.shape[0] == 1:
        if box_array.ndim != 3 or score_array.ndim != 2:
            raise InvalidPageStructureError("detection arrays must share a batch dimension")
        if box_array.shape[0] != 1 or score_array.shape[0] != 1:
            raise InvalidPageStructureError("detection arrays must share a batch dimension")
        label_array = label_array[0]
        box_array = box_array[0]
        score_array = score_array[0]
    if box_array.ndim != 2 or box_array.shape[1] != 4:
        raise InvalidPageStructureError("boxes must be an N x 4 xyxy array")
    if label_array.ndim != 1 or score_array.ndim != 1:
        raise InvalidPageStructureError("detection arrays must share a length")
    if label_array.shape[0] != box_array.shape[0] or label_array.shape[0] != score_array.shape[0]:
        raise InvalidPageStructureError("detection arrays must share a length")
    if not np.isfinite(letterbox.scale) or letterbox.scale <= 0.0:
        return ()
    # Apply identical float64 geometry in bulk before allocating Python objects.
    valid = (
        np.isfinite(box_array).all(axis=1)
        & np.isfinite(score_array)
        & (score_array >= threshold_value)
        & (score_array <= 1.0)
    )
    indices = np.flatnonzero(valid)
    mapped = box_array[indices].copy()
    mapped -= [letterbox.pad_left, letterbox.pad_top,
               letterbox.pad_left, letterbox.pad_top]
    mapped /= letterbox.scale
    np.clip(mapped, 0.0, [page.width_px, page.height_px,
                         page.width_px, page.height_px], out=mapped)
    nonempty = (mapped[:, 0] < mapped[:, 2]) & (mapped[:, 1] < mapped[:, 3])
    kept: list[Element] = []
    for index, box in zip(indices[nonempty], mapped[nonempty], strict=True):
        try:
            class_id = int(label_array[index])
        except (TypeError, ValueError, OverflowError):
            continue
        if class_id not in CLASS_NAMES:
            continue
        kept.append(
            Element(
                id=f"{page.page_index}:{len(kept)}",
                class_id=class_id,
                confidence=float(score_array[index]),
                bbox=BoundingBox(*(float(value) for value in box)),
            )
        )
    return tuple(kept)


def page_structure_from_elements(
    *,
    page: PageImage,
    elements: Sequence[Element],
    identity: ModelIdentity,
    threshold_override: float | None = None,
) -> PageStructure:
    """Bind already-inverted elements to a caller-supplied model identity."""

    return PageStructure(
        model=identity,
        page_index=page.page_index,
        width_px=page.width_px,
        height_px=page.height_px,
        elements=tuple(elements),
        threshold_profile=DEFAULT_THRESHOLD_PROFILE,
        threshold_override=threshold_override,
    )