Added support for all preop brain tumor types

neukit/gui.py

@@ -4,7 +4,7 @@ from .inference import run_model
 
 
 class WebUI:
-    def __init__(self, model_name:str = None, class_name:str = "meningioma", cwd:str = "/home/user/app/"):
+    def __init__(self, model_name:str = None, cwd:str = "/home/user/app/"):
         # global states
         self.images = []
         self.pred_images = []
@@ -13,9 +13,25 @@ class WebUI:
         self.nb_slider_items = 150
 
         self.model_name = model_name
-        self.class_name = class_name
         self.cwd = cwd
 
+        self.class_name = "meningioma"  # default - but can be updated based on which task is chosen from dropdown
+        self.class_names = {
+            "meningioma": "MRI_Meningioma",
+            "low-grade": "MRI_LGGlioma",
+            "metastasis": "MRI_Metastasis",
+            "high-grade": "MRI_GBM",
+            "brain": "MRI_Brain",
+        }
+
+        self.result_names = {
+            "meningioma": "Tumor",
+            "low-grade": "Tumor",
+            "metastasis": "Tumor",
+            "high-grade": "Tumor",
+            "brain": "Brain",
+        }
+
         # define widgets not to be rendered immediantly, but later on
         self.slider = gr.Slider(1, self.nb_slider_items, value=1, step=1, label="Which 2D slice to show")
         self.volume_renderer = gr.Model3D(
@@ -24,6 +40,10 @@ class WebUI:
             visible=True,
             elem_id="model-3d",
         ).style(height=512)
+    
+    def set_class_name(self, value):
+        print("Changed task to:", value)
+        self.class_name = value
 
     def combine_ct_and_seg(self, img, pred):
         return (img, [(pred, self.class_name)])
@@ -33,7 +53,7 @@ class WebUI:
     
     def load_mesh(self, mesh_file_name):
         path = mesh_file_name.name
-        run_model(path, model_path=self.cwd + "resources/models/")
+        run_model(path, model_path=self.cwd + "resources/models/", task=self.class_names[self.class_name], name=self.result_names[self.class_name])
         nifti_to_glb("prediction.nii.gz")
 
         self.images = load_ct_to_numpy(path)
@@ -56,25 +76,42 @@ class WebUI:
         height: 512px;
         margin: auto;
         }
+        #upload {
+        height: 120px;
+        }
         """
         with gr.Blocks(css=css) as demo:
 
             with gr.Row():
-                file_output = gr.File(
-                    file_count="single"
-                ).style(full_width=False, size="sm")
+                
+                file_output = gr.File(file_count="single", elem_id="upload")  # elem_id="upload"
                 file_output.upload(self.upload_file, file_output, file_output)
 
-                run_btn = gr.Button("Run analysis").style(full_width=False, size="sm")
+                # with gr.Column():
+
+                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",
+                    multiselect=False,
+                    size="sm",
+                )
+                model_selector.input(
+                    fn=lambda x: self.set_class_name(x),
+                    inputs=model_selector,
+                    outputs=None,
+                )
+
+                run_btn = gr.Button("Run analysis").style(full_width=False, size="lg")
                 run_btn.click(
                     fn=lambda x: self.load_mesh(x),
                     inputs=file_output,
-                    outputs=self.volume_renderer
+                    outputs=self.volume_renderer,
                 )
             
             with gr.Row():
                 gr.Examples(
-                    examples=[self.cwd + "RegLib_C01_2.nii"],
+                    examples=[self.cwd + "RegLib_C01_1.nii", self.cwd + "RegLib_C01_2.nii"],
                     inputs=file_output,
                     outputs=file_output,
                     fn=self.upload_file,

neukit/inference.py

@@ -5,7 +5,7 @@ import logging
 import traceback
 
 
-def run_model(input_path: str, model_path: str, verbose: str = "info", task: str = "MRI_Meningioma"):
+def run_model(input_path: str, model_path: str, verbose: str = "info", task: str = "MRI_Meningioma", name: str = "Tumor"):
     logging.basicConfig()
     logging.getLogger().setLevel(logging.WARNING)
 
@@ -17,49 +17,56 @@ def run_model(input_path: str, model_path: str, verbose: str = "info", task: str
         logging.getLogger().setLevel(logging.ERROR)
     else:
         raise ValueError("Unsupported verbose value provided:", verbose)
-
-    # create sequence folder, rename patient, and add to temporary patient directory
-    filename = input_path.split("/")[-1]
-    splits = filename.split(".")
-    extension = ".".join(splits[1:])
-    patient_directory = "./patient/"
-    os.makedirs(patient_directory + "T0/", exist_ok=True)
-    shutil.copy(input_path, patient_directory + "T0/" + splits[0] + "-t1gd." + extension)
     
-    # define output directory to save results
-    output_path = "./result/prediction-" + splits[0] + "/"
-    os.makedirs(output_path, exist_ok=True)
+    # delete patient/result folder if they exist
+    if os.path.exists("./patient/"):
+        shutil.rmtree("./patient/")
+    if os.path.exists("./result/"):
+        shutil.rmtree("./result/")
+
+    try:
+        # create sequence folder, rename patient, and add to temporary patient directory
+        filename = input_path.split("/")[-1]
+        splits = filename.split(".")
+        extension = ".".join(splits[1:])
+        patient_directory = "./patient/"
+        os.makedirs(patient_directory + "T0/", exist_ok=True)
+        shutil.copy(input_path, patient_directory + "T0/" + splits[0] + "-t1gd." + extension)
+        
+        # define output directory to save results
+        output_path = "./result/prediction-" + splits[0] + "/"
+        os.makedirs(output_path, exist_ok=True)
 
-    # Setting up the configuration file
-    rads_config = configparser.ConfigParser()
-    rads_config.add_section('Default')
-    rads_config.set('Default', 'task', 'neuro_diagnosis')
-    rads_config.set('Default', 'caller', '')
-    rads_config.add_section('System')
-    rads_config.set('System', 'gpu_id', "-1")
-    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.add_section('Runtime')
-    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')
+        # Setting up the configuration file
+        rads_config = configparser.ConfigParser()
+        rads_config.add_section('Default')
+        rads_config.set('Default', 'task', 'neuro_diagnosis')
+        rads_config.set('Default', 'caller', '')
+        rads_config.add_section('System')
+        rads_config.set('System', 'gpu_id', "-1")
+        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.add_section('Runtime')
+        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)
+        with open("rads_config.ini", "w") as f:
+            rads_config.write(f)
 
-    # finally, run inference
-    from raidionicsrads.compute import run_rads
+        # finally, run inference
+        from raidionicsrads.compute import run_rads
 
-    try:
         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", "./prediction.nii.gz")
+    
     except Exception as e:
         print(e)
-    
-    # rename and move final result
-    os.rename("./result/prediction-" + splits[0] + "/T0/" + splits[0] + "-t1gd_annotation-Tumor.nii.gz", "./prediction.nii.gz")
-    
+
     # Clean-up
     if os.path.exists(patient_directory):
         shutil.rmtree(patient_directory)