bugfix in load ext header with only one label

CHANGELOG.md

@@ -1,6 +1,8 @@
 ## Master
 * also return statistics from python api
 * add `totalseg_get_phase`
+* major bugfix: rib labels were in wrong order
+
 
 ## Release 2.1.0
 * Bugfix: add flush to DummyFile

tests/test_locally.py

@@ -193,6 +193,7 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
         print("Run totalsegmentator...")
         reset_monitors()
         for img_fn in tqdm(img_dir.glob("*.nii.gz")):
+        # for img_fn in tqdm(list(img_dir.glob("*.nii.gz"))[:1]):
             fast = resolution == "3mm"
             st = time.time()
             totalsegmentator(img_fn, pred_dir / img_fn.name, fast=fast, ml=True, device=device)

totalsegmentator/nifti_ext_header.py

@@ -67,6 +67,11 @@ def load_multilabel_nifti(img_path):
     ext_header = img.header.extensions[0].get_content()
     ext_header = xmltodict.parse(ext_header)
     ext_header = ext_header["CaretExtension"]["VolumeInformation"]["LabelTable"]["Label"]
+
+    # If only one label, ext_header is a dict instead of a list (because of xmltodict.parse()) -> convert to list
+    if isinstance(ext_header, dict):
+        ext_header = [ext_header]
+
     label_map = {int(e["@Key"]): e["#text"] for e in ext_header}
     return img, label_map