From c2101db5d24731df48b16aefb5acfd3a4132c0bd Mon Sep 17 00:00:00 2001 From: xenteros Date: Thu, 16 Nov 2023 22:27:32 +0100 Subject: [PATCH] Optimize sv.MaskAnnotator by reducing number of calls to cv2.addWeighted * Combined all colored masks before applying them to the scene * Made sure that returned ndarray is np.uint8 for video writer simplicity --- supervision/annotators/core.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 82d299da..2bdf25c2 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -156,6 +156,8 @@ class MaskAnnotator(BaseAnnotator): if detections.mask is None: return scene + colored_mask = np.array(scene, copy=True, dtype=np.uint8) + for detection_idx in np.flip(np.argsort(detections.area)): color = resolve_color( color=self.color, @@ -166,13 +168,12 @@ class MaskAnnotator(BaseAnnotator): else custom_color_lookup, ) mask = detections.mask[detection_idx] - colored_mask = np.zeros_like(scene, dtype=np.uint8) - colored_mask[:] = color.as_bgr() - scene[mask] = cv2.addWeighted( - colored_mask, self.opacity, scene, 1 - self.opacity, 0 - )[mask] + colored_mask[mask] = color.as_bgr() - return scene + scene = cv2.addWeighted( + colored_mask, self.opacity, scene, 1 - self.opacity, 0 + ) + return scene.astype(np.uint8) class PolygonAnnotator(BaseAnnotator):