diff --git a/fastembed/image/transform/functional.py b/fastembed/image/transform/functional.py index af28fb0b..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,9 +94,9 @@ 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) + std_arr = np.array(std_list, dtype=np.float32).reshape(-1, 1, 1) - image_upd = ((image.T - mean_arr) / std_arr).T + 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 6a203238..d152719c 100644 --- a/tests/test_image_transform.py +++ b/tests/test_image_transform.py @@ -1,7 +1,8 @@ +import numpy as np import pytest from PIL import Image -from fastembed.image.transform.functional import resize +from fastembed.image.transform.functional import normalize, resize @pytest.mark.parametrize( @@ -26,3 +27,64 @@ def test_resize_int_keeps_shortest_edge_behaviour() -> None: # size sets the shortest edge, and the aspect ratio is preserved. assert resize(landscape, size=100).size == (200, 100) assert resize(portrait, size=100).size == (100, 200) + + +@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) + 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) + + for c in range(3): + assert np.allclose(result[c], (image[c] - means[c]) / stds[c], atol=1e-6) + + +@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. + + 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) + 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(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) + + +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) + + +@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, match=expected): + normalize(image, mean=mean, std=std)