Annotator txt_color updates (#13842)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Muhammad Rizwan Munawar <muhammadrizwanmunawar123@gmail.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>

docs/en/usage/simple-utilities.md

@@ -444,6 +444,86 @@ for obb in obb_boxes:
 image_with_obb = ann.result()
 ```
 
+#### Bounding Boxes Circle Annotation ([Circle Label](https://docs.ultralytics.com/reference/utils/plotting/#ultralytics.utils.plotting.Annotator.circle_label))
+
+```python
+import cv2
+
+from ultralytics import YOLO
+from ultralytics.utils.plotting import Annotator, colors
+
+model = YOLO("yolov8s.pt")
+cap = cv2.VideoCapture("path/to/video/file.mp4")
+
+w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
+writer = cv2.VideoWriter("Ultralytics circle annotation.avi", cv2.VideoWriter_fourcc(*"MJPG"), fps, (w, h))
+
+while True:
+    ret, im0 = cap.read()
+    if not ret:
+        break
+
+    annotator = Annotator(im0, line_width=2)
+
+    results = model.predict(im0)
+    boxes = results[0].boxes.xyxy.cpu()
+    clss = results[0].boxes.cls.cpu().tolist()
+
+    for box, cls in zip(boxes, clss):
+        x1, y1 = int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)
+        annotator.circle_label(box, label=model.names[int(cls)], color=colors(int(cls), True))
+
+    writer.write(im0)
+    cv2.imshow("Ultralytics circle annotation", im0)
+
+    if cv2.waitKey(1) & 0xFF == ord("q"):
+        break
+
+writer.release()
+cap.release()
+cv2.destroyAllWindows()
+```
+
+#### Bounding Boxes Text Annotation ([Text Label](https://docs.ultralytics.com/reference/utils/plotting/#ultralytics.utils.plotting.Annotator.text_label))
+
+```python
+import cv2
+
+from ultralytics import YOLO
+from ultralytics.utils.plotting import Annotator, colors
+
+model = YOLO("yolov8s.pt")
+cap = cv2.VideoCapture("path/to/video/file.mp4")
+
+w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
+writer = cv2.VideoWriter("Ultralytics text annotation.avi", cv2.VideoWriter_fourcc(*"MJPG"), fps, (w, h))
+
+while True:
+    ret, im0 = cap.read()
+    if not ret:
+        break
+
+    annotator = Annotator(im0, line_width=2)
+
+    results = model.predict(im0)
+    boxes = results[0].boxes.xyxy.cpu()
+    clss = results[0].boxes.cls.cpu().tolist()
+
+    for box, cls in zip(boxes, clss):
+        x1, y1 = int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)
+        annotator.text_label(box, label=model.names[int(cls)], color=colors(int(cls), True))
+
+    writer.write(im0)
+    cv2.imshow("Ultralytics text annotation", im0)
+
+    if cv2.waitKey(1) & 0xFF == ord("q"):
+        break
+
+writer.release()
+cap.release()
+cv2.destroyAllWindows()
+```
+
 See the [`Annotator` Reference Page](../reference/utils/plotting.md#ultralytics.utils.plotting.Annotator) for additional insight.
 
 ## Miscellaneous

ultralytics/utils/plotting.py

@@ -183,11 +183,108 @@ class Annotator:
             (104, 31, 17),
         }
 
-    def box_label(self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255), rotated=False):
+    def get_txt_color(self, color=(128, 128, 128), txt_color=(255, 255, 255)):
-        """Add one xyxy box to image with label."""
+        """Assign text color based on background color."""
-        txt_color = (
+        if color in self.dark_colors:
-            (104, 31, 17) if color in self.dark_colors else (255, 255, 255) if color in self.light_colors else txt_color
+            return 104, 31, 17
+        elif color in self.light_colors:
+            return 255, 255, 255
+        else:
+            return txt_color
+
+    def circle_label(self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255), margin=2):
+        """
+        Draws a label with a background rectangle centered within a given bounding box.
+
+        Args:
+            box (tuple): The bounding box coordinates (x1, y1, x2, y2).
+            label (str): The text label to be displayed.
+            color (tuple, optional): The background color of the rectangle (R, G, B).
+            txt_color (tuple, optional): The color of the text (R, G, B).
+            margin (int, optional): The margin between the text and the rectangle border.
+        """
+
+        # If label have more than 3 characters, skip other characters, due to circle size
+        if len(label) > 3:
+            print(
+                f"Length of label is {len(label)}, initial 3 label characters will be considered for circle annotation!"
             )
+            label = label[:3]
+
+        # Calculate the center of the box
+        x_center, y_center = int((box[0] + box[2]) / 2), int((box[1] + box[3]) / 2)
+        # Get the text size
+        text_size = cv2.getTextSize(str(label), cv2.FONT_HERSHEY_SIMPLEX, self.sf - 0.15, self.tf)[0]
+        # Calculate the required radius to fit the text with the margin
+        required_radius = int(((text_size[0] ** 2 + text_size[1] ** 2) ** 0.5) / 2) + margin
+        # Draw the circle with the required radius
+        cv2.circle(self.im, (x_center, y_center), required_radius, color, -1)
+        # Calculate the position for the text
+        text_x = x_center - text_size[0] // 2
+        text_y = y_center + text_size[1] // 2
+        # Draw the text
+        cv2.putText(
+            self.im,
+            str(label),
+            (text_x, text_y),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            self.sf - 0.15,
+            self.get_txt_color(color, txt_color),
+            self.tf,
+            lineType=cv2.LINE_AA,
+        )
+
+    def text_label(self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255), margin=5):
+        """
+        Draws a label with a background rectangle centered within a given bounding box.
+
+        Args:
+            box (tuple): The bounding box coordinates (x1, y1, x2, y2).
+            label (str): The text label to be displayed.
+            color (tuple, optional): The background color of the rectangle (R, G, B).
+            txt_color (tuple, optional): The color of the text (R, G, B).
+            margin (int, optional): The margin between the text and the rectangle border.
+        """
+
+        # Calculate the center of the bounding box
+        x_center, y_center = int((box[0] + box[2]) / 2), int((box[1] + box[3]) / 2)
+        # Get the size of the text
+        text_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, self.sf - 0.1, self.tf)[0]
+        # Calculate the top-left corner of the text (to center it)
+        text_x = x_center - text_size[0] // 2
+        text_y = y_center + text_size[1] // 2
+        # Calculate the coordinates of the background rectangle
+        rect_x1 = text_x - margin
+        rect_y1 = text_y - text_size[1] - margin
+        rect_x2 = text_x + text_size[0] + margin
+        rect_y2 = text_y + margin
+        # Draw the background rectangle
+        cv2.rectangle(self.im, (rect_x1, rect_y1), (rect_x2, rect_y2), color, -1)
+        # Draw the text on top of the rectangle
+        cv2.putText(
+            self.im,
+            label,
+            (text_x, text_y),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            self.sf - 0.1,
+            self.get_txt_color(color, txt_color),
+            self.tf,
+            lineType=cv2.LINE_AA,
+        )
+
+    def box_label(self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255), rotated=False):
+        """
+        Draws a bounding box to image with label.
+
+        Args:
+            box (tuple): The bounding box coordinates (x1, y1, x2, y2).
+            label (str): The text label to be displayed.
+            color (tuple, optional): The background color of the rectangle (R, G, B).
+            txt_color (tuple, optional): The color of the text (R, G, B).
+            rotated (bool, optional): Variable used to check if task is OBB
+        """
+
+        txt_color = self.get_txt_color(color, txt_color)
         if isinstance(box, torch.Tensor):
             box = box.tolist()
         if self.pil or not is_ascii(label):
@@ -242,6 +339,7 @@ class Annotator:
             alpha (float): Mask transparency: 0.0 fully transparent, 1.0 opaque
             retina_masks (bool): Whether to use high resolution masks or not. Defaults to False.
         """
+
         if self.pil:
             # Convert to numpy first
             self.im = np.asarray(self.im).copy()
@@ -281,6 +379,7 @@ class Annotator:
         Note:
             `kpt_line=True` currently only supports human pose plotting.
         """
+
         if self.pil:
             # Convert to numpy first
             self.im = np.asarray(self.im).copy()
@@ -376,6 +475,7 @@ class Annotator:
         Returns:
             angle (degree): Degree value of angle between three points
         """
+
         x_min, y_min, x_max, y_max = bbox
         width = x_max - x_min
         height = y_max - y_min
@@ -390,6 +490,7 @@ class Annotator:
             color (tuple): Region Color value
             thickness (int): Region area thickness value
         """
+
         cv2.polylines(self.im, [np.array(reg_pts, dtype=np.int32)], isClosed=True, color=color, thickness=thickness)
 
     def draw_centroid_and_tracks(self, track, color=(255, 0, 255), track_thickness=2):
@@ -401,6 +502,7 @@ class Annotator:
             color (tuple): tracks line color
             track_thickness (int): track line thickness value
         """
+
         points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
         cv2.polylines(self.im, [points], isClosed=False, color=color, thickness=track_thickness)
         cv2.circle(self.im, (int(track[-1][0]), int(track[-1][1])), track_thickness * 2, color, -1)
@@ -513,6 +615,7 @@ class Annotator:
         Returns:
             angle (degree): Degree value of angle between three points
         """
+
         a, b, c = np.array(a), np.array(b), np.array(c)
         radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
         angle = np.abs(radians * 180.0 / np.pi)
@@ -530,6 +633,7 @@ class Annotator:
             shape (tuple): imgsz for model inference
             radius (int): Keypoint radius value
         """
+
         if indices is None:
             indices = [2, 5, 7]
         for i, k in enumerate(keypoints):
@@ -626,6 +730,7 @@ class Annotator:
             det_label (str): Detection label text
             track_label (str): Tracking label text
         """
+
         cv2.polylines(self.im, [np.int32([mask])], isClosed=True, color=mask_color, thickness=2)
 
         label = f"Track ID: {track_label}" if track_label else det_label
@@ -695,6 +800,7 @@ class Annotator:
             color (tuple): object centroid and line color value
             pin_color (tuple): visioneye point color value
         """
+
         center_bbox = int((box[0] + box[2]) / 2), int((box[1] + box[3]) / 2)
         cv2.circle(self.im, center_point, self.tf * 2, pin_color, -1)
         cv2.circle(self.im, center_bbox, self.tf * 2, color, -1)