add dataset brightness variance

granatpy/metrics.py

@@ -9,7 +9,7 @@
 import lpips
 
 
-__all__ = ["compute_all_metrics", "compare_images", "image_entropy", "dataset_entropy", "compute_lpips"]
+__all__ = ["compute_all_metrics", "compare_images", "image_entropy", "dataset_entropy", "brightness_var", "compute_lpips"]
 
 
 def _compute_nf(image: np.ndarray) -> float:
@@ -197,6 +197,43 @@ def dataset_entropy(image_paths: List[str], size: tuple, bins=256) -> float:
 
     return float(ent_map.mean())
 
+def brightness_var(image_paths: List[str], size: tuple) -> float:
+    """
+    Compute the variance of brightness across a dataset of images.
+
+    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 brightness variance.
+
+    Returns:
+        Brightness variance as float, computed as the mean variance 
+        across all pixel positions in the dataset.
+    """
+    imgs = []
+
+    for path in image_paths:
+        try:
+            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"Skip {path}: {e}")
+
+    data = np.stack(imgs, axis=0)
+    N, H, W = data.shape
+    var_map = np.zeros((H, W), dtype=np.float64)
+
+    for i in range(H):
+        for j in range(W):
+            values = data[:, i, j]
+            var_map[i, j] = np.var(values)
+
+    return float(var_map.mean())
+
 def compute_lpips(real: str, synthetic: str, use_gpu: bool = True):
     """
     Compute the Learned Perceptual Image Patch Similarity (LPIPS) metric for two images.

tests/test_metrics.py

@@ -3,7 +3,7 @@
 from PIL import Image
 from skimage.measure import shannon_entropy
 from skimage import img_as_ubyte
-from granatpy.metrics import compute_all_metrics, image_entropy, dataset_entropy, load_image
+from granatpy.metrics import compute_all_metrics, image_entropy, dataset_entropy, brightness_var, load_image
 
 def test_compute_all_metrics():
     # Create two identical arrays
@@ -135,4 +135,55 @@ def test_dataset_entropy(tmp_path):
     # The output should be identical to the one without invalid path since it skips invalid
     assert entropy_val_with_invalid == entropy_val
 
+def test_brightness_var(tmp_path):
+    # Prepare minimal RGB, RGBA, and 16-bit image paths
+    rgb_path = tmp_path / "test_rgb.png"
+    rgba_path = tmp_path / "test_rgba.png"
+    gray16_path = tmp_path / "test_gray16.png"
+
+    # Create and save minimal images (e.g. 2x2 pixels)
+    # 1. RGB (2x2x3)
+    rgb_data = np.array([
+        [[255, 0, 0], [0, 255, 0]],
+        [[0, 0, 255], [128, 128, 128]]
+    ], dtype=np.uint8)
+    Image.fromarray(rgb_data).save(rgb_path)
+
+    # 2. RGBA (2x2x4)
+    rgba_data = np.array([
+        [[255, 0, 0, 255], [0, 255, 0, 255]],
+        [[0, 0, 255, 255], [128, 128, 128, 0]]
+    ], dtype=np.uint8)
+    Image.fromarray(rgba_data).save(rgba_path)
 
+    # 3. 16-bit grayscale (2x2)
+    gray16_data = np.array([
+        [0, 65535],
+        [32768, 16384]
+    ], dtype=np.uint16)
+    Image.fromarray(gray16_data).save(gray16_path)
+
+    # Calculate dataset-wide brightness variance using all 3 paths
+    paths = [str(rgb_path), str(rgba_path), str(gray16_path)]
+    brightness_val = brightness_var(image_paths=paths, size=(2, 2))
+
+    assert isinstance(brightness_val, float)
+    assert brightness_val > 0.0
+
+    # Pre-calculated Brightness Variance (base 2) for pooled pixels under current Pillow logic:
+    # - test_rgb: loaded shape (2,2,3) -> {0: 6, 128: 3, 255: 3}
+    # - test_rgba: loaded shape (2,2,3) (converted to RGB, alpha discarded) -> {0: 6, 128: 3, 255: 3}
+    # - 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)) = 4667.000000000001
+    assert np.isclose(brightness_val, 4667.000000000001)
+
+    # Test handling of invalid paths / error handling in brightness_var
+    invalid_path = tmp_path / "non_existent.png"
+    paths_with_invalid = paths + [str(invalid_path)]
+    brightness_val_with_invalid = brightness_var(image_paths=paths_with_invalid, size=(2, 2))
+
+    assert isinstance(brightness_val_with_invalid, float)
+    # The output should be identical to the one without invalid path since it skips invalid
+    assert brightness_val_with_invalid == brightness_val