update entropy function

granatpy/metrics.py

@@ -148,43 +148,50 @@ def image_entropy(img: np.ndarray) -> float:
 
     return float(entropy(hist, base=2))
 
- 
-def dataset_entropy(image_paths: List[str]) -> float:
+
+def dataset_entropy(image_paths: List[str], size: tuple, bins=256) -> float:
     """
-    Compute the Shannon entropy over an entire dataset of images.
+    Compute the Shannon entropy for an entire dataset of images.
 
-    All pixel values across every image are pooled into a single histogram before
-    computing entropy, so the result reflects diversity across the full dataset
-    rather than an average of per-image entropies.
+    Pixel values across every image for each pixel position are pooled together to
+    compute entropy, then an average is taken across all pixel positions. This ensures 
+    that the entropy reflects the diversity of pixel values across the entire dataset, 
+    rather than just per-image entropy.
+
 
     Args:
         image_paths: List of paths to image files. Files that cannot be read are
                      skipped with a printed warning.
+        size: The size to which all images will be resized before computing entropy.
+        bins: The number of bins for the histogram.
 
     Returns:
-        Shannon entropy as float, computed over the pooled pixel
-        intensity distribution of the entire dataset.
+        Shannon entropy as float, computed as the mean entropy 
+        across all pixel positions in the dataset.
     """
-    all_pixels = []
- 
+    imgs = []
+
     for path in image_paths:
         try:
-            img = load_image(path)
-            img = img_as_ubyte(img)
-            all_pixels.append(img.flatten())
+            img = Image.open(path).convert("L")
+            if size is not None:
+                img = img.resize(size)
+
+            imgs.append(np.array(img, dtype=np.uint8))
+
         except Exception as e:
-            print(f"Error: {path} -> {e}")
- 
-    all_pixels = np.concatenate(all_pixels)
-
-    hist, _ = np.histogram(
-        all_pixels,
-        bins=256,
-        range=(0, 255),
-        density=True
-    )
-
-    hist = hist[hist > 0]
-
-    return float(entropy(hist, base=2))
-
+            print(f"Skip {path}: {e}")
+
+    data = np.stack(imgs, axis=0)
+    N, H, W = data.shape
+    ent_map = np.zeros((H, W), dtype=np.float64)
+    bin_edges = np.linspace(0, 256, bins + 1)
+
+    for i in range(H):
+        for j in range(W):
+            values = data[:, i, j]
+            hist, _ = np.histogram(values, bins=bin_edges, density=True)
+            hist = hist[hist > 0]
+            ent_map[i, j] = entropy(hist, base=2)
+
+    return float(ent_map.mean())

tests/test_metrics.py

@@ -112,7 +112,7 @@ def test_dataset_entropy(tmp_path):
 
     # Calculate dataset-wide entropy using all 3 paths
     paths = [str(rgb_path), str(rgba_path), str(gray16_path)]
-    entropy_val = dataset_entropy(paths)
+    entropy_val = dataset_entropy(image_paths=paths, size=(2, 2), bins=256)
 
     assert isinstance(entropy_val, float)
     assert entropy_val > 0.0
@@ -123,13 +123,13 @@ def test_dataset_entropy(tmp_path):
     # - test_gray16: loaded shape (2,2,3) (converted to RGB, Pillow clips values > 255 to 255) -> {0: 3, 255: 9}
     # Total counts: {0: 15, 128: 6, 255: 15} out of 36 pixels.
     # Probabilities: {0: 15/36, 128: 6/36, 255: 15/36} = {0: 5/12, 128: 1/6, 255: 5/12}
-    # - (2 * 5/12 * log2(5/12) + 1/6 * log2(1/6)) = 1.4833557549816876
-    assert np.isclose(entropy_val, 1.4833557549816876)
+    # - (2 * 5/12 * log2(5/12) + 1/6 * log2(1/6)) = 0.9182958340544894
+    assert np.isclose(entropy_val, 0.9182958340544894)
 
     # Test handling of invalid paths / error handling in dataset_entropy
     invalid_path = tmp_path / "non_existent.png"
     paths_with_invalid = paths + [str(invalid_path)]
-    entropy_val_with_invalid = dataset_entropy(paths_with_invalid)
+    entropy_val_with_invalid = dataset_entropy(image_paths=paths_with_invalid, size=(2, 2), bins=256)
 
     assert isinstance(entropy_val_with_invalid, float)
     # The output should be identical to the one without invalid path since it skips invalid