Speed up semantic segmentation annotations loading

luxonis_ml/data/datasets/annotation.py

@@ -294,12 +294,21 @@ def combine_to_numpy(
         height: int,
         width: int,
     ) -> np.ndarray:
-        seg = np.zeros((len(class_mapping), height, width), dtype=np.bool_)
-        for ann in annotations:
-            class_ = class_mapping.get(ann.class_, 0)
-            seg[class_, ...] |= ann.to_numpy(class_mapping, width, height)
+        seg = np.zeros((len(class_mapping), height, width), dtype=np.uint8)
+
+        masks = np.stack(
+            [ann.to_numpy(class_mapping, width, height) for ann in annotations]
+        )
+        classes = np.array(
+            [class_mapping.get(ann.class_, 0) for ann in annotations]
+        )
+
+        for i, class_ in enumerate(classes):
+            seg[class_, ...] = np.maximum(
+                seg[class_, ...], masks[i].astype(np.uint8)
+            )
 
-        return seg.astype(np.uint8)
+        return seg
 
 
 class RLESegmentationAnnotation(SegmentationAnnotation):

luxonis_ml/data/loaders/luxonis_loader.py

@@ -318,7 +318,7 @@ def _load_image_with_annotations(
                 anns, self.class_mappings[task], width=width, height=height
             )
             if self.add_background and task == LabelType.SEGMENTATION:
-                unassigned_pixels = np.sum(array, axis=0) == 0
+                unassigned_pixels = ~np.any(array, axis=0)
                 background_idx = self.class_mappings[task]["background"]
                 array[background_idx, unassigned_pixels] = 1