Added support to download prediction

demo/src/gui.py

@@ -59,7 +59,8 @@ class WebUI:
             visible=True,
             elem_id="model-3d",
             camera_position=[90, 180, 768],
-        ).style(height=512)
+            height=512,
+        )
 
     def set_class_name(self, value):
         LOGGER.info(f"Changed task to: {value}")
@@ -75,22 +76,31 @@ class WebUI:
 
     def process(self, mesh_file_name):
         path = mesh_file_name.name
+        curr = path.split("/")[-1]
+        self.extension = ".".join(curr.split(".")[1:])
+        self.filename = curr.split(".")[0] + "-" + self.class_names[self.class_name]
         run_model(
             path,
             model_path=os.path.join(self.cwd, "resources/models/"),
             task=self.class_names[self.class_name],
             name=self.result_names[self.class_name],
+            output_filename=self.filename + "." + self.extension
         )
         LOGGER.info("Converting prediction NIfTI to OBJ...")
-        nifti_to_obj("prediction.nii.gz")
+        nifti_to_obj(path=self.filename + "." + self.extension)
 
         LOGGER.info("Loading CT to numpy...")
         self.images = load_ct_to_numpy(path)
 
         LOGGER.info("Loading prediction volume to numpy..")
-        self.pred_images = load_pred_volume_to_numpy("./prediction.nii.gz")
+        self.pred_images = load_pred_volume_to_numpy(self.filename + "." + self.extension)
 
         return "./prediction.obj"
+    
+    def download_prediction(self):
+        if (not self.filename) or (not self.extension):
+            LOGGER.error("The prediction is not available or ready to download. Wait until the result is available in the 3D viewer.")
+        return self.filename + "." + self.extension
 
     def get_img_pred_pair(self, k):
         k = int(k)
@@ -98,7 +108,6 @@ class WebUI:
             self.combine_ct_and_seg(self.images[k], self.pred_images[k]),
             visible=True,
             elem_id="model-2d",
-        ).style(
             color_map={self.class_name: "#ffae00"},
             height=512,
             width=512,
@@ -122,9 +131,7 @@ class WebUI:
                         autoscroll=True,
                         elem_id="logs",
                         show_copy_button=True,
-                        scroll_to_output=False,
                         container=True,
-                        line_breaks=True,
                     )
                     demo.load(read_logs, None, logs, every=1)
 
@@ -160,7 +167,7 @@ class WebUI:
                             label="Task",
                             info="Which structure to segment.",
                             multiselect=False,
-                            size="sm",
+                            scale=1.0,
                         )
                         model_selector.input(
                             fn=lambda x: self.set_class_name(x),
@@ -173,16 +180,33 @@ class WebUI:
                                 "Run analysis",
                                 variant="primary",
                                 elem_id="run-button",
-                            ).style(
-                                full_width=False,
+                                #scale=1.0,
+                                # size=1.0,
                                 size="lg",
                             )
+                            #.style(
+                            #    full_width=False,
+                            #    size="lg",
+                            #)
                             run_btn.click(
                                 fn=lambda x: self.process(x),
                                 inputs=file_output,
                                 outputs=self.volume_renderer,
                             )
 
+                            download_btn = gr.DownloadButton(
+                                "Download the result as NIfTI",
+                                visible=True,
+                                variant="secondary",
+                                # scale=1.0,
+                                size="sm",
+                            )
+                            download_btn.click(
+                                fn=self.download_prediction,
+                                inputs=None,
+                                outputs=download_btn,
+                            )
+
                     with gr.Row():
                         gr.Examples(
                             examples=[
@@ -202,16 +226,17 @@ class WebUI:
                         )
 
                     with gr.Row():
-                        with gr.Box():
+                        with gr.Group(): #gr.Box():
                             with gr.Column():
                                 # create dummy image to be replaced by loaded images
                                 t = gr.AnnotatedImage(
                                     visible=True, elem_id="model-2d"
-                                ).style(
-                                    color_map={self.class_name: "#ffae00"},
-                                    height=512,
-                                    width=512,
                                 )
+                                #.style(
+                                #    color_map={self.class_name: "#ffae00"},
+                                #    height=512,
+                                #    width=512,
+                                #)
 
                                 self.slider.input(
                                     self.get_img_pred_pair,
@@ -221,7 +246,7 @@ class WebUI:
 
                                 self.slider.render()
 
-                        with gr.Box():
+                        with gr.Group(): #gr.Box():
                             self.volume_renderer.render()
 
         # sharing app publicly -> share=True:

demo/src/inference.py

@@ -11,6 +11,7 @@ def run_model(
     verbose: str = "info",
     task: str = "CT_Airways",
     name: str = "Airways",
+    output_filename: str = None,
 ):
     if verbose == "debug":
         logging.getLogger().setLevel(logging.DEBUG)
@@ -27,6 +28,9 @@ def run_model(
     if os.path.exists("./result/"):
         shutil.rmtree("./result/")
 
+    if output_filename is None:
+        raise ValueError("Please, set output_filename.")
+
     patient_directory = ""
     output_path = ""
     try:
@@ -84,7 +88,7 @@ def run_model(
             + "-t1gd_annotation-"
             + name
             + ".nii.gz",
-            "./prediction.nii.gz",
+            output_filename,
         )
         # Clean-up
         if os.path.exists(patient_directory):