From 825414b0c4b79723605ef07a4575fc78b639ac3c Mon Sep 17 00:00:00 2001 From: Hulk Date: Sat, 3 Oct 2026 13:40:46 +0800 Subject: [PATCH] fix: preserve ColModernVBERT image patch geometry --- fastembed/image/transform/operators.py | 29 +++- .../colmodernvbert.py | 55 +++--- tests/test_colmodernvbert_preprocessing.py | 159 ++++++++++++++++++ tests/test_image_transform.py | 68 ++++++++ 4 files changed, 274 insertions(+), 37 deletions(-) create mode 100644 tests/test_colmodernvbert_preprocessing.py diff --git a/fastembed/image/transform/operators.py b/fastembed/image/transform/operators.py index c5155be4..d3fe9c58 100644 --- a/fastembed/image/transform/operators.py +++ b/fastembed/image/transform/operators.py @@ -191,7 +191,15 @@ def __init__( self.resample = resample def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]: # type: ignore[override] + patches, _ = self.split_with_grid(images) + return patches + + def split_with_grid( + self, images: list[NumpyArray] + ) -> tuple[list[list[NumpyArray]], list[tuple[int, int]]]: + """Return patches and their (rows, cols), or (0, 0) for unsplit images.""" result = [] + image_grid = [] for image in images: # Assume (C, H, W) format @@ -204,6 +212,7 @@ def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]: # type: # Calculate the number of splits needed num_splits_h = math.ceil(height / max_height) num_splits_w = math.ceil(width / max_width) + image_grid.append((num_splits_h, num_splits_w)) # Calculate optimal patch dimensions optimal_height = math.ceil(height / num_splits_h) @@ -234,12 +243,13 @@ def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]: # type: frames.append(global_view) else: # Image is small enough, no splitting needed + image_grid.append((0, 0)) frames.append(image) # Append (not extend) to preserve per-image grouping result.append(frames) - return result + return result, image_grid class SquareResize(Transform): @@ -272,10 +282,23 @@ def __init__(self, transforms: list[Transform]): self.transforms = transforms def __call__( - self, images: list[Image.Image] | list[NumpyArray] + self, + images: list[Image.Image] | list[NumpyArray], + *, + metadata: dict[str, Any] | None = None, ) -> list[NumpyArray] | list[Image.Image]: + """Apply transforms without changing their usual return type. + + If supplied, metadata["image_grid"] records (rows, cols) for each image, + or (0, 0) for images that are not split. + """ + if metadata is not None: + metadata["image_grid"] = [(0, 0) for _ in images] for transform in self.transforms: - images = transform(images) + if metadata is not None and isinstance(transform, ImageSplitter): + images, metadata["image_grid"] = transform.split_with_grid(images) # type: ignore[assignment,arg-type] + else: + images = transform(images) return images @classmethod diff --git a/fastembed/late_interaction_multimodal/colmodernvbert.py b/fastembed/late_interaction_multimodal/colmodernvbert.py index f8fe0154..7e09caad 100644 --- a/fastembed/late_interaction_multimodal/colmodernvbert.py +++ b/fastembed/late_interaction_multimodal/colmodernvbert.py @@ -222,11 +222,15 @@ def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutpu for image in images ] assert self.processor is not None, "Processor is not initialized" - processed = self.processor(image_files) + processor_metadata: dict[str, Any] = {} + processed = self.processor(image_files, metadata=processor_metadata) encoded, attention_mask, metadata = self._process_nested_patches(processed) # type: ignore[arg-type] + metadata.update(processor_metadata) onnx_input = {"pixel_values": encoded, "attention_mask": attention_mask} - onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs) + onnx_input = self._preprocess_onnx_image_input( + onnx_input, image_grid=metadata["image_grid"], **kwargs + ) model_output = self.model.run(None, onnx_input) # type: ignore[union-attr] return OnnxOutputContext( @@ -276,17 +280,22 @@ def _process_nested_patches( return encoded, attention_mask, metadata # type: ignore[return-value] def _preprocess_onnx_image_input( - self, onnx_input: dict[str, np.ndarray], **kwargs: Any + self, + onnx_input: dict[str, np.ndarray], + *, + image_grid: list[tuple[int, int]] | None = None, + **kwargs: Any, ) -> dict[str, NumpyArray]: """ Add text input placeholders for image data, following Idefics3 processing logic. - Constructs input_ids dynamically based on the actual number of image patches, - using the same token expansion logic as Idefics3Processor. + Constructs input_ids from the actual image patch grid, using the same + token expansion logic as Idefics3Processor. Args: onnx_input: Dict with 'pixel_values' (batch, num_patches, C, H, W) and 'attention_mask' (batch, num_patches) indicating real patches + image_grid: Actual (rows, cols) from image splitting, (0, 0) if unsplit **kwargs: Additional arguments Returns: @@ -298,15 +307,18 @@ def _preprocess_onnx_image_input( pixel_values = onnx_input["pixel_values"] batch_size = pixel_values.shape[0] + if image_grid is None or len(image_grid) != batch_size: + raise ValueError("The image patch grid is required for each image in the batch") batch_input_ids = [] - # Build input_ids for each image based on its actual patch count - for i in range(batch_size): + # A patch count cannot distinguish portrait, landscape, and square grids. + for i, (rows, cols) in enumerate(image_grid): # Count real patches (non-padded) from attention mask patch_count = int(np.sum(patch_attention_mask[i])) - # Compute rows/cols from patch count - rows, cols = self._compute_rows_cols_from_patches(patch_count) + valid_grid = (rows == 0 and cols == 0) or (rows > 0 and cols > 0) + if not valid_grid or rows * cols + 1 != patch_count: + raise ValueError("The image patch grid does not match the number of patches") # Build input_ids for this image input_ids = self._build_input_ids_for_image(rows, cols) @@ -340,31 +352,6 @@ def _preprocess_onnx_image_input( onnx_input["attention_mask"] = attention_mask return onnx_input - @staticmethod - def _compute_rows_cols_from_patches(patch_count: int) -> tuple[int, int]: - if patch_count <= 1: - return 0, 0 - - # Subtract 1 for the global image - grid_patches = patch_count - 1 - - # Find rows and cols (assume square or near-square grid) - rows = int(grid_patches**0.5) - cols = grid_patches // rows - - # Verify the calculation - if rows * cols + 1 != patch_count: - # Handle non-square grids - for r in range(1, grid_patches + 1): - if grid_patches % r == 0: - c = grid_patches // r - if r * c + 1 == patch_count: - return r, c - # Fallback: treat as unsplit - return 0, 0 - - return rows, cols - def _create_single_image_prompt_string(self) -> str: return ( "" diff --git a/tests/test_colmodernvbert_preprocessing.py b/tests/test_colmodernvbert_preprocessing.py new file mode 100644 index 00000000..3b41fd10 --- /dev/null +++ b/tests/test_colmodernvbert_preprocessing.py @@ -0,0 +1,159 @@ +"""Offline preprocessing regressions: real transforms/tokenizer, no model weights.""" + +from unittest.mock import Mock + +import numpy as np +import pytest +from PIL import Image +from tokenizers import Tokenizer +from tokenizers.models import BPE +from tokenizers.pre_tokenizers import ByteLevel +from tokenizers.trainers import BpeTrainer + +from fastembed.image.transform.operators import Compose +from fastembed.late_interaction_multimodal.colmodernvbert import ColModernVBERT + + +def expected_prompt(rows: int, cols: int, image_seq_len: int = 2) -> str: + """Idefics3's row-major local crops followed by the global image.""" + image_tokens = "" * image_seq_len + fake = "" + local_rows = [ + "".join(f"{fake}{image_tokens}" for col in range(1, cols + 1)) + for row in range(1, rows + 1) + ] + local = "\n".join(local_rows) + "\n\n" if local_rows else "" + image = local + f"{fake}{image_tokens}{fake}" + return f"<|begin_of_text|>User:{image}Describe the image.\nAssistant:" + + +@pytest.fixture +def model() -> ColModernVBERT: + # Bypass downloading weights and constructing an ONNX session. Every preprocessing + # component below is real; only the final inference call is a recording test double. + model = object.__new__(ColModernVBERT) + tokenizer = Tokenizer(BPE(unk_token="[UNK]")) + tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False) + special_tokens = [ + "[UNK]", + "[PAD]", + "", + "", + "", + "<|begin_of_text|>", + "", + ] + [f"" for row in range(1, 5) for col in range(1, 5)] + tokenizer.train_from_iterator( + [expected_prompt(4, 4)], BpeTrainer(vocab_size=256, special_tokens=special_tokens) + ) + tokenizer.enable_padding(pad_id=tokenizer.token_to_id("[PAD]"), pad_token="[PAD]") + model.tokenizer = tokenizer + model.image_seq_len = 2 + model.model = Mock() + model.model.run.return_value = [np.empty((0,), dtype=np.float32)] + return model + + +def configure(model: ColModernVBERT, *, splitting: bool = True, resize: bool = False) -> None: + model.processor = Compose.from_config( + { + "image_processor_type": "Idefics3ImageProcessor", + "do_resize": resize, + "size": {"longest_edge": 16}, + "do_image_splitting": splitting, + "max_image_size": {"longest_edge": 4}, + "do_rescale": True, + "do_normalize": True, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + } + ) + + +def assert_prompt( + model: ColModernVBERT, onnx_input: dict, index: int, grid: tuple[int, int] +) -> None: + expected = model.tokenizer.encode(expected_prompt(*grid)).ids + active = onnx_input["attention_mask"][index].astype(bool) + actual = onnx_input["input_ids"][index][active] + # Assert actual special-token content as well as exact IDs and row separators. + actual_tokens = [model.tokenizer.id_to_token(int(token)) for token in actual] + expected_tokens = [model.tokenizer.id_to_token(token) for token in expected] + assert actual_tokens == expected_tokens + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + ("size", "grid"), + [ + ((4, 8), (2, 1)), + ((8, 4), (1, 2)), + ((16, 4), (1, 4)), + ((4, 16), (4, 1)), + ((8, 8), (2, 2)), + ((4, 4), (0, 0)), + ((5, 9), (3, 2)), + ], + ids=["portrait", "landscape", "wide", "tall", "square", "unsplit", "rounded"], +) +def test_image_prompt_retains_patch_grid(model, size, grid) -> None: + configure(model) + output = model.onnx_embed_image([Image.new("RGB", size, (32, 64, 128))]) + onnx_input = model.model.run.call_args.args[1] + assert_prompt(model, onnx_input, 0, grid) + assert set(onnx_input) == {"pixel_values", "input_ids", "attention_mask"} + assert output.metadata["image_grid"] == [grid] + count = grid[0] * grid[1] + 1 + assert output.metadata["patch_counts"] == [count] + assert onnx_input["pixel_values"].shape == (1, count, 3, 4, 4) + + +@pytest.mark.parametrize("direction", ["left", "right"]) +def test_mixed_image_batch_preserves_grids_and_padding(model, direction) -> None: + configure(model) + model.tokenizer.enable_padding(direction=direction, pad_id=1, pad_token="[PAD]") + sizes = [(4, 8), (8, 4), (16, 4), (8, 8), (4, 4)] + grids = [(2, 1), (1, 2), (1, 4), (2, 2), (0, 0)] + output = model.onnx_embed_image([Image.new("RGB", size, (32, 64, 128)) for size in sizes]) + onnx_input = model.model.run.call_args.args[1] + assert output.metadata["patch_counts"] == [3, 3, 5, 5, 1] + assert onnx_input["pixel_values"].shape == (5, 5, 3, 4, 4) + for i, grid in enumerate(grids): + assert_prompt(model, onnx_input, i, grid) + count = grid[0] * grid[1] + 1 + assert np.count_nonzero(onnx_input["pixel_values"][i, count:]) == 0 + mask = onnx_input["attention_mask"][i] + assert np.all(onnx_input["input_ids"][i][mask == 0] == 1) + padding_count = int(np.sum(mask == 0)) + if padding_count: + padded = mask[:padding_count] if direction == "left" else mask[-padding_count:] + assert not np.any(padded) + + +@pytest.mark.parametrize("size", [(4, 8), (8, 4)]) +def test_grid_is_retained_after_longest_edge_resize(model, size) -> None: + configure(model, resize=True) + model.onnx_embed_image([Image.new("RGB", size)]) + grid = (4, 2) if size[1] > size[0] else (2, 4) + assert_prompt(model, model.model.run.call_args.args[1], 0, grid) + + +def test_disabled_splitting_uses_only_global_prompt(model) -> None: + configure(model, splitting=False) + model.onnx_embed_image([Image.new("RGB", (4, 16))]) + onnx_input = model.model.run.call_args.args[1] + assert_prompt(model, onnx_input, 0, (0, 0)) + assert onnx_input["pixel_values"].shape == (1, 1, 3, 4, 4) + + +@pytest.mark.parametrize( + ("grid", "count"), + [(None, 3), ([], 3), ([(2, 2)], 3), ([(-1, -2)], 3), ([(0, 2)], 1), ([(2, 0)], 1)], +) +def test_missing_or_inconsistent_grid_is_rejected(model, grid, count) -> None: + inputs = { + "pixel_values": np.zeros((1, count, 3, 4, 4), dtype=np.float32), + "attention_mask": np.ones((1, count), dtype=np.int64), + } + with pytest.raises(ValueError, match="image patch grid"): + model._preprocess_onnx_image_input(inputs, image_grid=grid) diff --git a/tests/test_image_transform.py b/tests/test_image_transform.py index ceed9e35..b78be05f 100644 --- a/tests/test_image_transform.py +++ b/tests/test_image_transform.py @@ -109,3 +109,71 @@ def test_normalize_channel_count_mismatch_raises( image = np.zeros((3, 4, 4), dtype=np.float32) with pytest.raises(ValueError, match=expected): normalize(image, mean=mean, std=std) + + +@pytest.mark.parametrize("splitting", [False, True]) +@pytest.mark.parametrize("rescale_and_normalize", [False, True]) +def test_compose_grid_metadata_keeps_default_outputs(splitting, rescale_and_normalize) -> None: + processor = Compose.from_config( + { + "image_processor_type": "Idefics3ImageProcessor", + "do_resize": False, + "do_image_splitting": splitting, + "max_image_size": {"longest_edge": 4}, + "do_rescale": rescale_and_normalize, + "do_normalize": rescale_and_normalize, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + } + ) + images = [Image.new("RGB", size, (32, 64, 128)) for size in [(4, 8), (8, 4), (4, 4)]] + ordinary = processor(images) + metadata = {"existing": "preserved"} + with_metadata = processor(images, metadata=metadata) + + assert isinstance(ordinary, list) and isinstance(with_metadata, list) + assert metadata["image_grid"] == ([(2, 1), (1, 2), (0, 0)] if splitting else [(0, 0)] * 3) + assert metadata["existing"] == "preserved" + for expected, actual in zip(ordinary, with_metadata): + assert isinstance(expected, list) and isinstance(actual, list) + assert len(expected) == len(actual) + for expected_patch, actual_patch in zip(expected, actual): + assert isinstance(expected_patch, np.ndarray) and isinstance(actual_patch, np.ndarray) + np.testing.assert_array_equal(actual_patch, expected_patch) + + # Geometry belongs to each call, including repeated use of a metadata dictionary. + processor([images[-1]], metadata=metadata) + assert metadata["image_grid"] == [(0, 0)] + + +def test_splitter_grid_keeps_row_major_pixels_and_global_last() -> None: + from fastembed.image.transform.operators import ImageSplitter + + image = np.arange(3 * 8 * 12, dtype=np.float32).reshape(3, 8, 12) + splitter = ImageSplitter(max_size=4) + ordinary = splitter([image]) + patches, grids = splitter.split_with_grid([image]) + + assert grids == [(2, 3)] + assert isinstance(ordinary, list) and isinstance(ordinary[0], list) + assert len(patches[0]) == len(ordinary[0]) == 7 + for row in range(2): + for col in range(3): + index = row * 3 + col + np.testing.assert_array_equal( + patches[0][index], image[:, row * 4 : (row + 1) * 4, col * 4 : (col + 1) * 4] + ) + for actual, expected in zip(patches[0], ordinary[0]): + np.testing.assert_array_equal(actual, expected) + assert patches[0][-1].shape == (3, 4, 4) + + +def test_compose_metadata_does_not_change_flat_image_batches() -> None: + processor = Compose.from_config({"do_resize": False, "do_rescale": False}) + images = [Image.new("RGB", (4, 4))] + metadata = {} + actual = processor(images, metadata=metadata) + + assert isinstance(actual, list) and isinstance(actual[0], np.ndarray) + assert metadata == {"image_grid": [(0, 0)]} + np.testing.assert_array_equal(actual, processor(images))