Comments

app.py

@@ -19,9 +19,11 @@ from swinunetr import SwinUnetrModelForInference, SwinUnetrConfig
 ffmpeg_path = shutil.which('ffmpeg')
 mediapy.set_ffmpeg(ffmpeg_path)
 
+# Load model
 model = SwinUnetrModelForInference.from_pretrained('darragh/swinunetr-btcv-tiny')
 model.eval()
 
+# Pull files from github
 input_files = glob.glob('pmrc/SwinUNETR/BTCV/dataset/imagesSampleTs/*.nii.gz')
 input_files = dict((f.split('/')[-1], f) for f in input_files)
 
@@ -43,8 +45,6 @@ test_transform = transforms.Compose(
         transforms.ToTensord(keys=["image"]),
     ])
 
-
-
 # Create Data Loader
 def create_dl(test_files):
     ds = test_transform(test_files)
@@ -56,14 +56,17 @@ def create_dl(test_files):
 # Inference and video generation
 def generate_dicom_video(selected_file, n_frames):
     
+    # Data processor
     test_file = input_files[selected_file]
     test_files = [{'image': test_file}]
     dl = create_dl(test_files)
     batch = next(iter(dl))
     
+    # Select dicom slices
     tst_inputs = batch["image"]
-    tst_inputs = tst_inputs[:,:,:,:,-32:]
-    
+    tst_inputs = tst_inputs[:,:,:,:,-n_frames:]
+
+    # Inference    
     with torch.no_grad():
         outputs = model(tst_inputs,
                             (96,96,96),
@@ -75,7 +78,6 @@ def generate_dicom_video(selected_file, n_frames):
     
     # Write frames to video
     for inp, outp in zip(tst_inputs, tst_outputs):
-
         frames = []
         for idx in range(inp.shape[-1]):
             # Segmentation
@@ -96,7 +98,7 @@ theme = 'dark-peach'
 with gr.Blocks(theme=theme) as demo:
 
     gr.Markdown('''<center><h1>SwinUnetr BTCV</h1></center>
-	This is a Gradio Blocks app of the winning transformer in the Beyond the Cranial Vault (BTCV) Segmentation Challenge, <a href="https://github.com/darraghdog/Project-MONAI-research-contributions/tree/main/SwinUNETR/BTCV">SwinUnetr (tiny version).</a>.
+	This is a Gradio Blocks app of the winning transformer in the Beyond the Cranial Vault (BTCV) Segmentation Challenge, <a href="https://github.com/darraghdog/Project-MONAI-research-contributions/tree/main/SwinUNETR/BTCV">SwinUnetr</a> (tiny version).
 	''')
     selected_dicom_key = gr.inputs.Dropdown(
             choices=sorted(input_files),
@@ -105,9 +107,6 @@ with gr.Blocks(theme=theme) as demo:
     n_frames = gr.Slider(1, 100, value=32, label="Number of dicom slices")
     button_gen_video = gr.Button("Generate Video")
     output_interpolation = gr.Video(label="Generated Video")
-    gr.Markdown(
-            '<center><img src="https://visitor-badge.glitch.me/badge?page_id=gradio-blocks.anime-biggan" alt="visitor badge"/></center>'
-        )
     button_gen_video.click(fn=generate_dicom_video, 
                            inputs=[selected_dicom_key, n_frames], 
                            outputs=output_interpolation)