Merge pull request #17 from andreped/demo-hf

Dockerfile

@@ -65,4 +65,4 @@ RUN rm -r *.zip
 RUN wget "https://github.com/andreped/neukit/releases/download/test-data/test_thorax_CT.nii.gz"
 
 # CMD ["/bin/bash"]
-CMD ["python3", "app.py"]
+CMD ["python3", "demo/app.py"]

README.md

@@ -7,7 +7,7 @@ app_port: 7860
 emoji: 🧠
 pinned: false
 license: mit
-app_file: app.py
+app_file: demo/app.py
 ---
 
 <div align="center">
@@ -58,7 +58,7 @@ virtualenv -python3 venv --clear
 source venv/bin/activate
 pip install -r ./demo/requirements.txt
 
-python app.py --cwd ./
+python demo/app.py --cwd ./
 ```
 
 ## [Troubleshooting](https://github.com/raidionics/AeroPath#troubleshooting)

demo/src/gui.py

@@ -20,7 +20,7 @@ class WebUI:
         self.pred_images = []
 
         # @TODO: This should be dynamically set based on chosen volume size
-        self.nb_slider_items = 300
+        self.nb_slider_items = 415
 
         self.model_name = model_name
         self.cwd = cwd
@@ -29,10 +29,12 @@ class WebUI:
         self.class_name = "airways"  # default
         self.class_names = {
             "airways": "CT_Airways",
+            "lungs": "CT_Lungs",
         }
 
         self.result_names = {
-            "airways": "Airway",
+            "airways": "Airways",
+            "lungs": "Lungs",
         }
 
         # define widgets not to be rendered immediantly, but later on
@@ -104,8 +106,7 @@ class WebUI:
                 model_selector = gr.Dropdown(
                     list(self.class_names.keys()),
                     label="Task",
-                    info="Which task to perform - one model for"
-                    "each brain tumor type and brain extraction",
+                    info="Which task to perform",
                     multiselect=False,
                     size="sm",
                 )

demo/src/inference.py

@@ -2,6 +2,7 @@ import configparser
 import logging
 import os
 import shutil
+import traceback
 
 
 def run_model(
@@ -29,6 +30,8 @@ def run_model(
     if os.path.exists("./result/"):
         shutil.rmtree("./result/")
 
+    patient_directory = ''
+    output_path = ''
     try:
         # setup temporary patient directory
         filename = input_path.split("/")[-1]
@@ -86,12 +89,15 @@ def run_model(
             + ".nii.gz",
             "./prediction.nii.gz",
         )
-
+        # Clean-up
+        if os.path.exists(patient_directory):
+            shutil.rmtree(patient_directory)
+        if os.path.exists(output_path):
+            shutil.rmtree(output_path)
     except Exception as e:
-        print(e)
-
-    # Clean-up
-    if os.path.exists(patient_directory):
-        shutil.rmtree(patient_directory)
-    if os.path.exists(output_path):
-        shutil.rmtree(output_path)
+        print(traceback.format_exc())
+        # Clean-up
+        if os.path.exists(patient_directory):
+            shutil.rmtree(patient_directory)
+        if os.path.exists(output_path):
+            shutil.rmtree(output_path)