Formatting and examples

app.py

@@ -13,22 +13,26 @@ from monai.transforms import (
     ScaleIntensityd,
 )
 
+# Define the bundle name and path for downloading
 BUNDLE_NAME = 'spleen_ct_segmentation_v0.1.0'
 BUNDLE_PATH = os.path.join(torch.hub.get_dir(), 'bundle', BUNDLE_NAME)
 
-title = "Segment Brain Tumors with MONAI!"
+# Title and description
+title = "# Segment Brain Tumors with MONAI! 🧠"
 description = """
-## Brain Tumor Segmentation 🧠
-A pre-trained model for volumetric (3D) segmentation of brain tumor subregions from multimodal MRIs based on BraTS 2018 data.
-
 ## To run 🚀
 
-Upload a image file in the format: 4 channel MRI (4 aligned MRIs T1c, T1, T2, FLAIR at 1x1x1 mm)
+Upload a image file in the format: 4 channel MRI (4 aligned MRIs T1c, T1, T2, FLAIR at 1x1x1 mm), or try out one of the examples below! 
+If you want to see a different slice, update the slider and click the button.
+
+More details on the model can be found [here!](https://huggingface.co/katielink/brats_mri_segmentation_v0.1.0)
 
 ## Disclaimer ⚠️
 
 This is an example, not to be used for diagnostic purposes.
+"""
 
+references = """
 ## References 👀
 
 1. Myronenko, Andriy. "3D MRI brain tumor segmentation using autoencoder regularization." International MICCAI Brainlesion Workshop. Springer, Cham, 2018. https://arxiv.org/abs/1810.11654.
@@ -36,8 +40,12 @@ This is an example, not to be used for diagnostic purposes.
 3. Bakas S, et al. "Advancing The Cancer Genome Atlas glioma MRI collections with expert segmentation labels and radiomic features", Nature Scientific Data, 4:170117 (2017) DOI:10.1038/sdata.2017.117
 """
 
-#examples = 'examples/'
+examples = [
+    ['examples/BRATS_485.nii.gz', 100],
+    ['examples/BRATS_', 100]
+]
 
+# Load the MONAI pretrained model from Hugging Face Hub
 model, _, _ = bundle.load(
     name = BUNDLE_NAME,
     source = 'huggingface_hub',
@@ -45,10 +53,13 @@ model, _, _ = bundle.load(
     load_ts_module=True,
 )
 
+# Use GPU if available
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
+# Load the parser from the MONAI bundle's inference config
 parser = bundle.load_bundle_config(BUNDLE_PATH, 'inference.json')
 
+# Compose the preprocessing transforms
 preproc_transforms = Compose(
     [
         LoadImaged(keys=["image"]),
@@ -57,7 +68,14 @@ preproc_transforms = Compose(
         NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
     ]
 )
-inferer = parser.get_parsed_content('inferer', lazy=True, eval_expr=True, instantiate=True)
+
+# Get the inferer from the bundle's inference config
+inferer = parser.get_parsed_content(
+    'inferer',
+    lazy=True, eval_expr=True, instantiate=True
+)
+
+# Compose the postprocessing transforms
 post_transforms = Compose(
     [
         Activationsd(keys='pred', sigmoid=True),
@@ -66,10 +84,13 @@ post_transforms = Compose(
     ]
 )
 
+# Define the predict function for the demo
 def predict(input_file, z_axis, model=model, device=device):
+    # Load and process data in MONAI format
     data = {'image': [input_file.name]}
     data = preproc_transforms(data)
     
+    # Run inference and post-process predicted labels
     model.to(device)
     model.eval()
     with torch.no_grad():
@@ -77,34 +98,57 @@ def predict(input_file, z_axis, model=model, device=device):
         data['pred'] = inferer(inputs=inputs[None,...], network=model)
         data = post_transforms(data)
     
-    input_image = data['image'].numpy()
-    pred_image = data['pred'].cpu().detach().numpy()
+    # Convert tensors back to numpy arrays
+    data['image'] = data['image'].numpy()
+    data['pred'] = data['pred'].cpu().detach().numpy()
     
-    input_t1c_image = input_image[0, :, :, z_axis]
-    #input_t1_image = input_image[1, :, :, z_axis]
-    #input_t2_image = input_image[2, :, :, z_axis]
-    #input_flair_image = input_image[3, :, :, z_axis]
+    # Magnetic resonance imaging sequences
+    t1c = data['image'][0, :, :, z_axis]  # T1-weighted, post contrast
+    t1 = data['image'][1, :, :, z_axis]  # T1-weighted, pre contrast
+    t2 = data['image'][2, :, :, z_axis]  # T2-weighted
+    flair = data['image'][3, :, :, z_axis]  # FLAIR
     
-    pred_tc_image = pred_image[0, 0, :, :, z_axis]
-    #pred_et_image = pred_image[0, 1, :, :, z_axis]
-    #pred_wt_image = pred_image[0, 2, :, :, z_axis]
+    # BraTS labels
+    tc = data['pred'][0, 0, :, :, z_axis]  # Tumor core
+    wt = data['pred'][0, 1, :, :, z_axis]  # Whole tumor
+    et = data['pred'][0, 2, :, :, z_axis]  # Enhancing tumor
     
-    return input_t1c_image, pred_tc_image, 
-
-
-iface = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.File(label='Input file'),
-        gr.Slider(0, 200, label='z-axis', value=100)
-    ], 
-    outputs=[
-        gr.Image(label='T1C image'),
-        gr.Image(label='Segmentation'),
-    ],
-    title=title,
-    description=description,
-    #examples=examples,
-)
+    return [t1c, t1, t2, flair], [tc, wt, et] 
+
+# Use blocks to set up a more complex demo
+with gr.Blocks() as demo:
+
+    # Show title and description
+    gr.Markdown(title)
+    gr.Markdown(description)
+
+    # Get the input file and slice slider as inputs
+    input_file = gr.File(label='input file')
+    z_axis = gr.Slider(0, 200, label='z-axis', value=50)
+    
+    # Show the button with custom label
+    button = gr.Button("Segment Tumor!")
+
+    # Show examples for the user to try
+    gr.Markdown("Try some examples from MONAI's Decathlon Dataset:")
+    examples = gr.Examples(
+        examples=examples,
+        inputs=[gr.File(), gr.Slider()]
+    )
+    
+    # Show the input image with different MR sequences
+    input_image = gr.Gallery(label='input MRI sequences (T1+, T1, T2, FLAIR)')
+    output_segmentation = gr.Gallery(label='output segmentations (TC, EC, WT)')
+    
+    # Run prediction on button click
+    button.click(
+        predict, 
+        inputs=[input_file, z_axis],
+        outputs=[input_image, output_segmentation]
+    )
+    
+    # Show references at the bottom of the demo
+    gr.Markdown(references)
 
-iface.launch()
+# Launch the demo 
+demo.launch()

examples/BRATS_486.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8957d67a50b39afd8210f3ca51a20c77ef1c92642800f91b50f16b27778f2b2
+size 11111216

examples/log.csv

@@ -1,2 +0,0 @@
-input_file
-BRATS_485.nii.gz