processing_nanovdr_doc.py

"""NanoVDR document tower — standalone image processor for the Hub.

The document tower does not consume a page as one 448x448 view. It consumes a
variable number of aspect-ratio-matched crops plus a thumbnail, zero-padded to
a fixed tile budget and accompanied by a mask saying which tiles are real. That
is preprocessing, so it belongs here rather than in the modeling file or in a
training dataset, and it is the reason a stock image processor is not
interchangeable with this one: a stock processor emits a single view and the
model, given a single view, silently produces a materially worse embedding.

``NanoVDRDocImageProcessor`` therefore emits both of the model's inputs::

    pixel_values : (B, T, 3, H, W) float32   zero-padded in normalised space
    tile_mask    : (B, T) bool               True for real tiles

which is what makes ``processor(images=pages)`` -> ``model(**inputs)`` correct,
and what lets sentence-transformers drive the tower without a custom module.

This file is self-contained: it does not import from the rest of the package,
and it is the only source of truth for tiling.
"""

from __future__ import annotations

import copy
from typing import List, Optional, Sequence, Union

import numpy as np
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from transformers.image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    resize,
    to_channel_dimension_format,
)
from transformers.image_utils import (
    ChannelDimension,
    PILImageResampling,
    infer_channel_dimension_format,
    make_list_of_images,
    to_numpy_array,
)

IMAGE_SIZE = 448

# InternViT-300M-448px-V2_5's statistics. Kept explicit rather than inherited so
# that a checkpoint's preprocessing is fully described by its own config.
IMAGE_MEAN = [0.485, 0.456, 0.406]
IMAGE_STD = [0.229, 0.224, 0.225]

__all__ = ["NanoVDRDocImageProcessor", "dynamic_tile", "IMAGE_SIZE"]


# --------------------------------------------------------------------------
# Dynamic tiling (InternVL-V2 partition rule)
# --------------------------------------------------------------------------
def _closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    best_diff, best = float("inf"), (1, 1)
    area = width * height
    for ratio in target_ratios:
        target = ratio[0] / ratio[1]
        diff = abs(aspect_ratio - target)
        if diff < best_diff:
            best_diff, best = diff, ratio
        elif diff == best_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best = ratio
    return best


def dynamic_tile(
    image,
    min_num: int = 1,
    max_num: int = 6,
    image_size: int = IMAGE_SIZE,
    use_thumbnail: bool = True,
) -> list:
    """Split a page into aspect-ratio-matched ``image_size`` crops.

    The grid whose aspect ratio is closest to the page's is chosen, the page is
    resized to that grid's exact pixel extent and cut into tiles. The final
    element is a whole-page thumbnail when ``use_thumbnail`` is set and more
    than one crop was produced, which is what keeps global layout available
    after the page has been cut up. Returns a list of PIL images.
    """
    orig_w, orig_h = image.size
    aspect_ratio = orig_w / orig_h
    target_ratios = sorted(
        {
            (i, j)
            for n in range(min_num, max_num + 1)
            for i in range(1, n + 1)
            for j in range(1, n + 1)
            if i * j <= max_num and i * j >= min_num
        },
        key=lambda x: x[0] * x[1],
    )
    cols, rows = _closest_aspect_ratio(aspect_ratio, target_ratios, orig_w, orig_h, image_size)
    resized = image.resize((image_size * cols, image_size * rows))
    tiles = []
    for i in range(cols * rows):
        c, r = i % cols, i // cols
        box = (c * image_size, r * image_size, (c + 1) * image_size, (r + 1) * image_size)
        tiles.append(resized.crop(box))
    if use_thumbnail and len(tiles) > 1:
        tiles.append(image.resize((image_size, image_size)))
    return tiles


# --------------------------------------------------------------------------
# Processor
# --------------------------------------------------------------------------
class NanoVDRDocImageProcessor(BaseImageProcessor):
    """Turn page images into the document tower's exact inputs.

    ``model_input_names`` advertises ``tile_mask`` alongside ``pixel_values``,
    so frameworks that forward a processor's outputs into a model - including
    sentence-transformers, which filters kwargs against the forward signature -
    carry the mask through without special-casing this model.

    Setting ``do_tile=False`` reproduces a plain single-view processor. It is
    available for ablations and for encoders that were trained without tiling;
    it is not the right setting for any released NanoVDR document tower.
    """

    model_input_names = ["pixel_values", "tile_mask"]

    def __init__(
        self,
        do_tile: bool = True,
        tile_min_num: int = 1,
        tile_max_num: int = 6,
        tile_max_total: Optional[int] = None,
        tile_use_thumbnail: bool = True,
        image_size: int = IMAGE_SIZE,
        do_resize: bool = True,
        size: Optional[dict] = None,
        resample: PILImageResampling = PILImageResampling.BICUBIC,
        do_center_crop: bool = True,
        crop_size: Optional[dict] = None,
        do_rescale: bool = True,
        rescale_factor: float = 1 / 255,
        do_normalize: bool = True,
        image_mean: Optional[Union[float, Sequence[float]]] = None,
        image_std: Optional[Union[float, Sequence[float]]] = None,
        do_convert_rgb: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.do_tile = do_tile
        self.tile_min_num = tile_min_num
        self.tile_max_num = tile_max_num
        self.tile_use_thumbnail = tile_use_thumbnail
        self.image_size = image_size
        # The budget the model pads to. Derived from the tile settings when not
        # given, because a mismatch here is a silent shape bug at training time.
        self.tile_max_total = int(
            tile_max_total if tile_max_total is not None
            else tile_max_num + (1 if tile_use_thumbnail else 0)
        )

        self.do_resize = do_resize
        self.size = get_size_dict(size if size is not None else {"shortest_edge": image_size},
                                  default_to_square=False)
        self.resample = resample
        self.do_center_crop = do_center_crop
        self.crop_size = get_size_dict(
            crop_size if crop_size is not None else {"height": image_size, "width": image_size},
            param_name="crop_size",
        )
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.image_mean = list(image_mean) if image_mean is not None else list(IMAGE_MEAN)
        self.image_std = list(image_std) if image_std is not None else list(IMAGE_STD)
        self.do_convert_rgb = do_convert_rgb

    # -- construction ------------------------------------------------------
    @classmethod
    def from_image_processor(cls, processor, **tile_kwargs) -> "NanoVDRDocImageProcessor":
        """Adopt an existing processor's pixel statistics and add tiling.

        The document tower's pixel normalisation is its visual encoder's, so the
        statistics are taken from that encoder's own processor instead of being
        restated. Everything about tiling comes from ``tile_kwargs``.
        """
        src = processor.to_dict() if hasattr(processor, "to_dict") else dict(processor)
        pixel_keys = (
            "do_resize", "size", "resample", "do_center_crop", "crop_size",
            "do_rescale", "rescale_factor", "do_normalize", "image_mean",
            "image_std", "do_convert_rgb",
        )
        kept = {k: src[k] for k in pixel_keys if k in src and src[k] is not None}
        kept.update(tile_kwargs)
        return cls(**kept)

    # -- internals ---------------------------------------------------------
    def _to_pixel_array(
        self,
        image,
        input_data_format: Optional[ChannelDimension] = None,
    ) -> np.ndarray:
        """Apply the pixel pipeline to one already-cropped view -> (3, H, W)."""
        if self.do_convert_rgb:
            image = convert_to_rgb(image)
        image = to_numpy_array(image)
        fmt = input_data_format or infer_channel_dimension_format(image)

        if self.do_resize:
            out_size = get_resize_output_image_size(
                image,
                size=(self.size["shortest_edge"] if "shortest_edge" in self.size
                      else (self.size["height"], self.size["width"])),
                default_to_square="shortest_edge" not in self.size,
                input_data_format=fmt,
            )
            image = resize(image, size=out_size, resample=self.resample, input_data_format=fmt)
        if self.do_center_crop:
            image = center_crop(
                image,
                size=(self.crop_size["height"], self.crop_size["width"]),
                input_data_format=fmt,
            )
        if self.do_rescale:
            image = self.rescale(image, scale=self.rescale_factor, input_data_format=fmt)
        if self.do_normalize:
            image = self.normalize(
                image, mean=self.image_mean, std=self.image_std, input_data_format=fmt
            )
        return to_channel_dimension_format(image, ChannelDimension.FIRST, input_channel_dim=fmt)

    def _tiles_for(self, image) -> List:
        return dynamic_tile(
            image.convert("RGB") if hasattr(image, "convert") else image,
            min_num=self.tile_min_num,
            max_num=self.tile_max_num,
            image_size=self.image_size,
            use_thumbnail=self.tile_use_thumbnail,
        )[: self.tile_max_total]

    # -- entry point -------------------------------------------------------
    def preprocess(
        self,
        images,
        return_tensors: Optional[str] = None,
        input_data_format: Optional[ChannelDimension] = None,
        **kwargs,
    ) -> BatchFeature:
        """Tile and preprocess pages.

        Returns ``pixel_values`` of shape ``(B, T, 3, H, W)`` zero-padded to the
        tile budget and a ``(B, T)`` boolean ``tile_mask``, or ``(B, 3, H, W)``
        and no mask when ``do_tile`` is False. Padding is zeros *after*
        normalisation, matching what the model's masked pooling expects.
        """
        if kwargs:
            unknown = [k for k in kwargs if not hasattr(self, k)]
            if unknown:
                raise TypeError(f"unexpected preprocess argument(s): {unknown}")
            # Per-call overrides act on a copy, so a processor shared between
            # dataloader workers is never mutated mid-flight.
            clone = copy.copy(self)
            for k, v in kwargs.items():
                setattr(clone, k, v)
            return clone.preprocess(
                images, return_tensors=return_tensors, input_data_format=input_data_format
            )

        images = make_list_of_images(images)

        if not self.do_tile:
            px = np.stack([self._to_pixel_array(im, input_data_format) for im in images])
            return BatchFeature(data={"pixel_values": px}, tensor_type=return_tensors)

        T = self.tile_max_total
        pixel_values, tile_mask = [], []
        for image in images:
            views = [self._to_pixel_array(t, input_data_format) for t in self._tiles_for(image)]
            n = len(views)
            if n == 0:
                raise ValueError("tiling produced no views for an input image")
            stacked = np.stack(views)
            padded = np.zeros((T, *stacked.shape[1:]), dtype=stacked.dtype)
            padded[:n] = stacked
            mask = np.zeros(T, dtype=bool)
            mask[:n] = True
            pixel_values.append(padded)
            tile_mask.append(mask)

        return BatchFeature(
            data={
                "pixel_values": np.stack(pixel_values),
                "tile_mask": np.stack(tile_mask),
            },
            tensor_type=return_tensors,
        )