From 145e081b35b5ea15adaa4392e0afc950c539545b Mon Sep 17 00:00:00 2001
From: serhiizghama <zmrser@gmail.com>
Date: Wed, 19 Aug 2026 10:36:13 +0700
Subject: [PATCH 1/2] 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 9d9e2197e..9d82e3432 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 <zmrser@gmail.com>
Date: Wed, 19 Aug 2026 10:52:30 +0700
Subject: [PATCH 2/2] 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 000000000..6c5b0a2a6
--- /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])