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
This commit is contained in:
xenteros 2023-11-16 22:27:32 +01:00
parent fe4657f0d0
commit c2101db5d2
1 changed files with 7 additions and 6 deletions

View File

@ -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):