added list capabilities for the roi names dictionary

bhklab · Jun 16, 2022 · 2293242 · 2293242
1 parent 23b76e4
commit 2293242
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/imgtools/modules/structureset.py b/imgtools/modules/structureset.py
@@ -134,9 +134,16 @@ def to_segmentation(self, reference_image: sitk.Image,
             labels = self._assign_labels(roi_names, force_missing) #only the ones that match the regex
         elif isinstance(roi_names, dict):
             for name, pattern in roi_names.items():
-                matching_names = list(self._assign_labels([pattern], force_missing).keys())
-                if matching_names:
-                    labels[name] = matching_names #{"GTV": ["GTV1", "GTV2"]}
+                if isinstance(pattern, str):
+                    matching_names = list(self._assign_labels([pattern], force_missing).keys())
+                    if matching_names:
+                        labels[name] = matching_names #{"GTV": ["GTV1", "GTV2"]}
+                elif isinstance(pattern, list):
+                    labels[name] = []
+                    for pat in pattern:
+                        matching_names = list(self._assign_labels([pat], force_missing).keys())
+                        if matching_names:
+                            labels[name].extend(matching_names) #{"GTV": ["GTV1", "GTV2"]}
         if isinstance(roi_names, str):
             roi_names = [roi_names]
         if isinstance(roi_names, list):