From 8020da0cfdec918244c1b315b372d56c3165ef89 Mon Sep 17 00:00:00 2001 From: Florian Bruggisser Date: Fri, 1 Nov 2024 14:14:34 +0100 Subject: [PATCH] fixed masking issue if no mask has yet been detected --- spacestream/SpaceStreamGraph.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/spacestream/SpaceStreamGraph.py b/spacestream/SpaceStreamGraph.py index ef8446e..05035af 100644 --- a/spacestream/SpaceStreamGraph.py +++ b/spacestream/SpaceStreamGraph.py @@ -218,8 +218,9 @@ def _process(self): self.recorder.close() self.recorder = None + segmentations: Optional[List[vg.InstanceSegmentationResult]] = None if self.config.masking.value: - segmentations: List[vg.InstanceSegmentationResult] = self.segmentation_network.process(frame) + segmentations = self.segmentation_network.process(frame) for segment in segmentations: frame = self.mask_image(frame, segment.mask) @@ -277,8 +278,9 @@ def _process(self): depth_map = cv2.resize(depth_map, (w, h), interpolation=cv2.INTER_AREA) if self.config.masking.value: - for segment in segmentations: - depth_map = self.mask_image(depth_map, segment.mask) + if segmentations is not None: + for segment in segmentations: + depth_map = self.mask_image(depth_map, segment.mask) rgbd = np.hstack((depth_map, frame)) else: