File size: 4,716 Bytes
1b4ed07 | 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 | # 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,
)
|