From bfd3bbc1dce081f18d0b0119d199721592dea8d6 Mon Sep 17 00:00:00 2001 From: Nikita Sokovnin <49622375+sokovninn@users.noreply.github.com> Date: Fri, 8 Nov 2024 12:49:54 +0100 Subject: [PATCH] Speed up semantic segmentation annotations loading (#205) Co-authored-by: Martin Kozlovsky --- luxonis_ml/data/datasets/annotation.py | 19 ++++++++++++++----- luxonis_ml/data/loaders/luxonis_loader.py | 2 +- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/luxonis_ml/data/datasets/annotation.py b/luxonis_ml/data/datasets/annotation.py index 9a11bc76..fe5e70ec 100644 --- a/luxonis_ml/data/datasets/annotation.py +++ b/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): diff --git a/luxonis_ml/data/loaders/luxonis_loader.py b/luxonis_ml/data/loaders/luxonis_loader.py index 4c66c39c..adab865c 100644 --- a/luxonis_ml/data/loaders/luxonis_loader.py +++ b/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