bugfixes

wasserth · Aug 24, 2023 · 5d088d9 · 5d088d9
1 parent c6343a6
commit 5d088d9
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -37,18 +37,20 @@ TotalSegmentator -i ct.nii.gz -o segmentations
 ```
 > Note: A Nifti file or a folder of DICOM images is allowed as input
 
-> Note: If a CUDA compatible GPU is available TotalSegmentator will automatically use it. Otherwise it will use the CPU, which is a lot slower and should only be used with the `--fast` option.  
+> Note: If you run on CPU you can use the option `--fast` or `--roi_subset liver` to greatly decrease 
+runtime.
 
 > Note: You can also try it online: [www.totalsegmentator.com](https://totalsegmentator.com/) (supports dicom files)
 
 > Note: This is not a medical device and not intended for clinical usage.
 
 
 ### Advanced settings
+* `--device`: Choose `cpu` or `gpu`
 * `--fast`: For faster runtime and less memory requirements use this option. It will run a lower resolution model (3mm instead of 1.5mm). 
+* `--roi_subset`: Takes a space separated list of class names (e.g. `spleen colon brain`) and only predicts those classes. Saves a lot of runtime and memory.
 * `--preview`: This will generate a 3D rendering of all classes, giving you a quick overview if the segmentation worked and where it failed (see `preview.png` in output directory).
 * `--ml`: This will save one nifti file containing all labels instead of one file for each class. Saves runtime during saving of nifti files. (see [here](https://github.com/wasserth/TotalSegmentator#class-details) for index to class name mapping).
-* `--roi_subset`: Takes a space separated list of class names (e.g. `spleen colon brain`) and only saves those classes. Saves runtime during saving of nifti files.
 * `--statistics`: This will generate a file `statistics.json` with volume (in mm³) and mean intensity of each class.
 * `--radiomics`: This will generate a file `statistics_radiomics.json` with radiomics features of each class. You have to install pyradiomics to use this (`pip install pyradiomics`).
 

diff --git a/totalsegmentator/python_api.py b/totalsegmentator/python_api.py
@@ -251,7 +251,7 @@ def totalsegmentator(input, output, ml=False, nr_thr_resamp=1, nr_thr_saving=6,
             if not quiet: print("Generating rough body segmentation...")
             # todo: switch to ep4000 model!
             organ_seg = nnUNet_predict_image(input, None, 298, model="3d_fullres", folds=[0],
-                                trainer="nnUNetTrainerNoMirroring", tta=False, multilabel_image=True, resample=6.0,
+                                trainer="nnUNetTrainer_4000epochs_NoMirroring", tta=False, multilabel_image=True, resample=6.0,
                                 crop=None, crop_path=None, task_name="total", nora_tag="None", preview=False, 
                                 save_binary=False, nr_threads_resampling=nr_thr_resamp, nr_threads_saving=1, 
                                 crop_addon=crop_addon, output_type=output_type, statistics=False,
@@ -271,8 +271,8 @@ def totalsegmentator(input, output, ml=False, nr_thr_resamp=1, nr_thr_saving=6,
         download_pretrained_weights(300)
         st = time.time()
         if not quiet: print("Generating rough body segmentation...")
-        body_seg = nnUNet_predict_image(input, None, 269, model="3d_fullres", folds=[0],
-                            trainer="nnUNetTrainerV2", tta=False, multilabel_image=True, resample=6.0,
+        body_seg = nnUNet_predict_image(input, None, 300, model="3d_fullres", folds=[0],
+                            trainer="nnUNetTrainer", tta=False, multilabel_image=True, resample=6.0,
                             crop=None, crop_path=None, task_name="body", nora_tag="None", preview=False, 
                             save_binary=True, nr_threads_resampling=nr_thr_resamp, nr_threads_saving=1, 
                             crop_addon=crop_addon, output_type=output_type, statistics=False,