fix: move overlap check to combine_to_numpy()

luxonis_ml/data/datasets/annotation.py

@@ -313,10 +313,11 @@ def combine_to_numpy(
             [class_mapping.get(ann.class_, 0) for ann in annotations]
         )
 
+        assigned_pixels = np.zeros((height, width), dtype=bool)
         for i, class_ in enumerate(classes):
-            seg[class_, ...] = np.maximum(
-                seg[class_, ...], masks[i].astype(np.uint8)
-            )
+            mask = masks[i] & (assigned_pixels == 0)
+            seg[class_, ...] |= mask
+            assigned_pixels |= mask
 
         return seg
 

luxonis_ml/data/loaders/luxonis_loader.py

@@ -316,24 +316,10 @@ def _load_image_with_annotations(
             array = anns[0].combine_to_numpy(
                 anns, self.class_mappings[task], width=width, height=height
             )
-
-            if task == LabelType.SEGMENTATION:
-                mask_sum = np.sum(array, axis=0)
-
-                if self.add_background:
-                    unassigned_pixels = mask_sum == 0
-                    background_idx = self.class_mappings[task]["background"]
-                    array[background_idx, unassigned_pixels] = 1
-
-                # Check if there are overlaps in the masks
-                if np.any(mask_sum > 1):
-                    max_indices = np.argmax(array, axis=0)
-                    array.fill(0)
-                    array[
-                        max_indices,
-                        np.arange(array.shape[1])[:, None],
-                        np.arange(array.shape[2]),
-                    ] = 1
+            if self.add_background and task == LabelType.SEGMENTATION:
+                unassigned_pixels = ~np.any(array, axis=0)
+                background_idx = self.class_mappings[task]["background"]
+                array[background_idx, unassigned_pixels] = 1
 
             labels[task] = (array, anns[0]._label_type)