fix(image): normalize batched (N, C, H, W) input along the channel axis

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,17 +84,19 @@ 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:
         raise ValueError(
             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
 
 

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)