Formatting and examples

README.md

@@ -1,8 +1,8 @@
 ---
 title: Spleen Segmentation
-emoji: 👀
-colorFrom: indigo
-colorTo: red
+emoji: 🩸
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
 sdk_version: 3.1.1
 app_file: app.py

app.py

@@ -3,11 +3,17 @@ import gradio as gr
 import torch
 from monai import bundle
 
+# Set the bundle name and download path
 BUNDLE_NAME = 'spleen_ct_segmentation_v0.1.0'
 BUNDLE_PATH = os.path.join(torch.hub.get_dir(), 'bundle', BUNDLE_NAME)
 
-#examples = ['examples/spleen_46.nii.gz']
+# Set up some examples from the test set for better user experience
+examples = [
+    ['examples/spleen_1.nii.gz', 50],
+    ['examples/spleen_11.nii.gz', 50],
+]
 
+# Load the pretrained model from Hugging Face Hub
 model, _, _ = bundle.load(
     name = BUNDLE_NAME,
     source = 'huggingface_hub',
@@ -15,12 +21,21 @@ model, _, _ = bundle.load(
     load_ts_module=True,
 )
 
+# Use GPU if available
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
+# Load transforms and inferer directly from the bundle
 parser = bundle.load_bundle_config(BUNDLE_PATH, 'inference.json')
-preproc_transforms = parser.get_parsed_content('preprocessing', lazy=True, eval_expr=True, instantiate=True)
-inferer = parser.get_parsed_content('inferer', lazy=True, eval_expr=True, instantiate=True)
+preproc_transforms = parser.get_parsed_content(
+    'preprocessing',
+    lazy=True, eval_expr=True,instantiate=True
+)
+inferer = parser.get_parsed_content(
+    'inferer',
+    lazy=True, eval_expr=True, instantiate=True
+)
 
+# Define the prediction function
 def predict(input_file, z_axis, model=model, device=device):
     data = {'image': [input_file.name]}
     data = preproc_transforms(data)
@@ -36,15 +51,26 @@ def predict(input_file, z_axis, model=model, device=device):
     
     return input_image[0, :, :, z_axis], pred_image[0, :, :, z_axis]*255
 
+# Set up the demo interface
 iface = gr.Interface(
     fn=predict,
     inputs=[
-        gr.File(label='Nifti file'),
+        gr.File(label='input file'),
         gr.Slider(0, 200, label='z-axis', value=50)
     ], 
     outputs=['image', 'image'],
-    title='Segment the Spleen from a CT Scan using MONAI',
-    ##examples=examples,
+    title='Segment the Spleen using MONAI!',
+    description="""
+    ## 🚀 To run
+    Upload a abdominal CT scan, or try one of the examples below!
+    
+    More details on the model can be found [here!](https://huggingface.co/katielink/spleen_ct_segmentation_v0.1.0)
+
+    ## ⚠️ Disclaimer
+    Not to be used for diagnostic purposes.
+    """,
+    examples=examples,
 )
 
+# Launch the demo
 iface.launch()

examples/{spleen_46.nii.gz → spleen_1.nii.gz}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1dc87026596e782c25de81ab171cd2b456a8e2c76e304949e9318451406c8edf
-size 28610147
+oid sha256:a502f7e81d6b3742f5100e501b4e80180264d00e9aa04dfee43d821b1c6b07e7
+size 9862639

examples/spleen_11.nii.gz

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