From 2a1ba22f17ad31b84fa190e6cde09bc894d0c44f Mon Sep 17 00:00:00 2001 From: Nikita Date: Thu, 7 Nov 2024 22:23:40 +0000 Subject: [PATCH 1/2] feat: speedup background computation --- luxonis_ml/data/loaders/luxonis_loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From 1a9ed096a83b5197ddec12094ab49364a62cbf8b Mon Sep 17 00:00:00 2001 From: Nikita Date: Thu, 7 Nov 2024 22:24:37 +0000 Subject: [PATCH 2/2] feat: speedup SegmentationAnnotation.combine_to_numpy() function --- luxonis_ml/data/datasets/annotation.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 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):