Add view panel

app.py

@@ -15,6 +15,7 @@ from cinema.examples.inference.segmentation_lax_4c import (
     plot_volume_changes as plot_volume_changes_lax,
     post_process as post_process_lax_segmentation,
 )
+from cinema.examples.cine_cmr import plot_cmr_views
 from tqdm import tqdm
 import spaces
 import requests
@@ -39,6 +40,94 @@ theme = gr.themes.Ocean(
 )
 
 
+def load_nifti_from_github(name: str) -> sitk.Image:
+    path = cache_dir / name
+    if not path.exists():
+        image_url = f"https://raw.githubusercontent.com/mathpluscode/CineMA/main/cinema/examples/data/{name}"
+        response = requests.get(image_url)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "wb") as f:
+            f.write(response.content)
+    return sitk.ReadImage(path)
+
+
+def cmr_tab():
+    with gr.Blocks() as sax_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.
+            """
+        )
+        with gr.Row():
+            with gr.Column(scale=3):
+                gr.Markdown("## Views")
+                cmr_plot = gr.Plot(show_label=False)
+            with gr.Column(scale=1):
+                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,
+                )
+                # Placeholder for slice slider, will update dynamically
+                slice_idx = gr.Slider(
+                    minimum=0,
+                    maximum=8,
+                    step=1,
+                    label="SAX slice to visualize",
+                    value=0,
+                )
+
+        def get_num_slices(image_id):
+            sax_image = load_nifti_from_github(f"ukb/{image_id}/{image_id}_sax.nii.gz")
+            return sax_image.GetSize()[2]
+
+        def update_slice_slider(image_id):
+            num_slices = get_num_slices(image_id)
+            return gr.update(maximum=num_slices - 1, value=0, visible=True)
+
+        def fn(image_id, slice_idx):
+            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")
+            fig = plot_cmr_views(
+                lax_2c_image,
+                lax_3c_image,
+                lax_4c_image,
+                sax_image,
+                t_to_show=4,
+                depth_to_show=slice_idx,
+            )
+            fig.update_layout(height=600)
+            return fig
+
+        # When image changes, update the slice slider and plot
+        gr.on(
+            fn=lambda image_id: [update_slice_slider(image_id), fn(image_id, 0)],
+            inputs=[image_id],
+            outputs=[slice_idx, cmr_plot],
+        )
+
+        # When slice changes, update the plot
+        slice_idx.change(
+            fn=fn,
+            inputs=[image_id, slice_idx],
+            outputs=[cmr_plot],
+        )
+
+    return sax_interface
+
+
 @spaces.GPU
 def segmentation_sax_inference(
     images: torch.Tensor,
@@ -227,12 +316,6 @@ def segmentation_lax(seed, image_id, progress=gr.Progress()):
 
     # Download and load model
     progress(0, desc="Downloading model...")
-    image_url = f"https://raw.githubusercontent.com/mathpluscode/CineMA/main/cinema/examples/data/ukb/{image_id}/{image_id}_{view}.nii.gz"
-    image_path = cache_dir / f"{image_id}_{view}.nii.gz"
-    response = requests.get(image_url)
-    with open(image_path, "wb") as f:
-        f.write(response.content)
-
     model = ConvUNetR.from_finetuned(
         repo_id="mathpluscode/CineMA",
         model_filename=f"finetuned/segmentation/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
@@ -244,7 +327,11 @@ def segmentation_lax(seed, image_id, progress=gr.Progress()):
     progress(0, desc="Downloading data...")
     transform = ScaleIntensityd(keys=view)
 
-    images = np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(image_path)))
+    images = np.transpose(
+        sitk.GetArrayFromImage(
+            load_nifti_from_github(f"ukb/{image_id}/{image_id}_{view}.nii.gz")
+        )
+    )
     labels = segmentation_lax_inference(images, view, transform, model, progress)
 
     progress(1, desc="Plotting results...")
@@ -329,6 +416,8 @@ with gr.Blocks(
     )
 
     with gr.Tabs() as tabs:
+        with gr.TabItem("Cine CMR Views"):
+            cmr_tab()
         with gr.TabItem("Segmentation in SAX View"):
             segmentation_sax_tab()
         with gr.TabItem("Segmentation in LAX View"):