From 145e081b35b5ea15adaa4392e0afc950c539545b Mon Sep 17 00:00:00 2001 From: serhiizghama Date: Wed, 19 Aug 2026 10:36:13 +0700 Subject: [PATCH 1/3] fix(image): normalize batched input along the channel axis normalize() advertises 4D (N, C, H, W) support via its num_channels branch and the channel-count validation, but the actual math used ((image.T - mean) / std).T. Transpose reverses every axis, so on 4D input the channels no longer line up with mean/std: it raises when N != C and silently normalizes along the batch axis when N == C. Reshape mean/std to broadcast on the real channel axis instead; the (C, H, W) path is unchanged. --- fastembed/image/transform/functional.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/fastembed/image/transform/functional.py b/fastembed/image/transform/functional.py index 9d9e2197..9d82e343 100644 --- a/fastembed/image/transform/functional.py +++ b/fastembed/image/transform/functional.py @@ -88,7 +88,15 @@ def normalize( std_arr = np.array(std_list, dtype=np.float32) - image_upd = ((image.T - mean_arr) / std_arr).T + # Broadcast along the channel axis: 0 for (C, H, W), 1 for a (N, C, H, W) batch. + # Transposing instead would reverse every axis and misalign the channels on 4D input. + channel_axis = 1 if image.ndim == 4 else 0 + broadcast_shape = [1] * image.ndim + broadcast_shape[channel_axis] = num_channels + mean_arr = mean_arr.reshape(broadcast_shape) + std_arr = std_arr.reshape(broadcast_shape) + + image_upd = (image - mean_arr) / std_arr return image_upd From 038eae81d74857c55c5e72ef8f9c3f10451ecf12 Mon Sep 17 00:00:00 2001 From: serhiizghama Date: Wed, 19 Aug 2026 10:52:30 +0700 Subject: [PATCH 2/3] test(image): cover channel-wise normalize for 3D and batched input --- tests/test_image_transform.py | 63 +++++++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 tests/test_image_transform.py diff --git a/tests/test_image_transform.py b/tests/test_image_transform.py new file mode 100644 index 00000000..6c5b0a2a --- /dev/null +++ b/tests/test_image_transform.py @@ -0,0 +1,63 @@ +import numpy as np +import pytest + +from fastembed.image.transform.functional import normalize + + +def _reference_normalize(image, mean, std): + """Channel-wise normalization with explicit broadcasting, used as ground truth.""" + channel_axis = 1 if image.ndim == 4 else 0 + shape = [1] * image.ndim + shape[channel_axis] = image.shape[channel_axis] + mean_arr = np.asarray(mean, dtype=np.float32).reshape(shape) + std_arr = np.asarray(std, dtype=np.float32).reshape(shape) + return (image.astype(np.float32) - mean_arr) / std_arr + + +def test_normalize_chw_matches_channel_wise(): + rng = np.random.default_rng(0) + image = rng.random((3, 5, 7)).astype(np.float32) + mean, std = [0.1, 0.2, 0.3], [0.5, 0.6, 0.7] + + result = normalize(image, mean=mean, std=std) + + assert np.allclose(result, _reference_normalize(image, mean, std), atol=1e-6) + + +def test_normalize_scalar_mean_std(): + rng = np.random.default_rng(1) + image = rng.random((3, 4, 4)).astype(np.float32) + + result = normalize(image, mean=0.5, std=0.25) + + assert np.allclose(result, (image - 0.5) / 0.25, atol=1e-6) + + +def test_normalize_batched_input_normalizes_per_channel(): + # (N, C, H, W): every channel c is filled with the constant c, so subtracting + # mean == c and dividing by 1 must yield all zeros regardless of batch size. + image = np.zeros((3, 3, 2, 2), dtype=np.float32) + for c in range(3): + image[:, c] = c + + result = normalize(image, mean=[0.0, 1.0, 2.0], std=[1.0, 1.0, 1.0]) + + assert np.allclose(result, 0.0) + + +def test_normalize_batched_input_when_batch_differs_from_channels(): + # N != C used to raise because transposing reversed every axis. + rng = np.random.default_rng(2) + image = rng.random((2, 3, 4, 4)).astype(np.float32) + mean, std = [0.1, 0.2, 0.3], [0.5, 0.6, 0.7] + + result = normalize(image, mean=mean, std=std) + + assert result.shape == image.shape + assert np.allclose(result, _reference_normalize(image, mean, std), atol=1e-6) + + +def test_normalize_channel_count_mismatch_raises(): + image = np.zeros((3, 4, 4), dtype=np.float32) + with pytest.raises(ValueError): + normalize(image, mean=[0.1, 0.2], std=[1.0, 1.0, 1.0]) From 72726f1cfe24ac07e877429260ad90f818a35bd1 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Tue, 22 Sep 2026 14:50:06 +0700 Subject: [PATCH 3/3] refactor --- fastembed/image/transform/functional.py | 22 +++---- tests/test_image_transform.py | 84 +++++++++++++------------ 2 files changed, 54 insertions(+), 52 deletions(-) diff --git a/fastembed/image/transform/functional.py b/fastembed/image/transform/functional.py index 7ad3cc24..bfee639c 100644 --- a/fastembed/image/transform/functional.py +++ b/fastembed/image/transform/functional.py @@ -65,7 +65,13 @@ def normalize( mean: float | list[float], std: float | list[float], ) -> NumpyArray: - num_channels = image.shape[1] if len(image.shape) == 4 else image.shape[0] + if image.ndim < 3: + raise ValueError(f"image must be (C, H, W) or (N, C, H, W), got shape {image.shape}") + + # Channels sit on the third axis from the end, which covers (C, H, W) and + # (N, C, H, W) alike. Transposing instead reversed every axis, which put the + # batch dimension where the channels were meant to be. + num_channels = image.shape[-3] if not np.issubdtype(image.dtype, np.floating): image = image.astype(np.float32) @@ -78,7 +84,9 @@ def normalize( f"{len(mean_list)}" ) - mean_arr = np.array(mean_list, dtype=np.float32) + # (C, 1, 1) lines the channels up with the trailing (C, H, W) axes under numpy + # broadcasting, whatever batch dimensions lead them. + mean_arr = np.array(mean_list, dtype=np.float32).reshape(-1, 1, 1) std_list = std if isinstance(std, list) else [std] * num_channels if len(std_list) != num_channels: @@ -86,15 +94,7 @@ def normalize( f"std must have the same number of channels as the image, image has {num_channels} channels, got {len(std_list)}" ) - std_arr = np.array(std_list, dtype=np.float32) - - # Broadcast along the channel axis: 0 for (C, H, W), 1 for a (N, C, H, W) batch. - # Transposing instead would reverse every axis and misalign the channels on 4D input. - channel_axis = 1 if image.ndim == 4 else 0 - broadcast_shape = [1] * image.ndim - broadcast_shape[channel_axis] = num_channels - mean_arr = mean_arr.reshape(broadcast_shape) - std_arr = std_arr.reshape(broadcast_shape) + std_arr = np.array(std_list, dtype=np.float32).reshape(-1, 1, 1) image_upd = (image - mean_arr) / std_arr return image_upd diff --git a/tests/test_image_transform.py b/tests/test_image_transform.py index 4f2d140e..d152719c 100644 --- a/tests/test_image_transform.py +++ b/tests/test_image_transform.py @@ -29,60 +29,62 @@ def test_resize_int_keeps_shortest_edge_behaviour() -> None: assert resize(portrait, size=100).size == (100, 200) -def _reference_normalize(image, mean, std): - """Channel-wise normalization with explicit broadcasting, used as ground truth.""" - channel_axis = 1 if image.ndim == 4 else 0 - shape = [1] * image.ndim - shape[channel_axis] = image.shape[channel_axis] - mean_arr = np.asarray(mean, dtype=np.float32).reshape(shape) - std_arr = np.asarray(std, dtype=np.float32).reshape(shape) - return (image.astype(np.float32) - mean_arr) / std_arr - - -def test_normalize_chw_matches_channel_wise(): +@pytest.mark.parametrize( + ("mean", "std"), + [ + ([0.1, 0.2, 0.3], [0.5, 0.6, 0.7]), # per-channel, as every model config gives it + (0.5, 0.25), # scalar, expanded to one value per channel + ], +) +def test_normalize_chw_is_channel_wise( + mean: list[float] | float, std: list[float] | float +) -> None: + """Each channel must be normalized by its own mean/std, not by any other axis.""" rng = np.random.default_rng(0) image = rng.random((3, 5, 7)).astype(np.float32) - mean, std = [0.1, 0.2, 0.3], [0.5, 0.6, 0.7] + means = mean if isinstance(mean, list) else [mean] * 3 + stds = std if isinstance(std, list) else [std] * 3 result = normalize(image, mean=mean, std=std) - assert np.allclose(result, _reference_normalize(image, mean, std), atol=1e-6) - - -def test_normalize_scalar_mean_std(): - rng = np.random.default_rng(1) - image = rng.random((3, 4, 4)).astype(np.float32) - - result = normalize(image, mean=0.5, std=0.25) - - assert np.allclose(result, (image - 0.5) / 0.25, atol=1e-6) - - -def test_normalize_batched_input_normalizes_per_channel(): - # (N, C, H, W): every channel c is filled with the constant c, so subtracting - # mean == c and dividing by 1 must yield all zeros regardless of batch size. - image = np.zeros((3, 3, 2, 2), dtype=np.float32) for c in range(3): - image[:, c] = c - - result = normalize(image, mean=[0.0, 1.0, 2.0], std=[1.0, 1.0, 1.0]) + assert np.allclose(result[c], (image[c] - means[c]) / stds[c], atol=1e-6) - assert np.allclose(result, 0.0) +@pytest.mark.parametrize("batch_size", [2, 3]) +def test_normalize_batched_matches_per_image(batch_size: int) -> None: + """A batch must give exactly what the (C, H, W) path gives image by image. -def test_normalize_batched_input_when_batch_differs_from_channels(): - # N != C used to raise because transposing reversed every axis. + batch_size 2 used to raise, since transposing reversed every axis; batch_size 3 + matched the channel count and silently normalized along the batch axis instead. + """ rng = np.random.default_rng(2) - image = rng.random((2, 3, 4, 4)).astype(np.float32) + batch = rng.random((batch_size, 3, 4, 4)).astype(np.float32) mean, std = [0.1, 0.2, 0.3], [0.5, 0.6, 0.7] - result = normalize(image, mean=mean, std=std) + result = normalize(batch, mean=mean, std=std) + + per_image = np.stack([normalize(image, mean=mean, std=std) for image in batch]) + assert result.shape == batch.shape + assert np.array_equal(result, per_image) - assert result.shape == image.shape - assert np.allclose(result, _reference_normalize(image, mean, std), atol=1e-6) +def test_normalize_rejects_input_without_a_channel_axis() -> None: + """Every pipeline runs ConvertToRGB first, so normalize only ever sees (C, H, W).""" + with pytest.raises(ValueError, match=r"must be \(C, H, W\)"): + normalize(np.zeros((4, 6), dtype=np.float32), mean=0.5, std=0.25) -def test_normalize_channel_count_mismatch_raises(): + +@pytest.mark.parametrize( + ("mean", "std", "expected"), + [ + ([0.1, 0.2], [1.0, 1.0, 1.0], "mean must"), + ([0.1, 0.2, 0.3], [1.0, 1.0], "std must"), + ], +) +def test_normalize_channel_count_mismatch_raises( + mean: list[float], std: list[float], expected: str +) -> None: image = np.zeros((3, 4, 4), dtype=np.float32) - with pytest.raises(ValueError): - normalize(image, mean=[0.1, 0.2], std=[1.0, 1.0, 1.0]) + with pytest.raises(ValueError, match=expected): + normalize(image, mean=mean, std=std)