Brain tumor multiclass segmentation (#62)

Multiclass segmentation outputs supported, added user preferences options, updated backend accordingly [skip ci]

.github/workflows/build_macos.yml

@@ -42,7 +42,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install git+https://github.com/dbouget/raidionics_rads_lib.git@software_test
+        pip install git+https://github.com/dbouget/raidionics_rads_lib.git@software
         pip install -r assets/requirements.txt
         pip install matplotlib==3.3.4
         pip install --force-reinstall --no-cache-dir pyside6

.github/workflows/build_ubuntu.yml

@@ -69,7 +69,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install git+https://github.com/dbouget/raidionics_rads_lib.git@software_test
+        pip install git+https://github.com/dbouget/raidionics_rads_lib.git@software
         pip install -r assets/requirements.txt
         pip uninstall -y PySide6 PySide6-Addons PySide6-Essentials
         pip install --force-reinstall --no-cache-dir pyside6==6.2.4

.github/workflows/build_windows.yml

@@ -41,7 +41,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         ls
-        pip install git+https://github.com/dbouget/raidionics_rads_lib.git@software_test
+        pip install git+https://github.com/dbouget/raidionics_rads_lib.git@software
         pip install -r assets/requirements.txt
         pip install matplotlib==3.3.4
         pip install --force-reinstall --no-cache-dir pyside6

gui/SinglePatientComponent/CentralAreaExecutionWidget.py

@@ -16,6 +16,7 @@
 from gui.UtilsWidgets.CustomQDialog.TumorTypeSelectionQDialog import TumorTypeSelectionQDialog
 from utils.data_structures.PatientParametersStructure import MRISequenceType
 from utils.data_structures.AnnotationStructure import AnnotationGenerationType, AnnotationClassType
+from utils.data_structures.UserPreferencesStructure import UserPreferencesStructure
 
 
 class CentralAreaExecutionWidget(QLabel):
@@ -159,6 +160,11 @@ def on_pipeline_execution(self, pipeline_code: str) -> None:
                 self.model_name = "MRI_Metastasis"
             elif diag.tumor_type == 'Meningioma':
                 self.model_name = "MRI_Meningioma"
+
+            if UserPreferencesStructure.getInstance().segmentation_tumor_model_type != "Tumor":
+                self.model_name = self.model_name + '_multiclass'
+                if diag.tumor_type == 'Low-Grade Glioma':
+                    self.model_name = "MRI_GBM_multiclass"
         elif "Brain" in pipeline_code:
             self.model_name = "MRI_Brain"
         elif "postop" in pipeline_code:

gui/StudyBatchComponent/StudiesSidePanel/SingleStudyWidget.py

@@ -14,6 +14,7 @@
 from gui.UtilsWidgets.CustomQDialog.TumorTypeSelectionQDialog import TumorTypeSelectionQDialog
 from gui.UtilsWidgets.CustomQDialog.SavePatientChangesDialog import SavePatientChangesDialog
 from utils.software_config import SoftwareConfigResources
+from utils.data_structures.UserPreferencesStructure import UserPreferencesStructure
 
 
 class SingleStudyWidget(QCollapsibleWidget):
@@ -599,6 +600,10 @@ def __on_run_pipeline(self) -> None:
                 self.model_name = "MRI_Metastasis"
             elif diag.tumor_type == 'Meningioma':
                 self.model_name = "MRI_Meningioma"
+            if UserPreferencesStructure.getInstance().segmentation_tumor_model_type != "Tumor":
+                self.model_name = self.model_name + '_multiclass'
+                if diag.tumor_type == 'Low-Grade Glioma':
+                    self.model_name = "MRI_GBM_multiclass"
 
         self.on_processing_started()
         self.batch_pipeline_execution_requested.emit(self.uid, pipeline_task, self.model_name)

gui/UtilsWidgets/CustomQDialog/SoftwareSettingsDialog.py

@@ -178,6 +178,23 @@ def __set_processing_segmentation_options_interface(self):
         self.processing_segmentation_options_base_layout = QVBoxLayout()
         self.processing_segmentation_options_label = QLabel("Processing - Segmentation")
         self.processing_segmentation_options_base_layout.addWidget(self.processing_segmentation_options_label)
+        self.processing_segmentation_models_groupbox = QGroupBox()
+        self.processing_segmentation_models_groupbox.setTitle("Segmentation models")
+        self.processing_segmentation_models_groupboxlayout = QVBoxLayout()
+
+        self.processing_options_segmentation_models_layout = QHBoxLayout()
+        self.processing_options_segmentation_models_label = QLabel("Output classes")
+        self.processing_options_segmentation_models_label.setToolTip("Select the segmented output classes desired in the drop-down menu. N-B: all four MR sequences (i.e., T1-CE, T1-w, FLAIR, and T2) are required as input for the second choice.\n")
+        self.processing_options_segmentation_models_selector_combobox = QComboBox()
+        self.processing_options_segmentation_models_selector_combobox.addItems(["Tumor", "Tumor, Necrosis, Edema"])
+        self.processing_options_segmentation_models_selector_combobox.setCurrentText(UserPreferencesStructure.getInstance().segmentation_tumor_model_type)
+        self.processing_options_segmentation_models_layout.addWidget(self.processing_options_segmentation_models_label)
+        self.processing_options_segmentation_models_layout.addWidget(self.processing_options_segmentation_models_selector_combobox)
+        self.processing_options_segmentation_models_layout.addStretch(1)
+        self.processing_segmentation_models_groupboxlayout.addLayout(self.processing_options_segmentation_models_layout)
+        self.processing_segmentation_models_groupbox.setLayout(self.processing_segmentation_models_groupboxlayout)
+        self.processing_segmentation_options_base_layout.addWidget(self.processing_segmentation_models_groupbox)
+
         self.processing_segmentation_refinement_groupbox = QGroupBox()
         self.processing_segmentation_refinement_groupbox.setTitle("Refinement")
         self.processing_segmentation_refinement_groupboxlayout = QVBoxLayout()
@@ -375,6 +392,7 @@ def __set_connections(self):
         self.processing_options_use_stripped_inputs_checkbox.stateChanged.connect(self.__on_use_stripped_inputs_status_changed)
         self.processing_options_use_registered_inputs_checkbox.stateChanged.connect(self.__on_use_registered_inputs_status_changed)
         self.processing_options_export_results_rtstruct_checkbox.stateChanged.connect(self.__on_export_results_rtstruct_status_changed)
+        self.processing_options_segmentation_models_selector_combobox.currentTextChanged.connect(self.__on_segmentation_model_type_changed)
         self.processing_options_segmentation_refinement_checkbox.stateChanged.connect(self.__on_perform_segmentation_refinement_status_changed)
         self.processing_options_segmentation_refinement_dilation_threshold_spinbox.valueChanged.connect(self.__on_perform_segmentation_refinement_dilation_value_changed)
         self.processing_options_compute_corticalstructures_checkbox.stateChanged.connect(self.__on_compute_corticalstructures_status_changed)
@@ -633,19 +651,22 @@ def __on_model_purge_clicked(self) -> None:
                 os.makedirs(SoftwareConfigResources.getInstance().models_path)
 
     def __on_use_sequences_status_changed(self, status):
-        UserPreferencesStructure.getInstance().use_manual_sequences = status
+        UserPreferencesStructure.getInstance().use_manual_sequences = self.processing_options_use_sequences_checkbox.isChecked()
 
     def __on_use_manual_annotations_status_changed(self, status):
-        UserPreferencesStructure.getInstance().use_manual_annotations = status
+        UserPreferencesStructure.getInstance().use_manual_annotations = self.processing_options_use_annotations_checkbox.isChecked()
 
     def __on_use_stripped_inputs_status_changed(self, status):
-        UserPreferencesStructure.getInstance().use_stripped_inputs = status
+        UserPreferencesStructure.getInstance().use_stripped_inputs = self.processing_options_use_stripped_inputs_checkbox.isChecked()
 
     def __on_use_registered_inputs_status_changed(self, status):
-        UserPreferencesStructure.getInstance().use_registered_inputs = status
+        UserPreferencesStructure.getInstance().use_registered_inputs = self.processing_options_use_registered_inputs_checkbox.isChecked()
 
     def __on_export_results_rtstruct_status_changed(self, status):
-        UserPreferencesStructure.getInstance().export_results_as_rtstruct = status
+        UserPreferencesStructure.getInstance().export_results_as_rtstruct = self.processing_options_export_results_rtstruct_checkbox.isChecked()
+
+    def __on_segmentation_model_type_changed(self, text):
+        UserPreferencesStructure.getInstance().segmentation_tumor_model_type = text
 
     def __on_perform_segmentation_refinement_status_changed(self, status):
         UserPreferencesStructure.getInstance().perform_segmentation_refinement = status
@@ -660,7 +681,7 @@ def __on_perform_segmentation_refinement_dilation_value_changed(self, value):
         UserPreferencesStructure.getInstance().segmentation_refinement_dilation_percentage = value
 
     def __on_compute_corticalstructures_status_changed(self, state):
-        UserPreferencesStructure.getInstance().compute_cortical_structures = state
+        UserPreferencesStructure.getInstance().compute_cortical_structures = self.processing_options_compute_corticalstructures_checkbox.isChecked()
         if state:
             self.corticalstructures_mni_checkbox.setEnabled(True)
             self.corticalstructures_mni_label.setEnabled(True)
@@ -725,7 +746,7 @@ def __on_corticalstructure_harvardoxford_status_changed(self, state):
         UserPreferencesStructure.getInstance().cortical_structures_list = structs
 
     def __on_compute_subcorticalstructures_status_changed(self, state):
-        UserPreferencesStructure.getInstance().compute_subcortical_structures = state
+        UserPreferencesStructure.getInstance().compute_subcortical_structures = self.processing_options_compute_subcorticalstructures_checkbox.isChecked()
         if state:
             self.subcorticalstructures_bcb_checkbox.setEnabled(True)
             self.subcorticalstructures_bcb_label.setEnabled(True)

utils/data_structures/AnnotationStructure.py

@@ -23,6 +23,8 @@ class AnnotationClassType(Enum):
 
     Brain = 0, 'Brain'
     Tumor = 1, 'Tumor'
+    Necrosis = 2, 'Necrosis'
+    Edema = 3, 'Edema'
 
     Lungs = 100, 'Lungs'
     Airways = 101, 'Airways'

utils/data_structures/UserPreferencesStructure.py

@@ -23,6 +23,7 @@ class UserPreferencesStructure:
     _export_results_as_rtstruct = False  # True to export all masks as DICOM RTStruct in addition
     _use_stripped_inputs = False  # True to use inputs already stripped (e.g., skull-stripped or lungs-stripped)
     _use_registered_inputs = False  # True to use inputs already registered (e.g., altas-registered, multi-sequences co-registered)
+    _segmentation_tumor_model_type = "Tumor"  # Type of output to expect from the tumor segmentation model (i.e., indicating if a BraTS model should be used)
     _perform_segmentation_refinement = False  # True to enable any kind of segmentation refinement
     _segmentation_refinement_type = "dilation"  # String indicating the type of refinement to perform, to select from ["dilation"]
     _segmentation_refinement_dilation_percentage = 0  # Integer indicating the volume percentage increase to reach after dilation
@@ -128,7 +129,7 @@ def use_manual_annotations(self) -> bool:
         return self._use_manual_annotations
 
     @use_manual_annotations.setter
-    def use_manual_annotations(self, state) -> None:
+    def use_manual_annotations(self, state: bool) -> None:
         self._use_manual_annotations = state
         self.save_preferences()
 
@@ -137,7 +138,7 @@ def use_stripped_inputs(self) -> bool:
         return self._use_stripped_inputs
 
     @use_stripped_inputs.setter
-    def use_stripped_inputs(self, state) -> None:
+    def use_stripped_inputs(self, state: bool) -> None:
         self._use_stripped_inputs = state
         self.save_preferences()
 
@@ -146,10 +147,19 @@ def use_registered_inputs(self) -> bool:
         return self._use_registered_inputs
 
     @use_registered_inputs.setter
-    def use_registered_inputs(self, state) -> None:
+    def use_registered_inputs(self, state: bool) -> None:
         self._use_registered_inputs = state
         self.save_preferences()
 
+    @property
+    def segmentation_tumor_model_type(self) -> str:
+        return self._segmentation_tumor_model_type
+
+    @segmentation_tumor_model_type.setter
+    def segmentation_tumor_model_type(self, model_type: str) -> None:
+        self._segmentation_tumor_model_type = model_type
+        self.save_preferences()
+
     @property
     def perform_segmentation_refinement(self) -> bool:
         return self._perform_segmentation_refinement
@@ -204,7 +214,11 @@ def subcortical_structures_list(self, structures: List[str]) -> None:
         self._subcortical_structures_list = structures
         self.save_preferences()
 
-    def __parse_preferences(self):
+    def __parse_preferences(self) -> None:
+        """
+        Loads the saved user preferences from disk (located in raidionics_preferences.json) and updates all internal
+        variables accordingly.
+        """
         with open(self._preferences_filename, 'r') as infile:
             preferences = json.load(infile)
 
@@ -223,6 +237,8 @@ def __parse_preferences(self):
                 self.use_registered_inputs = preferences['Processing']['use_registered_inputs']
             if 'export_results_as_rtstruct' in preferences['Processing'].keys():
                 self.export_results_as_rtstruct = preferences['Processing']['export_results_as_rtstruct']
+            if 'segmentation_tumor_model_type' in preferences['Processing'].keys():
+                self.segmentation_tumor_model_type = preferences['Processing']['segmentation_tumor_model_type']
             if 'perform_segmentation_refinement' in preferences['Processing'].keys():
                 self.perform_segmentation_refinement = preferences['Processing']['perform_segmentation_refinement']
             if 'SegmentationRefinement' in preferences['Processing'].keys():
@@ -244,7 +260,10 @@ def __parse_preferences(self):
             if 'dark_mode' in preferences['Appearance'].keys():
                 self._use_dark_mode = preferences['Appearance']['dark_mode']
 
-    def save_preferences(self):
+    def save_preferences(self) -> None:
+        """
+        Automatically saving on disk the user preferences inside the .raidionics folder, as raidionics_preferences.json
+        """
         preferences = {}
         preferences['System'] = {}
         preferences['System']['user_home_location'] = self._user_home_location
@@ -256,6 +275,7 @@ def save_preferences(self):
         preferences['Processing']['use_stripped_inputs'] = self._use_stripped_inputs
         preferences['Processing']['use_registered_inputs'] = self._use_registered_inputs
         preferences['Processing']['export_results_as_rtstruct'] = self._export_results_as_rtstruct
+        preferences['Processing']['segmentation_tumor_model_type'] = self.segmentation_tumor_model_type
         preferences['Processing']['perform_segmentation_refinement'] = self._perform_segmentation_refinement
         preferences['Processing']['SegmentationRefinement'] = {}
         preferences['Processing']['SegmentationRefinement']['type'] = self._segmentation_refinement_type

utils/logic/PipelineCreationHandler.py

@@ -111,6 +111,10 @@ def __create_segmentation_pipeline(model_name, patient_parameters):
         download_model(model_name='MRI_Sequence_Classifier')
 
     for steps in list(raw_pip.keys()):
+        # Excluding brain segmentation step if the inputs are already skull-stripped
+        if (UserPreferencesStructure.getInstance().use_stripped_inputs and
+                (raw_pip[steps]["task"] == "Segmentation" and raw_pip[steps]["model"] == "MRI_Brain")):
+            continue
         pip_num_int = pip_num_int + 1
         pip_num = str(pip_num_int)
         pip[pip_num] = raw_pip[steps]
@@ -197,6 +201,10 @@ def __create_postop_segmentation_pipeline(model_name, patient_parameters):
         download_model(model_name='MRI_Sequence_Classifier')
 
     for steps in list(raw_pip.keys()):
+        # Excluding brain segmentation step if the inputs are already skull-stripped
+        if (UserPreferencesStructure.getInstance().use_stripped_inputs and
+                (raw_pip[steps]["task"] == "Segmentation" and raw_pip[steps]["model"] == "MRI_Brain")):
+            continue
         pip_num_int = pip_num_int + 1
         pip_num = str(pip_num_int)
         pip[pip_num] = raw_pip[steps]
@@ -226,12 +234,16 @@ def __create_preop_reporting_pipeline(model_name, patient_parameters):
         download_model(model_name='MRI_Sequence_Classifier')
 
     for steps in list(raw_pip.keys()):
+        # Excluding brain segmentation step if the inputs are already skull-stripped
+        if (UserPreferencesStructure.getInstance().use_stripped_inputs and
+                (raw_pip[steps]["task"] == "Segmentation" and raw_pip[steps]["model"] == "MRI_Brain")):
+            continue
         pip_num_int = pip_num_int + 1
         pip_num = str(pip_num_int)
         pip[pip_num] = raw_pip[steps]
 
     # @TODO. Hard-coded, to remove/improve
-    if "Meningioma" in model_name:
+    if "Meningioma" in model_name and not UserPreferencesStructure.getInstance().use_stripped_inputs:
         pip_num_int = pip_num_int + 1
         pip_num = str(pip_num_int)
         pip[pip_num] = {}