diff --git a/deepinfer/ve/processor/image_processor.py b/deepinfer/ve/processor/image_processor.py index 8ef6cea..cf2c662 100644 --- a/deepinfer/ve/processor/image_processor.py +++ b/deepinfer/ve/processor/image_processor.py @@ -152,6 +152,16 @@ def safe_resize( return layout +def exceeds_max_wh_ratio(width: int, height: int, max_wh_ratio: int | None) -> bool: + """该图是否超出宽高比上限(即其宽度会被裁掉)。 + + compute_image_layout 会把这类图的宽度夹到 height*max_wh_ratio,resize 时要把 + 裁掉的宽度压回来(stretch)而不是等比内接(contain)。判定必须严格一致:恰好 + 等于上限的图没有被裁剪,应当保持长宽比,否则它会被无端拉伸变形。 + """ + return max_wh_ratio is not None and width > max_wh_ratio * height + + def compute_image_layout( width: int, height: int, @@ -163,7 +173,7 @@ def compute_image_layout( ) -> ImageLayout: """把原图尺寸规划成 patch 对齐的 best_width/best_height 及 LLM token 数。""" # 超宽图先把宽度夹到 height*max_wh_ratio;该图的 resize 相应走 stretch 而非等比内接 - if max_wh_ratio is not None and width > height * max_wh_ratio: + if exceeds_max_wh_ratio(width, height, max_wh_ratio): width = height * max_wh_ratio current_pixels = width * height @@ -264,8 +274,13 @@ class ImageProcessorAnyRes(BaseProcessor): self.max_wh_ratio = max_wh_ratio def _is_stretched(self, width: int, height: int) -> bool: - """该图是否触发了 compute_image_layout 的超宽裁剪,裁掉的宽度靠 stretch resize 压回来。""" - return self.max_wh_ratio is not None and width >= self.max_wh_ratio * height + """该图是否触发了 compute_image_layout 的超宽裁剪,裁掉的宽度靠 stretch resize 压回来。 + + 判定与 compute_image_layout 的裁剪条件共用 exceeds_max_wh_ratio:恰好等于上限 + 的图没有被裁剪,必须走等比内接,否则会被拉伸变形(patch 对齐带来的长宽比误差 + 全都压到图上,而只差一个像素的邻居却不会)。 + """ + return exceeds_max_wh_ratio(width, height, self.max_wh_ratio) def _prepare_image(self, image: Image.Image, target: tuple[int, int], *, stretch: bool = False) -> ResizedImage: """CPU resize 到 target 尺寸,保留 uint8 像素(pad / normalize 由 GPU 变换完成)。""" diff --git a/tests/ve/test_image_stretch_boundary.py b/tests/ve/test_image_stretch_boundary.py new file mode 100644 index 0000000..ed4ac17 --- /dev/null +++ b/tests/ve/test_image_stretch_boundary.py @@ -0,0 +1,112 @@ +"""CPU-only regression test for the ultra-wide resize boundary. + +``compute_image_layout`` clamps images that are wider than ``height * max_wh_ratio`` +and the CPU resize then has to squeeze that removed width back (stretch) instead of +keeping the aspect ratio (contain). The predicate that picks between the two has to +agree with the clamp condition exactly: an image whose ratio sits *at* the limit is +not clamped, so it must keep its aspect ratio. Otherwise it is distorted by the +patch-alignment rounding, while a one-pixel narrower neighbour is not. + +Only CPU geometry is exercised: no checkpoint, no accelerator. +""" + +from __future__ import annotations + +import pytest +import torch +from PIL import Image + +from deepinfer.ve.processor.image_processor import ( + BasicImageTransform, + ImageLayout, + ImageProcessorAnyRes, + compute_image_layout, + exceeds_max_wh_ratio, +) + +PATCH_SIZE = 14 +DOWNSAMPLE_RATIO = 3 +MAX_N_TOKEN = 4096 +MAX_WH_RATIOS = [4, 8, 16] +HEIGHTS = [100, 250] + + +def _processor(max_wh_ratio: int | None) -> ImageProcessorAnyRes: + return ImageProcessorAnyRes( + image_transform=BasicImageTransform(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), + patch_size=PATCH_SIZE, + downsample_ratio=DOWNSAMPLE_RATIO, + dtype=torch.bfloat16, + max_n_token=MAX_N_TOKEN, + max_wh_ratio=max_wh_ratio, + ) + + +def _layout(width: int, height: int, max_wh_ratio: int | None = None) -> ImageLayout: + return compute_image_layout( + width, + height, + patch_size=PATCH_SIZE, + downsample_ratio=DOWNSAMPLE_RATIO, + max_n_token=MAX_N_TOKEN, + max_wh_ratio=max_wh_ratio, + ) + + +@pytest.mark.parametrize("max_wh_ratio", MAX_WH_RATIOS) +@pytest.mark.parametrize("height", HEIGHTS) +def test_stretch_predicate_mirrors_the_layout_clamp(max_wh_ratio: int, height: int) -> None: + """The resize must stretch exactly the images whose width the layout clamped.""" + processor = _processor(max_wh_ratio) + limit = height * max_wh_ratio + + # The layout clamps the width to `height * max_wh_ratio`, so a wider image ends up + # with the very same geometry as the clamped one. + assert _layout(limit + 4 * PATCH_SIZE, height, max_wh_ratio) == _layout(limit, height) + + # Only images strictly beyond the limit are clamped, so only they may be stretched. + assert processor._is_stretched(limit - 1, height) is False + assert processor._is_stretched(limit, height) is False + assert processor._is_stretched(limit + 1, height) is True + + +@pytest.mark.parametrize("max_wh_ratio", MAX_WH_RATIOS) +@pytest.mark.parametrize("height", HEIGHTS) +def test_image_at_the_ratio_limit_keeps_its_aspect_ratio(max_wh_ratio: int, height: int) -> None: + """An image exactly at the limit is not clamped, so it must not be stretched.""" + processor = _processor(max_wh_ratio) + limit = height * max_wh_ratio + + layout = _layout(limit, height, max_wh_ratio) + resized = processor._prepare_image( + Image.new("RGB", (limit, height)), + (layout.best_width, layout.best_height), + stretch=processor._is_stretched(limit, height), + ) + width, out_height = resized.pixels.shape[1], resized.pixels.shape[0] + # Contain keeps the aspect ratio (the remainder is padded by the GPU transform). + assert abs(width / out_height - limit / height) / (limit / height) < 0.02 + + +@pytest.mark.parametrize("max_wh_ratio", MAX_WH_RATIOS) +@pytest.mark.parametrize("height", HEIGHTS) +def test_image_beyond_the_ratio_limit_is_squeezed_into_the_canvas(max_wh_ratio: int, height: int) -> None: + """A clamped image keeps the layout canvas and is squeezed into it on purpose.""" + processor = _processor(max_wh_ratio) + limit = height * max_wh_ratio + width = limit + 4 * PATCH_SIZE + + layout = _layout(width, height, max_wh_ratio) + resized = processor._prepare_image( + Image.new("RGB", (width, height)), + (layout.best_width, layout.best_height), + stretch=processor._is_stretched(width, height), + ) + assert (resized.pixels.shape[1], resized.pixels.shape[0]) == (layout.best_width, layout.best_height) + + +def test_no_ratio_limit_never_stretches() -> None: + """Without a limit there is no clamp, so the resize always keeps the aspect ratio.""" + processor = _processor(None) + assert processor._is_stretched(100_000, 1) is False + assert exceeds_max_wh_ratio(100_000, 1, None) is False