Segmentation object now correctly handles indexing

Former-commit-id: 5cee77d
bhklab · Apr 16, 2020 · 107c807 · 107c807
1 parent 26328f7
commit 107c807
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/imgtools/segmentation/segmentation.py b/imgtools/segmentation/segmentation.py
@@ -63,3 +63,9 @@ def to_label_image(self):
         label_arr = arr.argmax(-1)
         label_img = array_to_image(label_arr, reference_image=self)
         return label_img
+
+    def __getitem__(self, idx):
+        res = super().__getitem__(idx)
+        if isinstance(res, sitk.Image):
+            res = Segmentation(res, self.roi_names)
+        return res