Add MAE tab

.pre-commit-config.yaml

@@ -27,7 +27,7 @@ repos:
     hooks:
       # run the linter
       - id: ruff
-        args: [--fix]
+        args: ["--fix", "--select=I"]
       # run the formatter
       - id: ruff-format
   - repo: https://github.com/pre-commit/mirrors-prettier

README.md

@@ -1,11 +1,11 @@
 ---
 title: CineMA
 tags:
-- medical
-- cardiac
-- MRI
-- foundation model
-- MAE
+  - medical
+  - cardiac
+  - MRI
+  - foundation model
+  - MAE
 emoji: 🚀
 colorFrom: red
 colorTo: purple
@@ -23,4 +23,4 @@ thumbnail: >-
 # CineMA: A Foundation Model for Cine Cardiac MRI
 
 This is a demo of CineMA, a foundation model for cine cardiac MRI. For more details, checkout our
-[GitHub](https://github.com/mathpluscode/CineMA).
+[GitHub](https://github.com/mathpluscode/CineMA).

app.py

@@ -1,24 +1,32 @@
-import numpy as np
+from pathlib import Path
+
 import gradio as gr
-from huggingface_hub import hf_hub_download
+import numpy as np
+import requests
 import SimpleITK as sitk  # noqa: N813
+import spaces
 import torch
-from monai.transforms import Compose, ScaleIntensityd, SpatialPadd
-from cinema import ConvUNetR
-from pathlib import Path
-from cinema.examples.inference.segmentation_sax import (
-    plot_segmentations as plot_segmentations_sax,
-    plot_volume_changes as plot_volume_changes_sax,
-)
+from cinema import CineMA, ConvUNetR
+from cinema.examples.cine_cmr import plot_cmr_views
+from cinema.examples.inference.mae import plot_mae_reconstruction, reconstruct_images
 from cinema.examples.inference.segmentation_lax_4c import (
     plot_segmentations as plot_segmentations_lax,
+)
+from cinema.examples.inference.segmentation_lax_4c import (
     plot_volume_changes as plot_volume_changes_lax,
+)
+from cinema.examples.inference.segmentation_lax_4c import (
     post_process as post_process_lax_segmentation,
 )
-from cinema.examples.cine_cmr import plot_cmr_views
+from cinema.examples.inference.segmentation_sax import (
+    plot_segmentations as plot_segmentations_sax,
+)
+from cinema.examples.inference.segmentation_sax import (
+    plot_volume_changes as plot_volume_changes_sax,
+)
+from huggingface_hub import hf_hub_download
+from monai.transforms import Compose, ScaleIntensityd, SpatialPadd
 from tqdm import tqdm
-import spaces
-import requests
 
 # cache directories
 cache_dir = Path("/tmp/.cinema")
@@ -52,18 +60,17 @@ def load_nifti_from_github(name: str) -> sitk.Image:
 
 
 def cmr_tab():
-    with gr.Blocks() as sax_interface:
+    with gr.Blocks() as cmr_interface:
         gr.Markdown(
             """
-            This page demonstrates the geometry of SAX and LAX views in 3D spaces.
-            Please adjust the settings on the right panels to select images and slices.
+            This page illustrates the spatial orientation of short-axis (SAX) and long-axis (LAX) views in 3D. Use the control panels on the right to select specific images and slices.
             """
         )
         with gr.Row():
-            with gr.Column(scale=3):
+            with gr.Column(scale=5):
                 gr.Markdown("## Views")
                 cmr_plot = gr.Plot(show_label=False)
-            with gr.Column(scale=1):
+            with gr.Column(scale=3):
                 gr.Markdown("## Data Settings")
                 image_id = gr.Slider(
                     minimum=1,
@@ -125,7 +132,127 @@ def cmr_tab():
             outputs=[cmr_plot],
         )
 
-    return sax_interface
+    return cmr_interface
+
+
+@spaces.GPU
+def mae_inference(
+    batch: dict[str, torch.Tensor],
+    transform: Compose,
+    model: CineMA,
+    mask_ratio: float,
+) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict[str, np.ndarray]]:
+    model.to(device)
+    sax_slices = batch["sax"].shape[-1]
+    batch = transform(batch)
+    batch = {k: v[None, ...].to(device=device, dtype=dtype) for k, v in batch.items()}
+    with (
+        torch.no_grad(),
+        torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()),
+    ):
+        _, pred_dict, enc_mask_dict, _ = model(batch, enc_mask_ratio=mask_ratio)
+        grid_size_dict = {
+            k: v.patch_embed.grid_size for k, v in model.enc_down_dict.items()
+        }
+        reconstructed_dict, masks_dict = reconstruct_images(
+            batch,
+            pred_dict,
+            enc_mask_dict,
+            model.dec_patch_size_dict,
+            grid_size_dict,
+            sax_slices,
+        )
+        batch = {
+            k: v.detach().to(torch.float32).cpu().numpy()[0, 0]
+            for k, v in batch.items()
+        }
+        batch["sax"] = batch["sax"][..., :sax_slices]
+        return batch, reconstructed_dict, masks_dict
+
+
+def mae(image_id, mask_ratio, progress=gr.Progress()):
+    t = 4  # which time frame to use
+    progress(0, desc="Downloading model...")
+    model = CineMA.from_pretrained()
+    model.eval()
+
+    progress(0, desc="Downloading data...")
+    lax_2c_image = load_nifti_from_github(f"ukb/{image_id}/{image_id}_lax_2c.nii.gz")
+    lax_3c_image = load_nifti_from_github(f"ukb/{image_id}/{image_id}_lax_3c.nii.gz")
+    lax_4c_image = load_nifti_from_github(f"ukb/{image_id}/{image_id}_lax_4c.nii.gz")
+    sax_image = load_nifti_from_github(f"ukb/{image_id}/{image_id}_sax.nii.gz")
+    transform = Compose(
+        [
+            ScaleIntensityd(keys=("sax", "lax_2c", "lax_3c", "lax_4c")),
+            SpatialPadd(keys="sax", spatial_size=(192, 192, 16), method="end"),
+            SpatialPadd(
+                keys=("lax_2c", "lax_3c", "lax_4c"),
+                spatial_size=(256, 256),
+                method="end",
+            ),
+        ]
+    )
+    lax_2c_image_np = np.transpose(sitk.GetArrayFromImage(lax_2c_image))
+    lax_3c_image_np = np.transpose(sitk.GetArrayFromImage(lax_3c_image))
+    lax_4c_image_np = np.transpose(sitk.GetArrayFromImage(lax_4c_image))
+    sax_image_np = np.transpose(sitk.GetArrayFromImage(sax_image))
+    image_dict = {
+        "sax": sax_image_np[None, ..., t],
+        "lax_2c": lax_2c_image_np[None, ..., 0, t],
+        "lax_3c": lax_3c_image_np[None, ..., 0, t],
+        "lax_4c": lax_4c_image_np[None, ..., 0, t],
+    }
+    batch = {k: torch.from_numpy(v) for k, v in image_dict.items()}
+
+    progress(0.5, desc="Running inference...")
+    batch, reconstructed_dict, masks_dict = mae_inference(
+        batch, transform, model, mask_ratio
+    )
+    progress(1, desc="Plotting results...")
+
+    fig = plot_mae_reconstruction(
+        batch,
+        reconstructed_dict,
+        masks_dict,
+    )
+    return fig
+
+
+def mae_tab():
+    with gr.Blocks() as mae_interface:
+        gr.Markdown(
+            """
+            This page illustrates the masking and reconstruction process of the masked autoencoder. The model was trained with mask ratio 0.75 over 74,000 studies.
+            """
+        )
+        with gr.Row():
+            with gr.Column(scale=5):
+                gr.Markdown("## Reconstruction")
+                plot = gr.Plot(show_label=False)
+            with gr.Column(scale=3):
+                gr.Markdown("## Data Settings")
+                image_id = gr.Slider(
+                    minimum=1,
+                    maximum=4,
+                    step=1,
+                    label="Choose an image, ID is between 1 and 4",
+                    value=1,
+                )
+                mask_ratio = gr.Slider(
+                    minimum=0.05,
+                    maximum=1,
+                    step=0.05,
+                    label="Mask ratio",
+                    value=0.75,
+                )
+                run_button = gr.Button("Run Masked Autoencoder", variant="primary")
+        run_button.click(
+            fn=mae,
+            inputs=[image_id, mask_ratio],
+            outputs=[plot],
+        )
+
+    return mae_interface
 
 
 @spaces.GPU
@@ -152,7 +279,7 @@ def segmentation_sax_inference(
         ):
             logits = model(batch)[view]
         labels_list.append(torch.argmax(logits, dim=1)[0, ..., :n_slices])
-    labels = torch.stack(labels_list, dim=-1).detach().cpu().numpy()
+    labels = torch.stack(labels_list, dim=-1).detach().to(torch.float32).cpu().numpy()
     return labels
 
 
@@ -181,6 +308,7 @@ def segmentation_sax(trained_dataset, seed, image_id, t_step, progress=gr.Progre
         config_filename=f"finetuned/segmentation/{trained_dataset}_{view}/config.yaml",
         cache_dir=cache_dir,
     )
+    model.eval()
 
     # Inference
     progress(0, desc="Downloading data...")
@@ -218,7 +346,7 @@ def segmentation_sax_tab():
                 ### Data
 
                 The available data is from ACDC. All images have been resampled to 1 mm × 1 mm × 10 mm and centre-cropped to 192 mm × 192 mm for each SAX slice.
-                Image 1 - 100 are from the training set, and image 101 - 150 are from the test set.
+                Image 101 - 150 are from the test set.
 
                 ### Model
 
@@ -232,11 +360,11 @@ def segmentation_sax_tab():
             with gr.Column(scale=3):
                 gr.Markdown("## Data Settings")
                 image_id = gr.Slider(
-                    minimum=1,
+                    minimum=101,
                     maximum=150,
                     step=1,
-                    label="Choose an ACDC image, ID is between 1 and 150",
-                    value=150,
+                    label="Choose an ACDC image, ID is between 101 and 150",
+                    value=101,
                 )
                 t_step = gr.Slider(
                     minimum=1,
@@ -299,7 +427,9 @@ def segmentation_lax_inference(
             torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()),
         ):
             logits = model(batch)[view]  # (1, 4, x, y)
-        labels = torch.argmax(logits, dim=1)[0].detach().cpu().numpy()  # (x, y)
+        labels = (
+            torch.argmax(logits, dim=1)[0].detach().to(torch.float32).cpu().numpy()
+        )  # (x, y)
 
         # the model seems to hallucinate an additional right ventricle and myocardium sometimes
         # find the connected component that is closest to left ventricle
@@ -322,6 +452,7 @@ def segmentation_lax(seed, image_id, progress=gr.Progress()):
         config_filename=f"finetuned/segmentation/{trained_dataset}_{view}/config.yaml",
         cache_dir=cache_dir,
     )
+    model.eval()
 
     # Inference
     progress(0, desc="Downloading data...")
@@ -418,6 +549,8 @@ with gr.Blocks(
     with gr.Tabs() as tabs:
         with gr.TabItem("Cine CMR Views"):
             cmr_tab()
+        with gr.TabItem("Masked Autoencoder"):
+            mae_tab()
         with gr.TabItem("Segmentation in SAX View"):
             segmentation_sax_tab()
         with gr.TabItem("Segmentation in LAX View"):

requirements.txt

@@ -17,6 +17,6 @@ scikit-learn==1.6.1
 scipy==1.15.2
 spaces==0.36.0
 timm==1.0.15
-git+https://github.com/mathpluscode/CineMA@1ff0e2220676aeff34988614e458588fc8150473#egg=cinema
+git+https://github.com/mathpluscode/CineMA@3ace4d79ee037f95e8767b35c7bc97d511f8b9c1#egg=cinema
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch==2.5.1