Add landmark tab

app.py

@@ -6,8 +6,13 @@ import requests
 import SimpleITK as sitk  # noqa: N813
 import spaces
 import torch
-from cinema import CineMA, ConvUNetR
+from cinema import CineMA, ConvUNetR, ConvViT, heatmap_soft_argmax
 from cinema.examples.cine_cmr import plot_cmr_views
+from cinema.examples.inference.landmark_heatmap import (
+    plot_heatmaps,
+    plot_landmarks,
+    plot_lv,
+)
 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,
@@ -63,7 +68,7 @@ def cmr_tab():
     with gr.Blocks() as cmr_interface:
         gr.Markdown(
             """
-            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.
+            This page illustrates the spatial orientation of short-axis (SAX) and long-axis (LAX) views in 3D.
             """
         )
         with gr.Row():
@@ -222,7 +227,7 @@ 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.
+            This page demonstrates the masking and reconstruction process of the masked autoencoder. The model was trained with a mask ratio of 0.75 over 74,000 studies.
             """
         )
         with gr.Row():
@@ -334,8 +339,7 @@ def segmentation_sax_tab():
     with gr.Blocks() as sax_interface:
         gr.Markdown(
             """
-            This page demonstrates the segmentation of cardiac structures in the Short-Axis (SAX) view.
-            Please adjust the settings on the right panels and click the button to run the inference.
+            This page demonstrates the segmentation of cardiac structures in the short-axis (SAX) view.
             """
         )
 
@@ -345,17 +349,16 @@ def segmentation_sax_tab():
                 ## Description
                 ### 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 101 - 150 are from the test set.
+                Images 101–150 are from the test set of [ACDC](https://www.creatis.insa-lyon.fr/Challenge/acdc/).
 
                 ### Model
 
-                The available models are finetuned on different datasets ([ACDC](https://www.creatis.insa-lyon.fr/Challenge/acdc/), [M&Ms](https://www.ub.edu/mnms/), and [M&Ms2](https://www.ub.edu/mnms-2/)). For each dataset, there are 3 models finetuned on different seeds: 0, 1, 2.
+                The available models are finetuned on different datasets ([ACDC](https://www.creatis.insa-lyon.fr/Challenge/acdc/), [M&Ms](https://www.ub.edu/mnms/), and [M&Ms2](https://www.ub.edu/mnms-2/)). For each dataset, there are three models finetuned with seeds: 0, 1, 2.
 
-                ### Visualization
+                ### Visualisation
 
-                The left figure shows the segmentation of ventricles and myocardium every n time steps across all SAX slices.
-                The right figure plots the ventricle and mycoardium volumes across all inference time frames.
+                The left figure shows the segmentation of ventricles and myocardium at every n time step across all SAX slices.
+                The right figure shows the volumes across all time frames and estimates the ejection fraction (EF) for the left ventricle (LV) and right ventricle (RV).
                 """)
             with gr.Column(scale=3):
                 gr.Markdown("## Data Settings")
@@ -363,7 +366,7 @@ def segmentation_sax_tab():
                     minimum=101,
                     maximum=150,
                     step=1,
-                    label="Choose an ACDC image, ID is between 101 and 150",
+                    label="Choose an image, ID is between 101 and 150",
                     value=101,
                 )
                 t_step = gr.Slider(
@@ -374,10 +377,10 @@ def segmentation_sax_tab():
                     value=2,
                 )
             with gr.Column(scale=3):
-                gr.Markdown("## Model Setting")
+                gr.Markdown("## Model Settings")
                 trained_dataset = gr.Dropdown(
                     choices=["ACDC", "M&MS", "M&MS2"],
-                    label="Choose which dataset the segmentation model was finetuned on",
+                    label="Choose which dataset the model was finetuned on",
                     value="ACDC",
                 )
                 seed = gr.Slider(
@@ -394,7 +397,7 @@ def segmentation_sax_tab():
                 gr.Markdown("## Ventricle and Myocardium Segmentation")
                 segmentation_plot = gr.Plot(show_label=False)
             with gr.Column():
-                gr.Markdown("## Ejection Fraction Prediction")
+                gr.Markdown("## Volume Estimation")
                 volume_plot = gr.Plot(show_label=False)
 
         run_button.click(
@@ -476,8 +479,7 @@ def segmentation_lax_tab():
     with gr.Blocks() as lax_interface:
         gr.Markdown(
             """
-            This page demonstrates the segmentation of cardiac structures in the Long-Axis (LAX) view.
-            Please adjust the settings on the right panels and click the button to run the inference.
+            This page demonstrates the segmentation of cardiac structures in the long-axis (LAX) four-chamber (4C) view.
             """
         )
 
@@ -487,16 +489,16 @@ def segmentation_lax_tab():
                 ## Description
                 ### Data
 
-                There are four example samples. All images have been resampled to 1 mm × 1 mm and centre-cropped.
+                There are four example images from the UK Biobank.
 
                 ### Model
 
-                The available models are finetuned on [M&Ms2](https://www.ub.edu/mnms-2/). For each dataset, there are 3 models finetuned on different seeds: 0, 1, 2.
+                The available models are finetuned on [M&Ms2](https://www.ub.edu/mnms-2/). There are three models finetuned with seeds: 0, 1, 2.
 
-                ### Visualization
+                ### Visualisation
 
                 The left figure shows the segmentation of ventricles and myocardium across all time frames.
-                The right figure plots the ventricle and mycoardium volumes across all inference time frames.
+                The right figure shows the volumes across all time frames and estimates the ejection fraction (EF).
                 """)
             with gr.Column(scale=3):
                 gr.Markdown("## Data Settings")
@@ -505,10 +507,10 @@ def segmentation_lax_tab():
                     maximum=4,
                     step=1,
                     label="Choose an image, ID is between 1 and 4",
-                    value=4,
+                    value=1,
                 )
             with gr.Column(scale=3):
-                gr.Markdown("## Model Setting")
+                gr.Markdown("## Model Settings")
                 seed = gr.Slider(
                     minimum=0,
                     maximum=2,
@@ -516,7 +518,7 @@ def segmentation_lax_tab():
                     label="Choose which seed the finetuning used",
                     value=0,
                 )
-        run_button = gr.Button("Run LAX segmentation inference", variant="primary")
+        run_button = gr.Button("Run LAX 4C segmentation inference", variant="primary")
 
         with gr.Row():
             with gr.Column():
@@ -534,6 +536,194 @@ def segmentation_lax_tab():
     return lax_interface
 
 
+@spaces.GPU
+def landmark_heatmap_inference(
+    images: torch.Tensor,
+    view: str,
+    transform: Compose,
+    model: ConvUNetR,
+    progress=gr.Progress(),
+) -> tuple[np.ndarray, np.ndarray]:
+    model.to(device)
+
+    n_frames = images.shape[-1]
+    probs_list = []
+    coords_list = []
+    for t in tqdm(range(n_frames), total=n_frames):
+        progress((t + 1) / n_frames, desc=f"Processing frame {t + 1} / {n_frames}...")
+        batch = transform({view: torch.from_numpy(images[None, ..., 0, t])})
+        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()),
+        ):
+            logits = model(batch)[view]  # (1, 3, x, y)
+        probs = torch.sigmoid(logits)  # (1, 3, width, height)
+        probs_list.append(probs[0].detach().to(torch.float32).cpu().numpy())
+        coords = heatmap_soft_argmax(probs)[0].detach().to(torch.float32).cpu().numpy()
+        coords = [int(x) for x in coords]
+        coords_list.append(coords)
+    probs = np.stack(probs_list, axis=-1)  # (3, x, y, t)
+    coords = np.stack(coords_list, axis=-1)  # (6, t)
+    return probs, coords
+
+
+@spaces.GPU
+def landmark_coordinate_inference(
+    images: torch.Tensor,
+    view: str,
+    transform: Compose,
+    model: ConvViT,
+    progress=gr.Progress(),
+) -> np.ndarray:
+    model.to(device)
+
+    w, h, _, n_frames = images.shape
+    coords_list = []
+    for t in tqdm(range(n_frames), total=n_frames):
+        progress((t + 1) / n_frames, desc=f"Processing frame {t + 1} / {n_frames}...")
+        batch = transform({view: torch.from_numpy(images[None, ..., 0, t])})
+        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()),
+        ):
+            coords = model(batch)[0].detach().to(torch.float32).cpu().numpy()  # (6,)
+        coords *= np.array([w, h, w, h, w, h])
+        coords = [int(x) for x in coords]
+        coords_list.append(coords)
+    coords = np.stack(coords_list, axis=-1)  # (6, t)
+    return coords
+
+
+def landmark(image_id, view, method, seed, progress=gr.Progress()):
+    view = "lax_2c" if view == "LAX 2C" else "lax_4c"
+    method = method.lower()
+
+    # Download and load model
+    progress(0, desc="Downloading model...")
+    if method == "heatmap":
+        model = ConvUNetR.from_finetuned(
+            repo_id="mathpluscode/CineMA",
+            model_filename=f"finetuned/landmark_{method}/{view}/{view}_{seed}.safetensors",
+            config_filename=f"finetuned/landmark_{method}/{view}/config.yaml",
+            cache_dir=cache_dir,
+        )
+    elif method == "coordinate":
+        model = ConvViT.from_finetuned(
+            repo_id="mathpluscode/CineMA",
+            model_filename=f"finetuned/landmark_{method}/{view}/{view}_{seed}.safetensors",
+            config_filename=f"finetuned/landmark_{method}/{view}/config.yaml",
+            cache_dir=cache_dir,
+        )
+    else:
+        raise ValueError(f"Invalid method: {method}")
+    model.eval()
+
+    # Inference
+    progress(0, desc="Downloading data...")
+    transform = ScaleIntensityd(keys=view)
+    images = np.transpose(
+        sitk.GetArrayFromImage(
+            load_nifti_from_github(f"ukb/{image_id}/{image_id}_{view}.nii.gz")
+        )
+    )
+
+    if method == "heatmap":
+        _, coords = landmark_heatmap_inference(images, view, transform, model, progress)
+    elif method == "coordinate":
+        coords = landmark_coordinate_inference(images, view, transform, model, progress)
+    else:
+        raise ValueError(f"Invalid method: {method}")
+
+    landmark_fig = plot_landmarks(images, coords)
+    lv_fig = plot_lv(coords)
+    return landmark_fig, lv_fig
+
+
+def landmark_tab():
+    with gr.Blocks() as landmark_interface:
+        gr.Markdown(
+            """
+            This page demonstrates landmark localisation in the long-axis (LAX) two-chamber (2C) and four-chamber (4C) views
+            """
+        )
+
+        with gr.Row():
+            with gr.Column(scale=4):
+                gr.Markdown("""
+                ## Description
+                ### Data
+
+                There are four example images from the UK Biobank.
+
+                ### Model
+
+                The available models are finetuned on data from [Xue et al.](https://pubs.rsna.org/doi/10.1148/ryai.2021200197).
+                There are two types of landmark localisation models:
+
+                - **Heatmap**: predicts dense probability maps of landmarks
+                - **Coordinate**: predicts landmark coordinates directly
+
+                For each type, there are three models finetuned with seeds: 0, 1, 2.
+
+                ### Visualisation
+
+                The left figure shows the landmark positions across all time frames.
+                The right figure shows the length of the left ventricle across all time frames and the estimates of two metrics:
+                - Mitral annular plane systolic excursion (MAPSE)
+                - Global longitudinal shortening (GLS)
+                """)
+            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,
+                )
+                view = gr.Dropdown(
+                    choices=["LAX 2C", "LAX 4C"],
+                    label="Choose which view to localise the landmarks",
+                    value="LAX 2C",
+                )
+            with gr.Column(scale=3):
+                gr.Markdown("## Model Settings")
+                method = gr.Dropdown(
+                    choices=["Heatmap", "Coordinate"],
+                    label="Choose which method to use",
+                    value="Heatmap",
+                )
+                seed = gr.Slider(
+                    minimum=0,
+                    maximum=2,
+                    step=1,
+                    label="Choose which seed the finetuning used",
+                    value=0,
+                )
+        run_button = gr.Button("Run landmark localisation inference", variant="primary")
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("## Landmark Localisation")
+                landmark_plot = gr.Plot(show_label=False)
+            with gr.Column():
+                gr.Markdown("## Left Ventricle Length Estimation")
+                lv_plot = gr.Plot(show_label=False)
+
+        run_button.click(
+            fn=landmark,
+            inputs=[image_id, view, method, seed],
+            outputs=[landmark_plot, lv_plot],
+        )
+    return landmark_interface
+
+
 with gr.Blocks(
     theme=theme, title="CineMA: A Foundation Model for Cine Cardiac MRI"
 ) as demo:
@@ -542,7 +732,7 @@ with gr.Blocks(
         # CineMA: A Foundation Model for Cine Cardiac MRI 🎥🫀
 
         This demo showcases the capabilities of CineMA in multiple tasks.
-        For more details, checkout our [GitHub](https://github.com/mathpluscode/CineMA).
+        For more details, check out our [GitHub](https://github.com/mathpluscode/CineMA).
         """
     )
 
@@ -553,7 +743,8 @@ with gr.Blocks(
             mae_tab()
         with gr.TabItem("Segmentation in SAX View"):
             segmentation_sax_tab()
-        with gr.TabItem("Segmentation in LAX View"):
+        with gr.TabItem("Segmentation in LAX 4C View"):
             segmentation_lax_tab()
-
+        with gr.TabItem("Landmark Localisation in LAX 2C/4C View"):
+            landmark_tab()
 demo.launch()

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@3ace4d79ee037f95e8767b35c7bc97d511f8b9c1#egg=cinema
+git+https://github.com/mathpluscode/CineMA@edc0774baed3c2429a2a4ffaa36a3910f2780b2b#egg=cinema
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch==2.5.1