Upgrade to gradio 5

.gitignore

@@ -13,3 +13,4 @@ venv/
 *.obj
 *.zip
 *.txt
+*.idea/

Dockerfile

@@ -1,6 +1,6 @@
 # read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
 # you will also find guides on how best to write your Dockerfile
-FROM python:3.8-slim
+FROM python:3.10-slim
 
 # set language, format and stuff
 ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
@@ -50,10 +50,10 @@ WORKDIR $HOME/app
 COPY --chown=user . $HOME/app
 
 # Download pretrained models
-RUN wget "https://github.com/raidionics/Raidionics-models/releases/download/1.2.0/Raidionics-CT_Airways-ONNX-v12.zip" && \
-    unzip "Raidionics-CT_Airways-ONNX-v12.zip" && mkdir -p resources/models/ && mv CT_Airways/ resources/models/CT_Airways/
-RUN wget "https://github.com/raidionics/Raidionics-models/releases/download/1.2.0/Raidionics-CT_Lungs-ONNX-v12.zip" && \
-    unzip "Raidionics-CT_Lungs-ONNX-v12.zip" && mv CT_Lungs/ resources/models/CT_Lungs/
+RUN wget "https://github.com/raidionics/Raidionics-models/releases/download/v1.3.0-rc/Raidionics-CT_Airways-v13.zip" && \
+    unzip "Raidionics-CT_Airways-v13.zip" && mkdir -p resources/models/ && mv CT_Airways/ resources/models/CT_Airways/
+RUN wget "https://github.com/raidionics/Raidionics-models/releases/download/v1.3.0-rc/Raidionics-CT_Lungs-v13.zip" && \
+    unzip "Raidionics-CT_Lungs-v13.zip" && mv CT_Lungs/ resources/models/CT_Lungs/
 
 RUN rm -r *.zip
 

demo/requirements.txt

@@ -1,2 +1,2 @@
-raidionicsrads@git+https://github.com/dbouget/raidionics_rads_lib
-gradio==4.29.0
+raidionicsrads
+gradio

demo/src/gui.py

@@ -1,6 +1,8 @@
 import os
 
 import gradio as gr
+from zipfile import ZipFile
+from PIL import Image
 
 from .convert import nifti_to_obj
 from .css_style import css
@@ -22,48 +24,39 @@ class WebUI:
         cwd: str = "/home/user/app/",
         share: int = 1,
     ):
+        self.file_output = None
+        self.model_selector = None
+        self.stripped_cb = None
+        self.registered_cb = None
+        self.run_btn = None
+        self.slider = None
+        self.download_file = None
+
         # global states
         self.images = []
         self.pred_images = []
 
-        # @TODO: This should be dynamically set based on chosen volume size
-        self.nb_slider_items = 820
-
         self.model_name = model_name
         self.cwd = cwd
         self.share = share
 
-        self.filename = None
-        self.extension = None
-
-        self.class_name = "airways"  # default
+        self.class_name = "Airways"  # default
         self.class_names = {
-            "airways": "CT_Airways",
-            "lungs": "CT_Lungs",
+            "Airways": "CT_Airways",
         }
 
         self.result_names = {
-            "airways": "Airways",
-            "lungs": "Lungs",
+            "Airways": "Airways",
         }
 
-        # define widgets not to be rendered immediantly, but later on
-        self.slider = gr.Slider(
-            minimum=1,
-            maximum=self.nb_slider_items,
-            value=1,
-            step=1,
-            label="Which 2D slice to show",
-        )
         self.volume_renderer = gr.Model3D(
             clear_color=[0.0, 0.0, 0.0, 0.0],
             label="3D Model",
-            show_label=True,
             visible=True,
             elem_id="model-3d",
-            camera_position=[90, 180, 768],
             height=512,
         )
+        # self.volume_renderer = ShinyModel3D()
 
     def set_class_name(self, value):
         LOGGER.info(f"Changed task to: {value}")
@@ -79,75 +72,107 @@ 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(path=self.filename + "." + self.extension)
+        nifti_to_obj("prediction.nii.gz")
 
         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(
-            self.filename + "." + self.extension
-        )
+        self.pred_images = load_pred_volume_to_numpy("./prediction.nii.gz")
 
-        return "./prediction.obj"
+        slider = gr.Slider(
+            minimum=0,
+            maximum=len(self.images) - 1,
+            value=int(len(self.images) / 2),
+            step=1,
+            label="Which 2D slice to show",
+            interactive=True,
+        )
 
-    def download_prediction(self):
-        if (self.filename is None) or (self.extension is None):
-            LOGGER.error(
-                "The prediction is not available or ready to download. Wait until the result is available in the 3D viewer."
-            )
-            raise ValueError("Run inference before downloading!")
-        return self.filename + "." + self.extension
+        return "./prediction.obj", slider
 
     def get_img_pred_pair(self, k):
-        k = int(k)
-        out = gr.AnnotatedImage(
-            self.combine_ct_and_seg(self.images[k], self.pred_images[k]),
-            visible=True,
-            elem_id="model-2d",
-            color_map={self.class_name: "#ffae00"},
-            height=512,
-            width=512,
-        )
-        return out
+        img = self.images[k]
+        img_pil = Image.fromarray(img)
+        seg_list = []
+        seg_list.append((self.pred_images[k], self.class_name))
+        return img_pil, seg_list
 
     def toggle_sidebar(self, state):
         state = not state
         return gr.update(visible=state), state
 
+    def package_results(self):
+        """Generates text files and zips them."""
+        output_dir = "temp_output"
+        os.makedirs(output_dir, exist_ok=True)
+
+        zip_filename = os.path.join(output_dir, "generated_files.zip")
+        with ZipFile(zip_filename, 'w') as zf:
+            zf.write("./prediction.nii.gz")
+
+        return zip_filename
+
+    def setup_interface_outputs(self):
+        with gr.Row():
+            with gr.Group():
+                with gr.Column(scale=2):
+                    t = gr.AnnotatedImage(
+                        visible=True,
+                        elem_id="model-2d",
+                        color_map={self.class_name: "#ffae00"},
+                        height=512,
+                        width=512,
+                    )
+
+                    self.slider = gr.Slider(
+                        minimum=0,
+                        maximum=1,
+                        value=0,
+                        step=1,
+                        label="Which 2D slice to show",
+                        interactive=True,
+                    )
+
+                    self.slider.change(fn=self.get_img_pred_pair, inputs=self.slider, outputs=t)
+
+            with gr.Group():
+                self.volume_renderer.render()
+                self.download_btn = gr.DownloadButton(label="Download results", visible=False)
+                self.download_file = gr.File(label="Download Zip", interactive=True, visible=False)
+
+
     def run(self):
         with gr.Blocks(css=css) as demo:
             with gr.Row():
-                with gr.Column(visible=True, scale=0.2) as sidebar_left:
+                with gr.Column(scale=1, visible=True) as sidebar_left:
                     logs = gr.Textbox(
                         placeholder="\n" * 16,
                         label="Logs",
                         info="Verbose from inference will be displayed below.",
-                        lines=36,
-                        max_lines=36,
+                        lines=38,
+                        max_lines=38,
                         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)
+                    timer = gr.Timer(value=1, active=True)
+                    timer.tick(fn=read_logs, inputs=None, outputs=logs)
+                    # demo.load(read_logs, None, logs, every=0.5)
 
-                with gr.Column():
+                with gr.Column(scale=2):
                     with gr.Row():
-                        with gr.Column(scale=1, min_width=150):
+                        with gr.Column(min_width=150):
                             sidebar_state = gr.State(True)
 
                             btn_toggle_sidebar = gr.Button(
@@ -160,66 +185,30 @@ class WebUI:
                                 [sidebar_left, sidebar_state],
                             )
 
-                            btn_clear_logs = gr.Button(
-                                "Clear logs", elem_id="logs-button"
-                            )
+                            btn_clear_logs = gr.Button("Clear logs", elem_id="logs-button")
                             btn_clear_logs.click(flush_logs, [], [])
 
-                        file_output = gr.File(
-                            file_count="single",
-                            elem_id="upload",
-                            scale=3,
-                        )
-                        file_output.upload(
-                            self.upload_file, file_output, file_output
-                        )
+                        self.file_output = gr.File(file_count="single", elem_id="upload")
 
-                        model_selector = gr.Dropdown(
+                        self.model_selector = gr.Dropdown(
                             list(self.class_names.keys()),
                             label="Task",
                             info="Which structure to segment.",
                             multiselect=False,
-                            scale=1,
-                        )
-                        model_selector.input(
-                            fn=lambda x: self.set_class_name(x),
-                            inputs=model_selector,
-                            outputs=None,
                         )
 
-                        with gr.Column(scale=1, min_width=150):
-                            run_btn = gr.Button(
-                                "Run analysis",
-                                variant="primary",
-                                elem_id="run-button",
-                            )
-                            run_btn.click(
-                                fn=lambda x: self.process(x),
-                                inputs=file_output,
-                                outputs=self.volume_renderer,
-                            )
-
-                            download_btn = gr.DownloadButton(
-                                "Download prediction",
-                                visible=True,
-                                variant="secondary",
-                                elem_id="download",
-                            )
-                            download_btn.click(
-                                fn=self.download_prediction,
-                                inputs=None,
-                                outputs=download_btn,
-                            )
+                        with gr.Column(min_width=150):
+                            self.run_btn = gr.Button("Run segmentation", variant="primary", elem_id="run-button")
 
                     with gr.Row():
                         gr.Examples(
                             examples=[
                                 os.path.join(self.cwd, "test_thorax_CT.nii.gz"),
                             ],
-                            inputs=file_output,
-                            outputs=file_output,
+                            inputs=self.file_output,
+                            outputs=self.file_output,
                             fn=self.upload_file,
-                            cache_examples=True,
+                            cache_examples=False,
                         )
 
                         gr.Markdown(
@@ -229,32 +218,19 @@ class WebUI:
                             """
                         )
 
-                    with gr.Row():
-                        with gr.Group():
-                            with gr.Column():
-                                # create dummy image to be replaced by loaded images
-                                t = gr.AnnotatedImage(
-                                    visible=True,
-                                    elem_id="model-2d",
-                                    color_map={self.class_name: "#ffae00"},
-                                    # height=512,
-                                    # width=512,
-                                )
-                                self.slider.input(
-                                    self.get_img_pred_pair,
-                                    self.slider,
-                                    t,
-                                )
-
-                                self.slider.render()
-
-                        with gr.Group():  # gr.Box():
-                            self.volume_renderer.render()
-
+                    self.setup_interface_outputs()
+
+            # Define the signals/slots
+            self.file_output.upload(self.upload_file, self.file_output, self.file_output)
+            self.model_selector.input(fn=lambda x: self.set_class_name(x), inputs=self.model_selector, outputs=None)
+            self.run_btn.click(fn=self.process, inputs=[self.file_output],
+                               outputs=[self.volume_renderer, self.slider]).then(fn=lambda:
+            gr.DownloadButton(visible=True), inputs=None, outputs=self.download_btn)
+            self.download_btn.click(fn=self.package_results, inputs=[], outputs=self.download_file).then(fn=lambda
+                file_path: gr.File(label="Download Zip", visible=True, value=file_path), inputs=self.download_file,
+                                    outputs=self.download_file)
         # sharing app publicly -> share=True:
         # https://gradio.app/sharing-your-app/
         # inference times > 60 seconds -> need queue():
         # https://github.com/tloen/alpaca-lora/issues/60#issuecomment-1510006062
-        demo.queue().launch(
-            server_name="0.0.0.0", server_port=7860, share=self.share
-        )
+        demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=self.share)

demo/src/inference.py

@@ -3,6 +3,8 @@ import logging
 import os
 import shutil
 import traceback
+import json
+import fnmatch
 
 
 def run_model(
@@ -11,7 +13,6 @@ 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)
@@ -28,9 +29,6 @@ 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:
@@ -59,37 +57,51 @@ def run_model(
         rads_config.set("System", "input_folder", patient_directory)
         rads_config.set("System", "output_folder", output_path)
         rads_config.set("System", "model_folder", model_path)
-        rads_config.set(
-            "System",
-            "pipeline_filename",
-            os.path.join(model_path, task, "pipeline.json"),
-        )
+        rads_config.set('System', 'pipeline_filename', os.path.join(output_path,
+                                                                    'test_pipeline.json'))
         rads_config.add_section("Runtime")
-        rads_config.set(
-            "Runtime", "reconstruction_method", "thresholding"
-        )  # thresholding, probabilities
+        rads_config.set("Runtime", "reconstruction_method", "thresholding")  # thresholding, probabilities
         rads_config.set("Runtime", "reconstruction_order", "resample_first")
         rads_config.set("Runtime", "use_preprocessed_data", "False")
 
         with open("rads_config.ini", "w") as f:
             rads_config.write(f)
 
+        pip = {}
+        step_index = 1
+        pip_num = str(step_index)
+        pip[pip_num] = {}
+        pip[pip_num]["task"] = "Classification"
+        pip[pip_num]["inputs"] = {}  # Empty input means running it on all existing data for the patient
+        pip[pip_num]["target"] = ["MRSequence"]
+        pip[pip_num]["model"] = "MRI_SequenceClassifier"
+        pip[pip_num]["description"] = "Classification of the MRI sequence type for all input scans."
+
+        step_index = step_index + 1
+        pip_num = str(step_index)
+        pip[pip_num] = {}
+        pip[pip_num]["task"] = 'Model selection'
+        pip[pip_num]["model"] = task
+        pip[pip_num]["timestamp"] = 0
+        pip[pip_num]["format"] = "thresholding"
+        pip[pip_num]["description"] = f"Identifying the best {task} segmentation model for existing inputs"
+
+        with open(os.path.join(output_path, 'test_pipeline.json'), 'w', newline='\n') as outfile:
+            json.dump(pip, outfile, indent=4, sort_keys=True)
+
         # finally, run inference
         from raidionicsrads.compute import run_rads
-
         run_rads(config_filename="rads_config.ini")
 
-        # rename and move final result
-        os.rename(
-            "./result/prediction-"
-            + splits[0]
-            + "/T0/"
-            + splits[0]
-            + "-t1gd_annotation-"
-            + name
-            + ".nii.gz",
-            output_filename,
-        )
+        logging.info(f"Looking for the following pattern: {task}")
+        patterns = [f"*-{name}.*"]
+        existing_files = os.listdir(os.path.join(output_path, "T0"))
+        logging.info(f"Existing files: {existing_files}")
+        fileName = str(os.path.join(output_path, "T0",
+                                    [x for x in existing_files if
+                                     any(fnmatch.fnmatch(x, pattern) for pattern in patterns)][0]))
+        os.rename(src=fileName, dst="./prediction.nii.gz")
+
         # Clean-up
         if os.path.exists(patient_directory):
             shutil.rmtree(patient_directory)