wasserth / TotalSegmentator

Tool for robust segmentation of >100 important anatomical structures in CT and MR images
Apache License 2.0
1.52k stars 253 forks source link

Key Error: vertebrae_body #265

Closed grinerD closed 9 months ago

grinerD commented 9 months ago

I received a license for the vertebrae_body segmentations, but when I try running it I get the following error:

Predicting s01:
perform_everything_on_gpu: True
100%|██████████| 18/18 [00:00<00:00, 185.41it/s]
Prediction done, transferring to CPU if needed
sending off prediction to background worker for resampling and export
done with s01
  Predicted in 7.23s
Resampling...
  back to original shape: (512, 512, 467)
Undoing canonical...
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[10], line 28
     24 print('\nSaving to: {}'.format(segmentations_directories[directory] + '/spine.nii.gz\n'))
     26 torch.cuda.empty_cache()
---> 28 totalsegmentator(dicom_directories[directory], segmentations_directories[directory] + '/vert_body.nii.gz',
     29              ml=True,
     30              roi_subset=rois_vert_body,
     31              #statistics=True,
     32              #device="cpu",
     33              verbose=True)
     35 print('\nSaving to: {}'.format(segmentations_directories[directory] + '/vert_body.nii.gz\n'))
     37 torch.cuda.empty_cache()

File /opt/conda/envs/py310/lib/python3.10/site-packages/totalsegmentator/python_api.py:272, in totalsegmentator(input, output, ml, nr_thr_resamp, nr_thr_saving, fast, nora_tag, preview, task, roi_subset, statistics, radiomics, crop_path, body_seg, force_split, output_type, quiet, verbose, test, skip_saving, device, license_number, statistics_exclude_masks_at_border, no_derived_masks, v1_order)
    270 roi_subset_crop = crop if crop is not None else roi_subset
    271 for roi in roi_subset_crop:
--> 272     crop_mask[organ_seg_data == class_map_inv[roi]] = 1
    273 crop_mask = nib.Nifti1Image(crop_mask, organ_seg.affine)
    274 crop_addon = [20,20,20]

KeyError: 'vertebrae_body'
grinerD commented 9 months ago

Sorry, I just realized I need to use task and not roi_subset.