From 107c8074d0a00150e15eafb05f206304a674373a Mon Sep 17 00:00:00 2001 From: mkazmier Date: Thu, 16 Apr 2020 11:15:31 -0400 Subject: [PATCH] Segmentation object now correctly handles indexing Former-commit-id: 5cee77d4a553bae602cc019876754130055aad81 --- imgtools/segmentation/segmentation.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/imgtools/segmentation/segmentation.py b/imgtools/segmentation/segmentation.py index 72a3fbe..fa532bd 100644 --- 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