diff --git a/granatpy/metrics.py b/granatpy/metrics.py index 0aa329d..6ee9f8a 100644 --- a/granatpy/metrics.py +++ b/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. diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 103cae8..163a473 100644 --- a/tests/test_metrics.py +++ b/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