feat: evaluation routines

Author: anna-grim

src/aind_exaspim_image_compression/evaluate.py

@@ -0,0 +1,100 @@
+"""
+Created on Tue June 3 12:00:00 2025
+
+@author: Anna Grim
+@email: anna.grim@alleninstitute.org
+
+Code for evaluating a denoising model on an image dataset.
+
+"""
+
+from numcodecs import blosc
+from tqdm import tqdm
+
+import numpy as np
+import os
+import pandas as pd
+import torch
+
+from aind_exaspim_image_compression.inference import predict, stitch
+from aind_exaspim_image_compression.utils import img_util, util
+
+
+class Evaluator:
+    def __init__(self, img_paths, model, output_dir):
+        # Instance attributes
+        self.codec = blosc.Blosc(cname="zstd", clevel=6, shuffle=blosc.SHUFFLE)
+        self.img_paths = img_paths
+        self.model = model
+        self.model.eval().to("cuda")
+
+        # Initialize output directory
+        self.output_dir = output_dir
+        util.mkdir(output_dir, delete=True)
+
+        # Load images
+        self.load_images()
+
+    def load_images(self):
+        # Initialize MIPs directory
+        noise_dir = os.path.join(self.output_dir, "noise_mips")
+        util.mkdir(noise_dir)
+
+        # Read images
+        self.noise_imgs = dict()
+        self.noise_cratios = dict()
+        for img_path in self.img_paths:
+            block_id = self.find_img_name(img_path)
+            img = img_util.read(img_path)
+            self.noise_cratios[block_id] = img_util.compute_cratio(
+                img, self.codec
+            )
+            self.noise_imgs[block_id] = img
+
+            output_path = os.path.join(noise_dir, block_id)
+            img_util.plot_mips(img[0, 0, ...], output_path=output_path)
+            break
+
+    # --- Main ---
+    def run(self, model_path):
+        # Initializations
+        self.model.load_state_dict(torch.load(model_path))
+        model_name = os.path.basename(model_path)
+        results_dir = os.path.join(self.output_dir, model_name)
+        util.mkdir(results_dir)
+
+        # Generate prediction
+        rows = list(self.noise_imgs.keys())
+        df = pd.DataFrame(index=rows, columns=["cratio", "ssim"])
+        desc = "Denoise Blocks"
+        for block_id, noise in tqdm(self.noise_imgs.items(), desc=desc):
+            # Run model
+            coords, preds = predict(noise, self.model, verbose=False)
+            denoised = stitch(noise, coords, preds)
+            denoised = denoised.astype(np.uint16)
+
+            # Compute metrics
+            df.loc[block_id, "cratio"] = img_util.compute_cratio(
+                denoised, self.codec
+            )
+            df.loc[block_id, "ssim"] = img_util.compute_ssim3D(
+                noise[0, 0, ...],
+                denoised[0, 0, ...],
+                data_range=np.percentile(noise, 99.9),
+            )
+
+            # Save MIPs
+            output_path = os.path.join(results_dir, block_id)
+            img_util.plot_mips(denoised[0, 0, ...], output_path=output_path)
+
+        # Save metrics
+        path = os.path.join(results_dir, "results.csv")
+        df.to_csv(path, index=True)
+        return model_name, df
+
+    # --- Helpers ---
+    def find_img_name(self, img_path):
+        for part in img_path.split("/"):
+            if "block_" in part:
+                return part
+        raise Exception(f"Block ID not found in {img_path}")

src/aind_exaspim_image_compression/inference.py

@@ -11,25 +11,27 @@
 import numpy as np
 import torch
 
-from aind_exaspim_image_compression.utils import img_util, util
+from aind_exaspim_image_compression.utils import img_util
 
 
-def predict(img, model, batch_size=32, patch_size=64, overlap=16):
+def predict(
+    img, model, batch_size=32, patch_size=64, overlap=16, verbose=True
+):
     # Initializations
     batch_coords, batch_inputs, mn_mx = list(), list(), list()
     coords = generate_coords(img, patch_size, overlap)
 
     # Main
-    pbar = tqdm(total=len(coords), desc="Denoise")
+    pbar = tqdm(total=len(coords), desc="Denoise") if verbose else None
     preds = list()
     for idx, (i, j, k) in enumerate(coords):
         # Get end coord
-        i_end = min(i + patch_size, img.shape[0])
-        j_end = min(j + patch_size, img.shape[1])
-        k_end = min(k + patch_size, img.shape[2])
+        i_end = min(i + patch_size, img.shape[2])
+        j_end = min(j + patch_size, img.shape[3])
+        k_end = min(k + patch_size, img.shape[4])
 
         # Get patch
-        patch = img[i:i_end, j:j_end, k:k_end]
+        patch = img[0, 0, i:i_end, j:j_end, k:k_end]
         mn, mx = np.percentile(patch, 5), np.percentile(patch, 99.9)
         patch = (patch - mn) / mx
         mn_mx.append((mn, mx))
@@ -42,7 +44,7 @@ def predict(img, model, batch_size=32, patch_size=64, overlap=16):
         # If batch is full or it's the last patch
         if len(batch_inputs) == batch_size or idx == len(coords) - 1:
             # Run model
-            input_tensor =to_tensor(np.stack(batch_inputs))
+            input_tensor = to_tensor(np.stack(batch_inputs))
             with torch.no_grad():
                 output_tensor = model(input_tensor)
 
@@ -51,8 +53,8 @@ def predict(img, model, batch_size=32, patch_size=64, overlap=16):
             for cnt in range(output_tensor.shape[0]):
                 mn, mx = mn_mx[cnt]
                 patch = np.array(output_tensor[cnt, 0, ...])
-                preds.append(patch * mx + mn)
-                pbar.update(1)
+                preds.append(np.maximum(patch * mx + mn, 0))
+                pbar.update(1) if verbose else None
 
             batch_coords.clear()
             batch_inputs.clear()
@@ -63,14 +65,15 @@ def predict(img, model, batch_size=32, patch_size=64, overlap=16):
 def stitch(img, coords, preds, patch_size=64, trim=5):
     denoised_accum = np.zeros_like(img, dtype=np.float32)
     weight_map = np.zeros_like(img, dtype=np.float32)
-
     for (i, j, k), pred in zip(coords, preds):
         # Determine how much to trim
         trim_start = trim
         trim_end = patch_size - trim
 
         # Trim prediction
-        pred_trimmed = pred[trim_start:trim_end, trim_start:trim_end, trim_start:trim_end]
+        pred_trimmed = pred[
+            trim_start:trim_end, trim_start:trim_end, trim_start:trim_end
+        ]
 
         # Adjust insertion indices
         i_start = i + trim
@@ -82,23 +85,29 @@ def stitch(img, coords, preds, patch_size=64, trim=5):
         k_end = k_start + pred_trimmed.shape[2]
 
         # Clip to image bounds (for safety)
-        i_end = min(i_end, img.shape[0])
-        j_end = min(j_end, img.shape[1])
-        k_end = min(k_end, img.shape[2])
+        i_end = min(i_end, img.shape[2])
+        j_end = min(j_end, img.shape[3])
+        k_end = min(k_end, img.shape[4])
 
         i_start = max(i_start, 0)
         j_start = max(j_start, 0)
         k_start = max(k_start, 0)
 
-        denoised_accum[i_start:i_end, j_start:j_end, k_start:k_end] += pred_trimmed[:i_end - i_start, :j_end - j_start, :k_end - k_start]
-        weight_map[i_start:i_end, j_start:j_end, k_start:k_end] += 1
+        denoised_accum[
+            0, 0, i_start:i_end, j_start:j_end, k_start:k_end
+        ] += pred_trimmed[
+            : i_end - i_start, : j_end - j_start, : k_end - k_start
+        ]
+        weight_map[0, 0, i_start:i_end, j_start:j_end, k_start:k_end] += 1
 
     # Average accumulated
     weight_map[weight_map == 0] = 1
     denoised = denoised_accum / weight_map
 
     # Fill boundary trim
-    fill_value = np.percentile(denoised[trim:-trim, trim:-trim, trim:-trim], 10)
+    fill_value = np.percentile(
+        denoised[..., trim:-trim, trim:-trim, trim:-trim], 10
+    )
     return img_util.fill_boundary(denoised, trim, fill_value)
 
 
@@ -109,18 +118,19 @@ def add_padding(patch, patch_size):
         (0, patch_size - patch.shape[1]),
         (0, patch_size - patch.shape[2]),
     ]
-    return np.pad(patch, pad_width, mode='constant', constant_values=0)
+    return np.pad(patch, pad_width, mode="constant", constant_values=0)
 
 
 def generate_coords(img, patch_size, overlap):
     coords = list()
     stride = patch_size - overlap
-    for i in range(0, img.shape[0] - patch_size + stride, stride):
-        for j in range(0, img.shape[1] - patch_size + stride, stride):
-            for k in range(0, img.shape[2] - patch_size + stride, stride):
+    for i in range(0, img.shape[2] - patch_size + stride, stride):
+        for j in range(0, img.shape[3] - patch_size + stride, stride):
+            for k in range(0, img.shape[4] - patch_size + stride, stride):
                 coords.append((i, j, k))
     return coords
 
 
 def to_tensor(arr):
-    return torch.tensor(arr[:, np.newaxis, ...]).to("cuda")
+    dtype = torch.float
+    return torch.tensor(arr[:, np.newaxis, ...]).to("cuda", dtype=dtype)