File size: 5,104 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | # Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Public ``rgb-letterbox-1024/1`` geometry. Independent of any detector graph.
This is the signed AXDocLayout-M v1.2 preprocess: HWC uint8 RGB, bilinear resize of the
longer side to 1024, white pad, odd remainder bottom/right, then inverse
mapping of letterbox-space boxes back onto the original raster.
It is not a DocLayout-YOLO drop-in. That incumbent uses a different colour
order, a different pad value, and a different class vocabulary. Those
differences are left unmapped on purpose.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from PIL import Image
from ax_doclayout.contract import PREPROCESS_ID, PageImage
from ax_doclayout.errors import InvalidPageImageError
SIZE = 1024
PAD_RGB = (255, 255, 255)
FILL_SCALE = 1.0 / 255.0
@dataclass(frozen=True, slots=True)
class Letterbox:
"""One 1024-square canvas plus the exact inverse-mapping terms."""
canvas: np.ndarray
scale: float
pad_left: int
pad_top: int
pad_right: int
pad_bottom: int
size: int = SIZE
def __post_init__(self) -> None:
canvas = np.ascontiguousarray(self.canvas)
if canvas.ndim != 3 or canvas.shape != (self.size, self.size, 3):
raise InvalidPageImageError("letterbox canvas must be size x size x 3")
if canvas.dtype != np.uint8:
raise InvalidPageImageError("letterbox canvas must be uint8")
object.__setattr__(self, "canvas", canvas)
def letterbox_rgb(rgb: np.ndarray, *, size: int = SIZE) -> Letterbox:
"""Letterbox a decoded HWC uint8 RGB array under ``rgb-letterbox-1024/1``."""
if not isinstance(rgb, np.ndarray):
raise InvalidPageImageError("rgb must be a numpy array")
if rgb.ndim != 3 or rgb.shape[2] != 3:
raise InvalidPageImageError("rgb must be HWC with three channels")
if rgb.dtype != np.uint8:
raise InvalidPageImageError("rgb must be uint8")
height, width = (int(rgb.shape[0]), int(rgb.shape[1]))
if height < 1 or width < 1:
raise InvalidPageImageError("rgb must not be empty")
scale = size / max(width, height)
new_width = max(1, round(width * scale))
new_height = max(1, round(height * scale))
image = Image.fromarray(np.ascontiguousarray(rgb), mode="RGB")
resized = image.resize((new_width, new_height), Image.BILINEAR)
pad_width = size - new_width
pad_height = size - new_height
pad_left = pad_width // 2
pad_right = pad_width - pad_left
pad_top = pad_height // 2
pad_bottom = pad_height - pad_top
if pad_right < 0 or pad_bottom < 0:
raise InvalidPageImageError("letterbox overflow")
canvas = Image.new("RGB", (size, size), PAD_RGB)
canvas.paste(resized, (pad_left, pad_top))
return Letterbox(
canvas=np.asarray(canvas, dtype=np.uint8),
scale=float(scale),
pad_left=int(pad_left),
pad_top=int(pad_top),
pad_right=int(pad_right),
pad_bottom=int(pad_bottom),
size=int(size),
)
def letterbox_page(page: PageImage, *, size: int = SIZE) -> Letterbox:
"""Letterbox a validated ``PageImage``."""
return letterbox_rgb(page.rgb, size=size)
def letterbox_nchw(letterbox: Letterbox) -> np.ndarray:
"""Return a batch-1 NCHW float32 tensor in ``[0, 1]``."""
chw = np.transpose(letterbox.canvas, (2, 0, 1)).astype(np.float32)
chw *= np.float32(FILL_SCALE)
return np.ascontiguousarray(chw[None, ...])
def inverse_letterbox_box(
box: tuple[float, float, float, float] | list[float],
letterbox: Letterbox,
*,
width_px: int,
height_px: int,
) -> tuple[float, float, float, float] | None:
"""Map a letterbox xyxy box onto the original raster, or drop it."""
if width_px < 1 or height_px < 1:
return None
try:
x_min, y_min, x_max, y_max = (float(value) for value in box)
except (TypeError, ValueError):
return None
if not all(np.isfinite(value) for value in (x_min, y_min, x_max, y_max)):
return None
scale = letterbox.scale
if scale <= 0.0:
return None
mapped = (
(x_min - letterbox.pad_left) / scale,
(y_min - letterbox.pad_top) / scale,
(x_max - letterbox.pad_left) / scale,
(y_max - letterbox.pad_top) / scale,
)
clipped_x_min = min(max(mapped[0], 0.0), float(width_px))
clipped_y_min = min(max(mapped[1], 0.0), float(height_px))
clipped_x_max = min(max(mapped[2], 0.0), float(width_px))
clipped_y_max = min(max(mapped[3], 0.0), float(height_px))
if clipped_x_min >= clipped_x_max or clipped_y_min >= clipped_y_max:
return None
return clipped_x_min, clipped_y_min, clipped_x_max, clipped_y_max
def orig_target_sizes(letterbox: Letterbox) -> np.ndarray:
"""Sizes fed to the exported graph. The graph is bound to the letterbox square."""
side = np.int64(letterbox.size)
return np.ascontiguousarray([[side, side]])
PREPROCESS = PREPROCESS_ID
|