Update layout

.gitignore

@@ -2,3 +2,9 @@
 .DS_Store
 node_modules/
 src/cache/
+.locks/
+ukb/
+*.gif
+*.png
+models--mathpluscode--CineMA/
+datasets--mathpluscode--ACDC/

app.py

@@ -10,6 +10,7 @@ import torch
 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_heatmap_and_landmarks,
     plot_landmarks,
     plot_lv,
 )
@@ -81,8 +82,8 @@ def cmr_tab():
                     minimum=1,
                     maximum=4,
                     step=1,
-                    label="Choose an image, ID is between 1 and 4",
-                    value=1,
+                    label="Choose an image",
+                    value=2,
                 )
                 # Placeholder for slice slider, will update dynamically
                 slice_idx = gr.Slider(
@@ -90,7 +91,7 @@ def cmr_tab():
                     maximum=8,
                     step=1,
                     label="SAX slice to visualize",
-                    value=0,
+                    value=1,
                 )
 
         def get_num_slices(image_id):
@@ -99,7 +100,7 @@ def cmr_tab():
 
         def update_slice_slider(image_id):
             num_slices = get_num_slices(image_id)
-            return gr.update(maximum=num_slices - 1, value=0, visible=True)
+            return gr.update(maximum=num_slices - 1, value=1, visible=True)
 
         def fn(image_id, slice_idx):
             lax_2c_image = load_nifti_from_github(
@@ -125,7 +126,7 @@ def cmr_tab():
 
         # When image changes, update the slice slider and plot
         gr.on(
-            fn=lambda image_id: [update_slice_slider(image_id), fn(image_id, 0)],
+            fn=lambda image_id: [update_slice_slider(image_id), fn(image_id, 1)],
             inputs=[image_id],
             outputs=[slice_idx, cmr_plot],
         )
@@ -238,9 +239,11 @@ def mae_tab():
     with gr.Blocks() as mae_interface:
         gr.Markdown(
             """
-            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.
+            This page demonstrates the masking and reconstruction process. The model was trained with a mask ratio of 0.75. Click the button below to launch the inference. ⬇️
             """
         )
+        run_button = gr.Button("Launch reconstruction", variant="primary")
+
         with gr.Row():
             with gr.Column(scale=5):
                 gr.Markdown("## Reconstruction")
@@ -249,14 +252,14 @@ def mae_tab():
                     type="filepath",
                     label="Masked Autoencoder Reconstruction",
                 )
-            with gr.Column(scale=3):
+            with gr.Column(scale=5):
                 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,
+                    label="Choose an image",
+                    value=2,
                 )
                 mask_ratio = gr.Slider(
                     minimum=0.05,
@@ -265,7 +268,6 @@ def mae_tab():
                     label="Mask ratio",
                     value=0.75,
                 )
-                run_button = gr.Button("Run masked autoencoder", variant="primary")
 
         run_button.click(
             fn=mae,
@@ -376,9 +378,10 @@ 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.
+            This page demonstrates the segmentation of cardiac structures in the short-axis (SAX) view. Click the button below to launch the inference. ⬇️
             """
         )
+        run_button = gr.Button("Launch segmentation inference", variant="primary")
 
         with gr.Row():
             with gr.Column(scale=4):
@@ -391,43 +394,47 @@ def segmentation_sax_tab():
                 ### 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 three models finetuned with seeds: 0, 1, 2.
-
-                ### Visualisation
-
-                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")
-                image_id = gr.Slider(
-                    minimum=101,
-                    maximum=150,
-                    step=1,
-                    label="Choose an image, ID is between 101 and 150",
-                    value=101,
-                )
-                t_step = gr.Slider(
-                    minimum=1,
-                    maximum=10,
-                    step=1,
-                    label="Choose the gap between time frames",
-                    value=3,
-                )
-            with gr.Column(scale=3):
-                gr.Markdown("## Model Settings")
-                trained_dataset = gr.Dropdown(
-                    choices=["ACDC", "M&MS", "M&MS2"],
-                    label="Choose which dataset the model was finetuned on",
-                    value="ACDC",
-                )
-                seed = gr.Slider(
-                    minimum=0,
-                    maximum=2,
-                    step=1,
-                    label="Choose which seed the finetuning used",
-                    value=0,
-                )
-        run_button = gr.Button("Run SAX segmentation inference", variant="primary")
+            with gr.Column(scale=6):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        gr.Markdown("## Data Settings")
+                        image_id = gr.Slider(
+                            minimum=101,
+                            maximum=150,
+                            step=1,
+                            label="Choose an image",
+                            value=150,
+                        )
+                        t_step = gr.Slider(
+                            minimum=1,
+                            maximum=10,
+                            step=1,
+                            label="Choose the gap between time frames",
+                            value=3,
+                        )
+                    with gr.Column(scale=1):
+                        gr.Markdown("## Model Settings")
+                        trained_dataset = gr.Dropdown(
+                            choices=["ACDC", "M&MS", "M&MS2"],
+                            label="Choose which dataset the model was finetuned on",
+                            value="ACDC",
+                        )
+                        seed = gr.Slider(
+                            minimum=0,
+                            maximum=2,
+                            step=1,
+                            label="Choose which seed the finetuning used",
+                            value=1,
+                        )
+
+        # Visualisation description block
+        gr.Markdown("""
+        ## Visualisation
+
+        The left figure shows the segmentation at every n time step across all SAX slices.
+        The right figure shows the volumes across time frames and estimates the ejection fraction (EF) for the left ventricle (LV) and right ventricle (RV).
+        """)
 
         with gr.Row():
             with gr.Column():
@@ -535,9 +542,10 @@ 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) four-chamber (4C) view.
+            This page demonstrates the segmentation of cardiac structures in the long-axis (LAX) four-chamber (4C) view. Click the button below to launch the inference. ⬇️
             """
         )
+        run_button = gr.Button("Launch segmentation inference", variant="primary")
 
         with gr.Row():
             with gr.Column(scale=4):
@@ -550,31 +558,35 @@ def segmentation_lax_tab():
                 ### Model
 
                 The available models are finetuned on [M&Ms2](https://www.ub.edu/mnms-2/). There are three models finetuned with seeds: 0, 1, 2.
-
-                ### Visualisation
-
-                The left figure shows the segmentation of ventricles and myocardium across all 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")
-                image_id = gr.Slider(
-                    minimum=1,
-                    maximum=4,
-                    step=1,
-                    label="Choose an image, ID is between 1 and 4",
-                    value=1,
-                )
-            with gr.Column(scale=3):
-                gr.Markdown("## Model Settings")
-                seed = gr.Slider(
-                    minimum=0,
-                    maximum=2,
-                    step=1,
-                    label="Choose which seed the finetuning used",
-                    value=0,
-                )
-        run_button = gr.Button("Run LAX 4C segmentation inference", variant="primary")
+            with gr.Column(scale=6):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        gr.Markdown("## Data Settings")
+                        image_id = gr.Slider(
+                            minimum=1,
+                            maximum=4,
+                            step=1,
+                            label="Choose an image",
+                            value=2,
+                        )
+                    with gr.Column(scale=1):
+                        gr.Markdown("## Model Settings")
+                        seed = gr.Slider(
+                            minimum=0,
+                            maximum=2,
+                            step=1,
+                            label="Choose which seed the finetuning used",
+                            value=1,
+                        )
+
+        # Visualisation description block
+        gr.Markdown("""
+        ## Visualisation
+
+        The left figure shows the segmentation across time frames.
+        The right figure shows the volumes across time frames and estimates the ejection fraction (EF).
+        """)
 
         with gr.Row():
             with gr.Column():
@@ -707,17 +719,18 @@ def landmark(image_id, view, method, seed, progress=gr.Progress()):
     )
 
     if method == "heatmap":
-        _, coords = landmark_heatmap_inference(images, view, transform, model, progress)
+        probs, coords = landmark_heatmap_inference(
+            images, view, transform, model, progress
+        )
+        progress(1, desc="Inference finished. Plotting ...")
+        plot_heatmap_and_landmarks(images, probs, coords, landmark_path)
     elif method == "coordinate":
         coords = landmark_coordinate_inference(images, view, transform, model, progress)
+        progress(1, desc="Inference finished. Plotting ...")
+        plot_landmarks(images, coords, landmark_path)
     else:
         raise ValueError(f"Invalid method: {method}")
 
-    progress(1, desc="Inference finished. Plotting ...")
-
-    # Plot landmarks in GIF
-    plot_landmarks(images, coords, landmark_path)
-
     # Plot LV change in PNG
     plot_lv(coords, lv_path)
 
@@ -728,9 +741,12 @@ 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.
+            This page demonstrates landmark localisation in the long-axis (LAX) two-chamber (2C) and four-chamber (4C) views. Click the button below to launch the inference. ⬇️
             """
         )
+        run_button = gr.Button(
+            "Launch landmark localisation inference", variant="primary"
+        )
 
         with gr.Row():
             with gr.Column(scale=4):
@@ -749,43 +765,45 @@ def landmark_tab():
                 - **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.Column(scale=6):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        gr.Markdown("## Data Settings")
+                        image_id = gr.Slider(
+                            minimum=1,
+                            maximum=4,
+                            step=1,
+                            label="Choose an image",
+                            value=2,
+                        )
+                        view = gr.Dropdown(
+                            choices=["LAX 2C", "LAX 4C"],
+                            label="Choose which view to localise the landmarks",
+                            value="LAX 2C",
+                        )
+                    with gr.Column(scale=1):
+                        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=1,
+                        )
+
+        # Visualisation description block
+        gr.Markdown("""
+        ## Visualisation
+
+        The left figure shows the landmark positions across time frames.
+        The right figure shows the length of the left ventricle across time frames and estimates mitral annular plane systolic excursion (MAPSE) and global longitudinal shortening (GLS).
+        """)
 
         with gr.Row():
             with gr.Column():
@@ -816,9 +834,9 @@ with gr.Blocks(
         """
         # CineMA: A Foundation Model for Cine Cardiac MRI 🎥🫀
 
-        🚀 The following demos showcase the capabilities of CineMA in multiple tasks.<br>
+        🚀 The following demonstrations showcase the capabilities of CineMA in multiple tasks. Click the button to launch the inference.<br>
         ⏱️ The examples may take 10-60 seconds, if not cached, to download data and model, perform inference, and render plots.<br>
-        🔗 For more details, check out our [GitHub](https://github.com/mathpluscode/CineMA).
+        🔗 For more details, check out our [GitHub repository](https://github.com/mathpluscode/CineMA).
         """
     )
 

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