fixed masking issue if no mask has yet been detected

cansik · Nov 1, 2024 · 8020da0 · 8020da0
1 parent 107b4de
commit 8020da0
Showing 1 changed file with 5 additions and 3 deletions.
diff --git 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: