inference/src/ax_doclayout/contract.py

# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Typed values of the raster detection contract (spec sections 4, 5, 7, 9).

Everything here is the consumer-facing shape; nothing here runs a model.
Every violation raises an ``AXDocLayoutError`` subclass.
"""

from __future__ import annotations

import math
import re
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from numbers import Integral, Real
from typing import Any

import numpy as np

from ax_doclayout.errors import AXDocLayoutError, InvalidPageImageError

# stayed because the payload shape is unchanged. The previous prefix is
# not an alias.
DETECTION_SCHEMA_VERSION = "ax-doclayout.detection/1"
TAXONOMY_ID = "ax-doclayout-11/1"
PREPROCESS_ID = "rgb-letterbox-1024/1"
COORDINATE_SPACE = "raster-pixels-top-left"
DEFAULT_THRESHOLD_PROFILE = "default/1"

# Public class ids are preserved; id 0 is reserved and never returned.
CLASS_NAMES: Mapping[int, str] = {
    1: "caption",
    2: "footnote",
    3: "formula",
    4: "list_item",
    5: "page_footer",
    6: "page_header",
    7: "picture",
    8: "section_header",
    9: "table",
    10: "text",
    11: "title",
}

# Rasters past this many pixels are rejected before any preprocessing.
MAX_PAGE_PIXELS = 50_000_000

HUB_COMMIT_RE = re.compile(r"^[0-9a-f]{40}$")
SHA256_RE = re.compile(r"^[0-9a-f]{64}$")


class InvalidPageStructureError(AXDocLayoutError):
    """A detection result violates the canonical contract."""


def _integer(value: object, name: str, *, minimum: int = 0) -> int:
    """Return ``value`` as a plain int; bools and non-integral types are rejected."""

    if isinstance(value, bool) or not isinstance(value, Integral):
        raise InvalidPageStructureError(f"{name} must be an integer")
    result = int(value)
    if result < minimum:
        raise InvalidPageStructureError(f"{name} must be >= {minimum}")
    return result


def _finite(value: object, name: str) -> float:
    if isinstance(value, bool) or not isinstance(value, Real):
        raise InvalidPageStructureError(f"{name} must be a real number")
    result = float(value)
    if not math.isfinite(result):
        raise InvalidPageStructureError(f"{name} must be finite")
    return result


@dataclass(frozen=True, slots=True, eq=False)
class PageImage:
    """One decoded page raster: HWC ``uint8`` RGB plus its zero-based page index.

    The array is copied to C-contiguous, read-only form on construction so
    later caller mutation can neither change a validated page nor be frozen
    as a side effect of constructing one.
    """

    rgb: np.ndarray
    page_index: int

    def __post_init__(self) -> None:
        if not isinstance(self.rgb, np.ndarray):
            raise InvalidPageImageError("rgb must be a numpy array")
        if self.rgb.ndim != 3 or self.rgb.shape[2] != 3:
            raise InvalidPageImageError("rgb must be HWC with three channels")
        if self.rgb.dtype != np.uint8:
            raise InvalidPageImageError("rgb must be uint8")
        height, width = self.rgb.shape[:2]
        if height == 0 or width == 0:
            raise InvalidPageImageError("rgb must not be empty")
        if height * width > MAX_PAGE_PIXELS:
            raise InvalidPageImageError("rgb exceeds the page pixel limit")
        if isinstance(self.page_index, bool) or not isinstance(
            self.page_index, Integral
        ):
            raise InvalidPageImageError("page_index must be an int")
        if self.page_index < 0:
            raise InvalidPageImageError(
                "page_index must be zero-based and non-negative"
            )
        rgb = np.array(self.rgb, dtype=np.uint8, order="C", copy=True)
        rgb.setflags(write=False)
        object.__setattr__(self, "rgb", rgb)
        object.__setattr__(self, "page_index", int(self.page_index))

    @property
    def width_px(self) -> int:
        return int(self.rgb.shape[1])

    @property
    def height_px(self) -> int:
        return int(self.rgb.shape[0])


@dataclass(frozen=True, slots=True)
class BoundingBox:
    """Half-open pixel-edge box ``[x_min, x_max) x [y_min, y_max)`` on the original raster."""

    x_min: float
    y_min: float
    x_max: float
    y_max: float

    def validate(self, width_px: int, height_px: int) -> None:
        x_min = _finite(self.x_min, "x_min")
        y_min = _finite(self.y_min, "y_min")
        x_max = _finite(self.x_max, "x_max")
        y_max = _finite(self.y_max, "y_max")
        if not (0 <= x_min < x_max <= width_px):
            raise InvalidPageStructureError(
                "bounding box x range must satisfy 0 <= x_min < x_max <= width_px"
            )
        if not (0 <= y_min < y_max <= height_px):
            raise InvalidPageStructureError(
                "bounding box y range must satisfy 0 <= y_min < y_max <= height_px"
            )

    def to_json(self) -> dict[str, float]:
        return {
            "x_min": float(self.x_min),
            "y_min": float(self.y_min),
            "x_max": float(self.x_max),
            "y_max": float(self.y_max),
        }


@dataclass(frozen=True, slots=True)
class Element:
    """One detected region. ``id`` is opaque and stable only within its result."""

    id: str
    class_id: int
    confidence: float
    bbox: BoundingBox

    @property
    def class_name(self) -> str:
        if isinstance(self.class_id, bool) or self.class_id not in CLASS_NAMES:
            raise InvalidPageStructureError(f"unknown class id {self.class_id!r}")
        return CLASS_NAMES[int(self.class_id)]

    def to_json(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "class_id": int(self.class_id),
            "class_name": self.class_name,
            "confidence": float(self.confidence),
            "bbox": self.bbox.to_json(),
        }


@dataclass(frozen=True, slots=True)
class ModelIdentity:
    """The immutable artifact a result came from (spec section 7)."""

    id: str
    revision: str
    artifact_sha256: str

    def __post_init__(self) -> None:
        if not isinstance(self.id, str) or not self.id.strip():
            raise InvalidPageStructureError("model id must be a non-empty string")
        if (
            not isinstance(self.revision, str)
            or HUB_COMMIT_RE.fullmatch(self.revision) is None
        ):
            raise InvalidPageStructureError(
                "model revision must be a 40-character immutable Hub commit"
            )
        if (
            not isinstance(self.artifact_sha256, str)
            or SHA256_RE.fullmatch(self.artifact_sha256) is None
        ):
            raise InvalidPageStructureError(
                "artifact_sha256 must be a 64-character hex digest"
            )

    def to_json(self) -> dict[str, str]:
        return {
            "id": self.id,
            "revision": self.revision,
            "artifact_sha256": self.artifact_sha256,
            "taxonomy": TAXONOMY_ID,
            "preprocess": PREPROCESS_ID,
        }


@dataclass(frozen=True, slots=True)
class PageStructure:
    """The canonical detection result for one page (spec section 9)."""

    model: ModelIdentity
    page_index: int
    width_px: int
    height_px: int
    elements: Sequence[Element] = field(default_factory=tuple)
    threshold_profile: str = DEFAULT_THRESHOLD_PROFILE
    threshold_override: float | None = None

    def __post_init__(self) -> None:
        if not isinstance(self.model, ModelIdentity):
            raise InvalidPageStructureError("model must be a ModelIdentity")
        object.__setattr__(self, "page_index", _integer(self.page_index, "page_index"))
        object.__setattr__(
            self, "width_px", _integer(self.width_px, "width_px", minimum=1)
        )
        object.__setattr__(
            self, "height_px", _integer(self.height_px, "height_px", minimum=1)
        )
        if (
            not isinstance(self.threshold_profile, str)
            or not self.threshold_profile.strip()
        ):
            raise InvalidPageStructureError(
                "threshold_profile must be a non-empty string"
            )
        if self.threshold_override is not None:
            override = _finite(self.threshold_override, "threshold_override")
            if not 0.0 <= override <= 1.0:
                raise InvalidPageStructureError(
                    "threshold_override must be within [0, 1]"
                )
            object.__setattr__(self, "threshold_override", override)
        elements = tuple(self.elements)
        object.__setattr__(self, "elements", elements)
        seen: set[str] = set()
        for element in elements:
            if not isinstance(element, Element):
                raise InvalidPageStructureError("elements must be Element values")
            if not isinstance(element.id, str) or not element.id:
                raise InvalidPageStructureError("element ids must be non-empty strings")
            if element.id in seen:
                raise InvalidPageStructureError(f"duplicate element id {element.id!r}")
            seen.add(element.id)
            _ = element.class_name  # raises for reserved or unknown ids
            confidence = _finite(element.confidence, "confidence")
            if not 0.0 <= confidence <= 1.0:
                raise InvalidPageStructureError("confidence must be within [0, 1]")
            if not isinstance(element.bbox, BoundingBox):
                raise InvalidPageStructureError("element bbox must be a BoundingBox")
            element.bbox.validate(self.width_px, self.height_px)

    def to_json(self) -> dict[str, Any]:
        inference: dict[str, Any] = {"threshold_profile": self.threshold_profile}
        if self.threshold_override is not None:
            inference["threshold_override"] = float(self.threshold_override)
        return {
            "schema_version": DETECTION_SCHEMA_VERSION,
            "model": self.model.to_json(),
            "page": {
                "page_index": int(self.page_index),
                "width_px": int(self.width_px),
                "height_px": int(self.height_px),
                "coordinate_space": COORDINATE_SPACE,
            },
            "inference": inference,
            "elements": [element.to_json() for element in self.elements],
        }