File size: 10,398 Bytes
fc27bf1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-19 for v1.3; 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],
        }