diff --git a/models/totalsegmentator2/config/default.yml b/models/totalsegmentator2/config/default.yml new file mode 100644 index 00000000..3e12508a --- /dev/null +++ b/models/totalsegmentator2/config/default.yml @@ -0,0 +1,84 @@ +general: + data_base_dir: /app/data + version: 1.0.0 + description: TotalSegmentator default config (dicom to dicom) + +execute: +- DicomImporter +- NiftiConverter +- TotalSegmentatorMLRunner +- DsegConverter +- DataOrganizer + +segdb: + triplets: + C_BODY_STRUCTURE: + code: custom + segments: + thyroid_gland: + name: thyroid_gland + category: C_BODY_STRUCTURE + vertebrae_S1: + name: vertebrae_S1 + category: C_BODY_STRUCTURE + pulmonary_vein: + name: pulmonary_vein + category: C_BODY_STRUCTURE + brachiocephalic_trunk: + name: brachiocephalic_trunk + category: C_BODY_STRUCTURE + subclavian_artery_right: + name: subclavian_artery_right + category: C_BODY_STRUCTURE + subclavian_artery_left: + name: subclavian_artery_left + category: C_BODY_STRUCTURE + common_carotid_artery_right: + name: common_carotid_artery_right + category: C_BODY_STRUCTURE + common_carotid_artery_left: + name: common_carotid_artery_left + category: C_BODY_STRUCTURE + brachiocephalic_vein_left: + name: brachiocephalic_vein_left + category: C_BODY_STRUCTURE + brachiocephalic_vein_right: + name: brachiocephalic_vein_right + category: C_BODY_STRUCTURE + atrial_appendage_left: + name: atrial_appendage_left + category: C_BODY_STRUCTURE + spinal_cord: + name: spinal_cord + category: C_BODY_STRUCTURE + skull: + name: skull + category: C_BODY_STRUCTURE + sternum: + name: sternum + category: C_BODY_STRUCTURE + costal_cartilages: + name: costal_cartilages + category: C_BODY_STRUCTURE + +modules: + DicomImporter: + source_dir: input_data + import_dir: sorted_data + sort_data: true + meta: + mod: '%Modality' + + TotalSegmentatorMLRunner: + use_fast_mode: true + + DsegConverter: + model_name: TotalSegmentator2 + body_part_examined: WHOLEBODY + source_segs: nifti:mod=seg + skip_empty_slices: True + + DataOrganizer: + targets: + - dicomseg:mod=seg-->[i:sid]/TotalSegmentator2.seg.dcm + - nifti:mod=ct-->[i:sid]/image.nii.gz diff --git a/models/totalsegmentator2/dockerfiles/Dockerfile b/models/totalsegmentator2/dockerfiles/Dockerfile new file mode 100644 index 00000000..3b21db0c --- /dev/null +++ b/models/totalsegmentator2/dockerfiles/Dockerfile @@ -0,0 +1,33 @@ +FROM mhubai/base:latest + +# FIXME: set this environment variable as a shortcut to avoid nnunet crashing the build +# by pulling sklearn instead of scikit-learn +# N.B. this is a known issue: +# https://github.com/MIC-DKFZ/nnUNet/issues/1281 +# https://github.com/MIC-DKFZ/nnUNet/pull/1209 +ENV SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True + +# Install TotalSegmentator +RUN uv pip install -n totalsegmentator==2.2.1 + +# Download weights using totalsegmentator utility +# NOTE: only licence free models are included +RUN uv run totalseg_download_weights -t total \ +&& uv run totalseg_download_weights -t total_fast \ +&& uv run totalseg_download_weights -t total_mr \ +&& uv run totalseg_download_weights -t total_fast_mr \ +&& uv run totalseg_download_weights -t lung_vessels \ +&& uv run totalseg_download_weights -t cerebral_bleed \ +&& uv run totalseg_download_weights -t hip_implant \ +&& uv run totalseg_download_weights -t coronary_arteries \ +&& uv run totalseg_download_weights -t pleural_pericard_effusion \ +&& uv run totalseg_download_weights -t body \ +&& uv run totalseg_download_weights -t body_fast + +# Import the MHub model definiton +ARG MHUB_MODELS_REPO +RUN buildutils/import_mhub_model.sh totalsegmentator2 ${MHUB_MODELS_REPO} + +# Default run script +ENTRYPOINT ["mhub.run"] +CMD ["--workflow", "default"] \ No newline at end of file diff --git a/models/totalsegmentator2/utils/TotalSegmentatorMLRunner.py b/models/totalsegmentator2/utils/TotalSegmentatorMLRunner.py new file mode 100644 index 00000000..cfc704d6 --- /dev/null +++ b/models/totalsegmentator2/utils/TotalSegmentatorMLRunner.py @@ -0,0 +1,165 @@ +""" +------------------------------------------------- +MHub - Run Module for TotalSegmentator. +------------------------------------------------- + +------------------------------------------------- +Author: Leonard Nürnberg +Email: leonard.nuernberg@maastrichtuniversity.nl +------------------------------------------------- +""" + +from typing import Union, List +from mhubio.core import Module, Instance, InstanceData, IO + +#https://github.com/wasserth/TotalSegmentator/blob/master/totalsegmentator/map_to_binary.py +mapping = { + 'spleen': 'SPLEEN', + 'kidney_right': 'RIGHT_KIDNEY', + 'kidney_left': 'LEFT_KIDNEY', + 'gallbladder': 'GALLBLADDER', + 'liver': 'LIVER', + 'stomach': 'STOMACH', + 'pancreas': 'PANCREAS', + 'adrenal_gland_right': 'RIGHT_ADRENAL_GLAND', + 'adrenal_gland_left': 'LEFT_ADRENAL_GLAND', + 'lung_upper_lobe_left': 'LEFT_UPPER_LUNG_LOBE', + 'lung_lower_lobe_left': 'LEFT_LOWER_LUNG_LOBE', + 'lung_upper_lobe_right': 'RIGHT_UPPER_LUNG_LOBE', + 'lung_middle_lobe_right': 'RIGHT_MIDDLE_LUNG_LOBE', + 'lung_lower_lobe_right': 'RIGHT_LOWER_LUNG_LOBE', + 'vertebrae_L5': 'VERTEBRAE_L5', + 'vertebrae_L4': 'VERTEBRAE_L4', + 'vertebrae_L3': 'VERTEBRAE_L3', + 'vertebrae_L2': 'VERTEBRAE_L2', + 'vertebrae_L1': 'VERTEBRAE_L1', + 'vertebrae_T12': 'VERTEBRAE_T12', + 'vertebrae_T11': 'VERTEBRAE_T11', + 'vertebrae_T10': 'VERTEBRAE_T10', + 'vertebrae_T9': 'VERTEBRAE_T9', + 'vertebrae_T8': 'VERTEBRAE_T8', + 'vertebrae_T7': 'VERTEBRAE_T7', + 'vertebrae_T6': 'VERTEBRAE_T6', + 'vertebrae_T5': 'VERTEBRAE_T5', + 'vertebrae_T4': 'VERTEBRAE_T4', + 'vertebrae_T3': 'VERTEBRAE_T3', + 'vertebrae_T2': 'VERTEBRAE_T2', + 'vertebrae_T1': 'VERTEBRAE_T1', + 'vertebrae_C7': 'VERTEBRAE_C7', + 'vertebrae_C6': 'VERTEBRAE_C6', + 'vertebrae_C5': 'VERTEBRAE_C5', + 'vertebrae_C4': 'VERTEBRAE_C4', + 'vertebrae_C3': 'VERTEBRAE_C3', + 'vertebrae_C2': 'VERTEBRAE_C2', + 'vertebrae_C1': 'VERTEBRAE_C1', + 'esophagus': 'ESOPHAGUS', + 'trachea': 'TRACHEA', + 'heart_myocardium': 'MYOCARDIUM', + 'heart_atrium_left': 'LEFT_ATRIUM', + 'heart_ventricle_left': 'LEFT_VENTRICLE', + 'heart_atrium_right': 'RIGHT_ATRIUM', + 'heart_ventricle_right': 'RIGHT_VENTRICLE', + 'pulmonary_artery': 'PULMONARY_ARTERY', + 'brain': 'BRAIN', + 'iliac_artery_left': 'LEFT_ILIAC_ARTERY', + 'iliac_artery_right': 'RIGHT_ILIAC_ARTERY', + 'iliac_vena_left': 'LEFT_ILIAC_VEIN', + 'iliac_vena_right': 'RIGHT_ILIAC_VEIN', + 'small_bowel': 'SMALL_INTESTINE', + 'duodenum': 'DUODENUM', + 'colon': 'COLON', + 'rib_left_1': 'LEFT_RIB_1', + 'rib_left_2': 'LEFT_RIB_2', + 'rib_left_3': 'LEFT_RIB_3', + 'rib_left_4': 'LEFT_RIB_4', + 'rib_left_5': 'LEFT_RIB_5', + 'rib_left_6': 'LEFT_RIB_6', + 'rib_left_7': 'LEFT_RIB_7', + 'rib_left_8': 'LEFT_RIB_8', + 'rib_left_9': 'LEFT_RIB_9', + 'rib_left_10': 'LEFT_RIB_10', + 'rib_left_11': 'LEFT_RIB_11', + 'rib_left_12': 'LEFT_RIB_12', + 'rib_right_1': 'RIGHT_RIB_1', + 'rib_right_2': 'RIGHT_RIB_2', + 'rib_right_3': 'RIGHT_RIB_3', + 'rib_right_4': 'RIGHT_RIB_4', + 'rib_right_5': 'RIGHT_RIB_5', + 'rib_right_6': 'RIGHT_RIB_6', + 'rib_right_7': 'RIGHT_RIB_7', + 'rib_right_8': 'RIGHT_RIB_8', + 'rib_right_9': 'RIGHT_RIB_9', + 'rib_right_10': 'RIGHT_RIB_10', + 'rib_right_11': 'RIGHT_RIB_11', + 'rib_right_12': 'RIGHT_RIB_12', + 'humerus_left': 'LEFT_HUMERUS', + 'humerus_right': 'RIGHT_HUMERUS', + 'scapula_left': 'LEFT_SCAPULA', + 'scapula_right': 'RIGHT_SCAPULA', + 'clavicula_left': 'LEFT_CLAVICLE', + 'clavicula_right': 'RIGHT_CLAVICLE', + 'femur_left': 'LEFT_FEMUR', + 'femur_right': 'RIGHT_FEMUR', + 'hip_left': 'LEFT_HIP', + 'hip_right': 'RIGHT_HIP', + 'sacrum': 'SACRUM', + 'face': 'FACE', + 'gluteus_maximus_left': 'LEFT_GLUTEUS_MAXIMUS', + 'gluteus_maximus_right': 'RIGHT_GLUTEUS_MAXIMUS', + 'gluteus_medius_left': 'LEFT_GLUTEUS_MEDIUS', + 'gluteus_medius_right': 'RIGHT_GLUTEUS_MEDIUS', + 'gluteus_minimus_left': 'LEFT_GLUTEUS_MINIMUS', + 'gluteus_minimus_right': 'RIGHT_GLUTEUS_MINIMUS', + 'autochthon_left': 'LEFT_AUTOCHTHONOUS_BACK_MUSCLE', + 'autochthon_right': 'RIGHT_AUTOCHTHONOUS_BACK_MUSCLE', + 'iliopsoas_left': 'LEFT_ILIOPSOAS', + 'iliopsoas_right': 'RIGHT_ILIOPSOAS', + 'urinary_bladder': 'URINARY_BLADDER' +} + +def str2lst(string: Union[List[str], str]) -> list: + if isinstance(string, str): + return string.split(',') + else: + return string + +@IO.Config('use_fast_mode', bool, True, the="flag to set to run TotalSegmentator in fast mode") +@IO.Config('rois', list, [], factory=str2lst, the="comma separated list of rois to segment (if empty, all rois are segmented)") +class TotalSegmentatorMLRunner(Module): + + use_fast_mode: bool + rois: list + + @IO.Instance() + @IO.Input('in_data', 'nifti:mod=ct', the="input whole body ct scan") + @IO.Output('out_data', 'segmentations.nii.gz', 'nifti:mod=seg:model=TotalSegmentator:roi=SPLEEN,RIGHT_KIDNEY,LEFT_KIDNEY,GALLBLADDER,LIVER,STOMACH,PANCREAS,RIGHT_ADRENAL_GLAND,LEFT_ADRENAL_GLAND,LEFT_UPPER_LUNG_LOBE,LEFT_LOWER_LUNG_LOBE,RIGHT_UPPER_LUNG_LOBE,RIGHT_MIDDLE_LUNG_LOBE,RIGHT_LOWER_LUNG_LOBE,ESOPHAGUS,TRACHEA,thyroid_gland,SMALL_INTESTINE,DUODENUM,COLON,URINARY_BLADDER,PROSTATE,LEFT_KIDNEY+CYST,RIGHT_KIDNEY+CYST,SACRUM,vertebrae_S1,VERTEBRAE_L5,VERTEBRAE_L4,VERTEBRAE_L3,VERTEBRAE_L2,VERTEBRAE_L1,VERTEBRAE_T12,VERTEBRAE_T11,VERTEBRAE_T10,VERTEBRAE_T9,VERTEBRAE_T8,VERTEBRAE_T7,VERTEBRAE_T6,VERTEBRAE_T5,VERTEBRAE_T4,VERTEBRAE_T3,VERTEBRAE_T2,VERTEBRAE_T1,VERTEBRAE_C7,VERTEBRAE_C6,VERTEBRAE_C5,VERTEBRAE_C4,VERTEBRAE_C3,VERTEBRAE_C2,VERTEBRAE_C1,HEART,AORTA,pulmonary_vein,brachiocephalic_trunk,subclavian_artery_right,subclavian_artery_left,common_carotid_artery_right,common_carotid_artery_left,brachiocephalic_vein_left,brachiocephalic_vein_right,atrial_appendage_left,SUPERIOR_VENA_CAVA,INFERIOR_VENA_CAVA,PORTAL_AND_SPLENIC_VEIN,LEFT_ILIAC_ARTERY,RIGHT_ILIAC_ARTERY,LEFT_ILIAC_VEIN,RIGHT_ILIAC_VEIN,LEFT_HUMERUS,RIGHT_HUMERUS,LEFT_SCAPULA,RIGHT_SCAPULA,LEFT_CLAVICLE,RIGHT_CLAVICLE,LEFT_FEMUR,RIGHT_FEMUR,LEFT_HIP,RIGHT_HIP,spinal_cord,LEFT_GLUTEUS_MAXIMUS,RIGHT_GLUTEUS_MAXIMUS,LEFT_GLUTEUS_MEDIUS,RIGHT_GLUTEUS_MEDIUS,LEFT_GLUTEUS_MINIMUS,RIGHT_GLUTEUS_MINIMUS,LEFT_AUTOCHTHONOUS_BACK_MUSCLE,RIGHT_AUTOCHTHONOUS_BACK_MUSCLE,LEFT_ILIOPSOAS,RIGHT_ILIOPSOAS,BRAIN,skull,LEFT_RIB_1,LEFT_RIB_2,LEFT_RIB_3,LEFT_RIB_4,LEFT_RIB_5,LEFT_RIB_6,LEFT_RIB_7,LEFT_RIB_8,LEFT_RIB_9,LEFT_RIB_10,LEFT_RIB_11,LEFT_RIB_12,RIGHT_RIB_1,RIGHT_RIB_2,RIGHT_RIB_3,RIGHT_RIB_4,RIGHT_RIB_5,RIGHT_RIB_6,RIGHT_RIB_7,RIGHT_RIB_8,RIGHT_RIB_9,RIGHT_RIB_10,RIGHT_RIB_11,RIGHT_RIB_12,sternum,costal_cartilages', data='in_data', the="output segmentation mask containing all labels") + def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData) -> None: + + # build command + bash_command = ["TotalSegmentator"] + bash_command += ["-ta", "total"] + bash_command += ["-i", in_data.abspath] + + # multi-label output (one nifti file containing all labels instead of one nifti file per label) + self.v("Generating multi-label output ('--ml')") + bash_command += ["-o", out_data.abspath] + bash_command += ["--ml"] + + # fast mode + if self.use_fast_mode: + self.v("Running TotalSegmentator in fast mode ('--fast', 3mm)") + bash_command += ["--fast"] + else: + self.v("Running TotalSegmentator in default mode (1.5mm)") + + # roi subselection + if self.rois: + self.v("Subselecting ROIs: ", self.rois) + inv_mapping = {v: k for k, v in mapping.items()} + bash_command += ["--roi_subset", " ".join([inv_mapping[roi] for roi in self.rois])] + + # TODO: remove + self.v(">> run: ", " ".join(bash_command)) + + # run the model + self.subprocess(bash_command, text=True) \ No newline at end of file