Refactor

Author: mathpluscode

cinema/classification/eval.py

@@ -83,8 +83,8 @@ def classification_eval_dataset(  # pylint:disable=too-many-statements
     model = get_classification_or_regression_model(config)
     checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
     model.load_state_dict(checkpoint["model"])
-    model.to(device)
     model.eval()
+    model.to(device)
 
     # inference
     pred_labels = []

cinema/examples/inference/classification_cvd.py

@@ -28,6 +28,7 @@ def run(trained_dataset: str, view: str, seed: int, device: torch.device, dtype:
         model_filename=f"finetuned/classification_cvd/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
         config_filename=f"finetuned/classification_cvd/{trained_dataset}_{view}/config.yaml",
     )
+    model.eval()
     model.to(device)
 
     # load sample data from mnms2 of class HCM and form a batch of size 1

cinema/examples/inference/classification_sex.py

@@ -29,6 +29,7 @@ def run(seed: int, device: torch.device, dtype: torch.dtype) -> None:
         model_filename=f"finetuned/classification_sex/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
         config_filename=f"finetuned/classification_sex/{trained_dataset}_{view}/config.yaml",
     )
+    model.eval()
     model.to(device)
 
     # load sample data from mnms2 of class HCM and form a batch of size 1

cinema/examples/inference/classification_vendor.py

@@ -29,6 +29,7 @@ def run(view: str, seed: int, device: torch.device, dtype: torch.dtype) -> None:
         model_filename=f"finetuned/classification_vendor/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
         config_filename=f"finetuned/classification_vendor/{trained_dataset}_{view}/config.yaml",
     )
+    model.eval()
     model.to(device)
 
     # load sample data from mnms2 of class HCM and form a batch of size 1

cinema/examples/inference/landmark_coordinate.py

@@ -20,6 +20,7 @@ def run(view: str, seed: int, device: torch.device, dtype: torch.dtype) -> None:
         model_filename=f"finetuned/landmark_coordinate/{view}/{view}_{seed}.safetensors",
         config_filename=f"finetuned/landmark_coordinate/{view}/config.yaml",
     )
+    model.eval()
     model.to(device)
 
     # load sample data and form a batch of size 1

cinema/examples/inference/landmark_heatmap.py

@@ -20,6 +20,7 @@ def run(view: str, seed: int, device: torch.device, dtype: torch.dtype) -> None:
         model_filename=f"finetuned/landmark_heatmap/{view}/{view}_{seed}.safetensors",
         config_filename=f"finetuned/landmark_heatmap/{view}/config.yaml",
     )
+    model.eval()
     model.to(device)
 
     # load sample data and form a batch of size 1

cinema/examples/inference/mae.py

@@ -11,12 +11,73 @@
 from cinema import CineMA, patchify, unpatchify
 
 
+def plot_mae_reconstruction(
+    batch: dict[str, torch.Tensor],
+    pred_dict: dict[str, torch.Tensor],
+    enc_mask_dict: dict[str, torch.Tensor],
+    patch_size_dict: dict[str, tuple[int, ...]],
+    grid_size_dict: dict[str, tuple[int, ...]],
+    sax_slices: int,
+) -> plt.Figure:
+    """Plot MAE reconstruction."""
+    n_rows = sax_slices + 3
+    n_cols = 4
+    fig, axs = plt.subplots(n_rows, n_cols, figsize=(n_cols * 2, n_rows * 2), dpi=300)
+    for i, view in enumerate(["lax_2c", "lax_3c", "lax_4c", "sax"]):
+        patches = patchify(image=batch[view], patch_size=patch_size_dict[view])
+        patches[enc_mask_dict[view]] = pred_dict[view]
+        masks = torch.zeros_like(patches)
+        masks[enc_mask_dict[view]] = 1
+        masks = unpatchify(masks, patch_size=patch_size_dict[view], grid_size=grid_size_dict[view])
+        masks = masks[0, 0]
+        reconstructed = unpatchify(
+            patches,
+            patch_size=patch_size_dict[view],
+            grid_size=grid_size_dict[view],
+        )
+        reconstructed = reconstructed[0, 0].numpy()
+        image = batch[view][0, 0].numpy()
+        error = np.abs(reconstructed - image)
+
+        if view == "sax":
+            reconstructed = reconstructed[..., :sax_slices]
+            for j in range(sax_slices):
+                axs[3 + j, 0].set_ylabel(f"SAX slice {j}")
+                axs[3 + j, 0].imshow(image[..., j], cmap="gray")
+                axs[3 + j, 1].imshow(masks[..., j], cmap="gray")
+                axs[3 + j, 2].imshow(reconstructed[..., j], cmap="gray")
+                axs[3 + j, 3].imshow(error[..., j], cmap="gray")
+        else:
+            axs[i, 0].imshow(image, cmap="gray")
+            axs[i, 1].imshow(masks, cmap="gray")
+            axs[i, 2].imshow(reconstructed, cmap="gray")
+            axs[i, 3].imshow(error, cmap="gray")
+            axs[i, 0].set_ylabel({"lax_2c": "LAX 2C", "lax_3c": "LAX 3C", "lax_4c": "LAX 4C"}[view])
+            if i == 0:
+                axs[i, 0].set_title("Original")
+                axs[i, 1].set_title("Mask")
+                axs[i, 2].set_title("Reconstructed")
+                axs[i, 3].set_title("Error")
+    # remove the x and y ticks
+    for i in range(n_rows):
+        for j in range(n_cols):
+            axs[i, j].set_xticks([])
+            axs[i, j].set_yticks([])
+    fig.tight_layout()
+    fig.subplots_adjust(wspace=0, hspace=0)
+    return fig
+
+
 def run(device: torch.device, dtype: torch.dtype) -> None:
     """Run MAE reconstruction."""
+    t = 25  # which time frame to use
+
     # load model
     model = CineMA.from_pretrained()
-    model.to(device)
     model.eval()
+    patch_size_dict = model.dec_patch_size_dict
+    grid_size_dict = {k: v.patch_embed.grid_size for k, v in model.enc_down_dict.items()}
+    model.to(device)
 
     # load sample data and form a batch of size 1
     transform = Compose(
@@ -46,7 +107,7 @@ def run(device: torch.device, dtype: torch.dtype) -> None:
     lax_4c_image = torch.from_numpy(
         np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_lax_4c.nii.gz")))
     )
-    t = 25  # which time frame to use
+    sax_slices = sax_image.shape[-2]
     batch = {
         "sax": sax_image[None, ..., t],
         "lax_2c": lax_2c_image[None, ..., 0, t],
@@ -62,45 +123,9 @@ def run(device: torch.device, dtype: torch.dtype) -> None:
         _, pred_dict, enc_mask_dict, _ = model(batch, enc_mask_ratio=0.75)
 
     # visualize
-    _, axs = plt.subplots(6, 4, figsize=(12, 18))
-    for i, view in enumerate(["lax_2c", "lax_3c", "lax_4c", "sax"]):
-        patches = patchify(image=batch[view], patch_size=model.dec_patch_size_dict[view])
-        patches[enc_mask_dict[view]] = pred_dict[view]
-        masks = torch.zeros_like(patches)
-        masks[enc_mask_dict[view]] = 1
-        masks = unpatchify(
-            masks, patch_size=model.dec_patch_size_dict[view], grid_size=model.enc_down_dict[view].patch_embed.grid_size
-        )
-        masks = masks[0, 0]
-        reconstructed = unpatchify(
-            patches,
-            patch_size=model.dec_patch_size_dict[view],
-            grid_size=model.enc_down_dict[view].patch_embed.grid_size,
-        )
-        reconstructed = reconstructed[0, 0].detach().cpu().numpy()
-        image = batch[view][0, 0].detach().cpu().numpy()
-        error = np.abs(reconstructed - image)
-
-        if view == "sax":
-            for j in range(3):
-                z = j * 3
-                axs[3 + j, 0].set_ylabel(f"{view} slice {z}")
-                axs[3 + j, 0].imshow(image[..., z], cmap="gray")
-                axs[3 + j, 1].imshow(masks[..., z], cmap="gray")
-                axs[3 + j, 2].imshow(reconstructed[..., z], cmap="gray")
-                axs[3 + j, 3].imshow(error[..., z], cmap="gray")
-        else:
-            axs[i, 0].imshow(image, cmap="gray")
-            axs[i, 1].imshow(masks, cmap="gray")
-            axs[i, 2].imshow(reconstructed, cmap="gray")
-            axs[i, 3].imshow(error, cmap="gray")
-            axs[i, 0].set_ylabel(view)
-            if i == 0:
-                axs[i, 0].set_title("Original")
-                axs[i, 1].set_title("Mask")
-                axs[i, 2].set_title("Reconstructed")
-                axs[i, 3].set_title("Error")
-    plt.savefig("mae_reconstruction.png", dpi=300, bbox_inches="tight")
+    batch = {k: v.detach().cpu() for k, v in batch.items()}
+    fig = plot_mae_reconstruction(batch, pred_dict, enc_mask_dict, patch_size_dict, grid_size_dict, sax_slices)
+    fig.savefig("mae_reconstruction.png", dpi=300, bbox_inches="tight")
     plt.show(block=False)
 
 

cinema/examples/inference/mae_feature_extraction.py

@@ -14,8 +14,8 @@ def run(device: torch.device, dtype: torch.dtype) -> None:
     """Run MAE feature extraction."""
     # load model
     model = CineMA.from_pretrained()
-    model.to(device)
     model.eval()
+    model.to(device)
 
     # load sample data and form a batch of size 1
     transform = Compose(

cinema/examples/inference/regression_age.py

@@ -30,6 +30,7 @@ def run(seed: int, device: torch.device, dtype: torch.dtype) -> None:
         model_filename=f"finetuned/regression_age/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
         config_filename=f"finetuned/regression_age/{trained_dataset}_{view}/config.yaml",
     )
+    model.eval()
     model.to(device)
 
     # load sample data

cinema/examples/inference/regression_bmi.py

@@ -30,6 +30,7 @@ def run(seed: int, device: torch.device, dtype: torch.dtype) -> None:
         model_filename=f"finetuned/regression_bmi/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
         config_filename=f"finetuned/regression_bmi/{trained_dataset}_{view}/config.yaml",
     )
+    model.eval()
     model.to(device)
 
     # load sample data