inference/src/ax_doclayout/local_detection.py

# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Canonical JSON for unqualified local detections.

Both the local ONNX runner and the local PyTorch runner serialize through
this function. The payload is not a published ``PageStructure``.
"""

from __future__ import annotations

from collections.abc import Sequence

from ax_doclayout.artifact import canonical_json_bytes
from ax_doclayout.contract import Element

LOCAL_DETECTION_SCHEMA = "ax-doclayout.local-detection/1"
LOCAL_QUALIFICATION = "local-bundle-unqualified"
# Shared unqualified serialization. PyTorch and ONNX agree well inside
# these places; the published parity gate is coarser (0.5 px, 0.001 score).
SCORE_DECIMALS = 4
BOX_DECIMALS = 3


def _places(value: float, places: int) -> float:
    return round(float(value), places)


def canonical_local_detection_bytes(
    *,
    page_index: int,
    width_px: int,
    height_px: int,
    elements: Sequence[Element],
    threshold: float,
) -> bytes:
    """Serialize one unqualified page detection with sorted keys."""

    payload = {
        "elements": [
            {
                "bbox": {
                    "x_max": _places(element.bbox.x_max, BOX_DECIMALS),
                    "x_min": _places(element.bbox.x_min, BOX_DECIMALS),
                    "y_max": _places(element.bbox.y_max, BOX_DECIMALS),
                    "y_min": _places(element.bbox.y_min, BOX_DECIMALS),
                },
                "class_id": int(element.class_id),
                "class_name": element.class_name,
                "confidence": _places(element.confidence, SCORE_DECIMALS),
                "id": element.id,
            }
            for element in elements
        ],
        "height_px": int(height_px),
        "page_index": int(page_index),
        "qualification": LOCAL_QUALIFICATION,
        "schema_version": LOCAL_DETECTION_SCHEMA,
        "threshold": float(threshold),
        "width_px": int(width_px),
    }
    return canonical_json_bytes(payload)