From ff2788945daaf62f6cec727850cd46a3ce8925f1 Mon Sep 17 00:00:00 2001
From: SkalskiP <piotr.skalski92@gmail.com>
Date: Sun, 9 Apr 2023 22:57:29 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=98=B7=20Initial=20implementation=20of=20?=
 =?UTF-8?q?MaskAnnotator?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 supervision/detection/annotate.py | 41 +++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py
index 070a5d66..f733eec1 100644
--- a/supervision/detection/annotate.py
+++ b/supervision/detection/annotate.py
@@ -117,3 +117,44 @@ class BoxAnnotator:
                 lineType=cv2.LINE_AA,
             )
         return scene
+
+
+class MaskAnnotator:
+    def __init__(
+        self,
+        color: Union[Color, ColorPalette] = ColorPalette.default(),
+    ):
+        self.color: Union[Color, ColorPalette] = color
+
+    def annotate(
+        self,
+        scene: np.ndarray,
+        detections: Detections,
+        opacity: float = 0.5
+    ) -> np.ndarray:
+
+        for i in range(len(detections.xyxy)):
+            if detections.mask is None:
+                continue
+
+            class_id = (
+                detections.class_id[i] if detections.class_id is not None else None
+            )
+            idx = class_id if class_id is not None else i
+            color = (
+                self.color.by_idx(idx)
+                if isinstance(self.color, ColorPalette)
+                else self.color
+            )
+
+            mask = detections.mask[i]
+            colored_mask = np.zeros_like(scene, dtype=np.uint8)
+            colored_mask[:] = color.as_rgb()
+
+            scene = np.where(
+                np.expand_dims(mask, axis=-1),
+                np.uint8(opacity * colored_mask + (1 - opacity) * scene),
+                scene
+            )
+
+        return scene
\ No newline at end of file