diff --git a/CHANGELOG.md b/CHANGELOG.md index e0d97b556..9e085e7bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## Master * add brain structures * add liver vessels +* greatly improved phase classification model ## Release 2.3.0 diff --git a/README.md b/README.md index 5bcd25fe4..25a001767 100644 --- a/README.md +++ b/README.md @@ -133,15 +133,6 @@ If you want to reduce memory consumption you can use the following options: * `--nr_thr_saving 1`: Saving big images with several threads will take a lot of memory -### Train/validation/test split -The exact split of the dataset can be found in the file `meta.csv` inside of the [dataset](https://doi.org/10.5281/zenodo.6802613). This was used for the validation in our paper. -The exact numbers of the results for the high-resolution model (1.5mm) can be found [here](resources/results_all_classes_v1.json). The paper shows these numbers in the supplementary materials Figure 11. - - -### Retrain model and run evaluation -See [here](resources/train_nnunet.md) for more info on how to train a nnU-Net yourself on the TotalSegmentator dataset, how to split the data into train/validation/test set as in our paper, and how to run the same evaluation as in our paper. - - ### Python API You can run totalsegmentator via Python: ```python @@ -159,13 +150,12 @@ if __name__ == "__main__": ``` You can see all available arguments [here](https://github.com/wasserth/TotalSegmentator/blob/master/totalsegmentator/python_api.py). Running from within the main environment should avoid some multiprocessing issues. -The segmentation image contains the names of the classes in the extended header. If you want to load this additional header information you can use the following code: +The segmentation image contains the names of the classes in the extended header. If you want to load this additional header information you can use the following code (requires `pip install xmltodict`): ```python from totalsegmentator.nifti_ext_header import load_multilabel_nifti segmentation_nifti_img, label_map_dict = load_multilabel_nifti(image_path) ``` -The above code requires `pip install xmltodict`. ### Install latest master branch (contains latest bug fixes) @@ -175,6 +165,11 @@ pip install git+https://github.com/wasserth/TotalSegmentator.git ### Other commands +If you want to know which contrast phase a CT image is you can use the following command (requires `pip install xgboost`). More details can be found [here](resources/contrast_phase_prediction.md): +``` +totalseg_get_phase -i ct.nii.gz -o contrast_phase.json +``` + If you want to combine some subclasses (e.g. lung lobes) into one binary mask (e.g. entire lung) you can use the following command: ``` totalseg_combine_masks -i totalsegmentator_output_dir -o combined_mask.nii.gz -m lungcomm @@ -191,6 +186,15 @@ totalseg_set_license -l aca_12345678910 ``` +### Train/validation/test split +The exact split of the dataset can be found in the file `meta.csv` inside of the [dataset](https://doi.org/10.5281/zenodo.6802613). This was used for the validation in our paper. +The exact numbers of the results for the high-resolution model (1.5mm) can be found [here](resources/results_all_classes_v1.json). The paper shows these numbers in the supplementary materials Figure 11. + + +### Retrain model and run evaluation +See [here](resources/train_nnunet.md) for more info on how to train a nnU-Net yourself on the TotalSegmentator dataset, how to split the data into train/validation/test set as in our paper, and how to run the same evaluation as in our paper. + + ### Typical problems **ITK loading Error** diff --git a/resources/contrast_phase_classifiers.pkl b/resources/contrast_phase_classifiers.pkl deleted file mode 100644 index faeb6743a..000000000 Binary files a/resources/contrast_phase_classifiers.pkl and /dev/null differ diff --git a/resources/contrast_phase_classifiers_2024_07_19.pkl b/resources/contrast_phase_classifiers_2024_07_19.pkl new file mode 100644 index 000000000..76768c490 Binary files /dev/null and b/resources/contrast_phase_classifiers_2024_07_19.pkl differ diff --git a/resources/contrast_phase_prediction.md b/resources/contrast_phase_prediction.md new file mode 100644 index 000000000..26df5221e --- /dev/null +++ b/resources/contrast_phase_prediction.md @@ -0,0 +1,30 @@ +# Details on how the prediction of the contrast phase is done + +TotalSegmentator is used to predict the following structures: +```python +["liver", "pancreas", "urinary_bladder", "gallbladder", +"heart", "aorta", "inferior_vena_cava", "portal_vein_and_splenic_vein", +"iliac_vena_left", "iliac_vena_right", "iliac_artery_left", "iliac_artery_right", +"pulmonary_vein", "brain", "colon", "small_bowel", +"internal_carotid_artery_right", "internal_carotid_artery_left", +"internal_jugular_vein_right", "internal_jugular_vein_left"] +``` +Then the median intensity (HU value) of each structure is used as feature for a xgboost classifier +to predict the post injection time (pi_time). The pi_time can be mapped to the contrast phase +then. It classifies into `native`, `arterial_early`, `arterial_late`, and `portal_venous` phase. +The classifier was trained on the TotalSegmentator dataset and therefore works with all sorts +of different CT images. + +Results on 5-fold cross validation: + +- Mean absolute error (MAE): 5.55s +- F1 scores for each class: + - native: 0.980 + - arterial_early+late: 0.915 + - portal: 0.940 + +The results contain a probablity for each class which is high if the predicted pi_time is close to the ideal +pi_time for the given phase. Moreover, the classifier is an ensemble of 5 models. The output contains the +standard deviation of the predictions which can be used as a measure of confidence. If it is low the 5 models +give similar predictions which is a good sign. + diff --git a/setup.py b/setup.py index c4be90a93..5c93db495 100644 --- a/setup.py +++ b/setup.py @@ -11,7 +11,10 @@ python_requires='>=3.9', license='Apache 2.0', packages=find_packages(), - package_data={"totalsegmentator": ["resources/totalsegmentator_snomed_mapping.csv"]}, + package_data={"totalsegmentator": + ["resources/totalsegmentator_snomed_mapping.csv", + "resources/contrast_phase_classifiers_2024_07_19.pkl"] + }, install_requires=[ 'torch>=2.0.0', 'numpy<2', diff --git a/totalsegmentator/bin/totalseg_get_phase.py b/totalsegmentator/bin/totalseg_get_phase.py index 4955d4ca5..451052dcb 100644 --- a/totalsegmentator/bin/totalseg_get_phase.py +++ b/totalsegmentator/bin/totalseg_get_phase.py @@ -68,7 +68,7 @@ def get_ct_contrast_phase(ct_img: nib.Nifti1Image, model_file: Path = None): # print(f"ts took: {time.time()-st:.2f}s") if stats["brain"]["volume"] > 100: - # print(f"Brain in image, therefore also running headneck model.") + # print("Brain in image, therefore also running headneck model.") st = time.time() seg_img_hn, stats_hn = totalsegmentator(ct_img, None, ml=True, fast=False, statistics=True, task="headneck_bones_vessels", @@ -85,11 +85,9 @@ def get_ct_contrast_phase(ct_img: nib.Nifti1Image, model_file: Path = None): features.append(stats_hn[organ]["intensity"]) if model_file is None: - # weights from longitudinalliver dataset - classifier_path = Path(__file__).parents[2] / "resources" / "contrast_phase_classifiers.pkl" - else: - # weights from megaseg dataset - # classifier_path = "/mnt/nor/wasserthalj_data/classifiers_megaseg.pkl" + classifier_path = Path(__file__).parents[2] / "resources" / "contrast_phase_classifiers_2024_07_19.pkl" + else: + # manually set model file classifier_path = model_file clfs = pickle.load(open(classifier_path, "rb")) @@ -136,13 +134,17 @@ def main(): parser.add_argument("-m", metavar="filepath", dest="model_file", help="path to classifier model", type=lambda p: Path(p).absolute(), required=False, default=None) + + parser.add_argument("-q", dest="quiet", action="store_true", + help="Print no output to stdout", default=False) args = parser.parse_args() res = get_ct_contrast_phase(nib.load(args.input_file), args.model_file) - print("Result:") - pprint(res) + if not args.quiet: + print("Result:") + pprint(res) with open(args.output_file, "w") as f: f.write(json.dumps(res, indent=4))