| |
| |
| |
| """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_mlx.contract import PREPROCESS_ID, PageImage |
| from ax_doclayout_mlx.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 |
|
|
