diff --git a/docs/en/solutions/index.md b/docs/en/solutions/index.md index 6231461f05..63d3906b68 100644 --- a/docs/en/solutions/index.md +++ b/docs/en/solutions/index.md @@ -66,21 +66,20 @@ Here's our curated list of Ultralytics solutions that can be used to create awes All Ultralytics Solutions use the separate class [`SolutionAnnotator`](https://docs.ultralytics.com/reference/solutions/solutions/#ultralytics.solutions.solutions.SolutionAnnotator), that extends the main [`Annotator`](https://docs.ultralytics.com/reference/utils/plotting/#ultralytics.utils.plotting.Annotator) class, and have the following methods: -| Method | Return Type | Description | -| ---------------------------------- | ----------- | ---------------------------------------------------------------------- | -| `draw_region()` | `None` | Draws a region using specified points, colors, and thickness. | -| `queue_counts_display()` | `None` | Displays queue counts in the specified region. | -| `display_analytics()` | `None` | Displays overall statistics for parking lot management. | -| `estimate_pose_angle()` | `float` | Calculates the angle between three points in an object pose. | -| `draw_specific_points()` | `None` | Draws specific keypoints on the image. | -| `plot_workout_information()` | `None` | Draws a labeled text box on the image. | -| `plot_angle_and_count_and_stage()` | `None` | Visualizes angle, step count, and stage for workout monitoring. | -| `plot_distance_and_line()` | `None` | Displays the distance between centroids and connects them with a line. | -| `display_objects_labels()` | `None` | Annotates bounding boxes with object class labels. | -| `seg_bbox()` | `None` | Draws contours for segmented objects and optionally labels them. | -| `visioneye()` | `None` | Maps and connects object centroids to a visual "eye" point. | -| `circle_label()` | `None` | Draws a circular label in the place of a bounding box. | -| `text_label()` | `None` | Draws a rectangular label in the place of a bounding box. | +| Method | Return Type | Description | +| ---------------------------------- | ----------- | -------------------------------------------------------------------------------- | +| `draw_region()` | `None` | Draws a region using specified points, colors, and thickness. | +| `queue_counts_display()` | `None` | Displays queue counts in the specified region. | +| `display_analytics()` | `None` | Displays overall statistics for parking lot management. | +| `estimate_pose_angle()` | `float` | Calculates the angle between three points in an object pose. | +| `draw_specific_points()` | `None` | Draws specific keypoints on the image. | +| `plot_workout_information()` | `None` | Draws a labeled text box on the image. | +| `plot_angle_and_count_and_stage()` | `None` | Visualizes angle, step count, and stage for workout monitoring. | +| `plot_distance_and_line()` | `None` | Displays the distance between centroids and connects them with a line. | +| `display_objects_labels()` | `None` | Annotates bounding boxes with object class labels. | +| `sweep_annotator()` | `None` | Visualize a vertical sweep line and optional label. | +| `visioneye()` | `None` | Maps and connects object centroids to a visual "eye" point. | +| `adaptive_label()` | `None` | Draw a circular or rectangle background shape label in center of a bounding box. | ### Working with SolutionResults diff --git a/docs/en/usage/simple-utilities.md b/docs/en/usage/simple-utilities.md index b8b6b46e21..cb7c44068e 100644 --- a/docs/en/usage/simple-utilities.md +++ b/docs/en/usage/simple-utilities.md @@ -379,189 +379,199 @@ See the docstring for each function or visit the `ultralytics.utils.ops` [refere ## Plotting -### Drawing Annotations +### Annotation utilities Ultralytics includes an `Annotator` class for annotating various data types. It's best used with [object detection bounding boxes](../modes/predict.md#boxes), [pose keypoints](../modes/predict.md#keypoints), and [oriented bounding boxes](../modes/predict.md#obb). -#### Ultralytics Sweep Annotation +#### Box Annotation !!! example "Python Examples using Ultralytics YOLO 🚀" - === "Python" + === "Horizontal Bounding Boxes" ```python - import cv2 + import cv2 as cv import numpy as np - from ultralytics import YOLO - from ultralytics.solutions.solutions import SolutionAnnotator - from ultralytics.utils.plotting import colors + from ultralytics.utils.plotting import Annotator, colors - # User defined video path and model file - cap = cv2.VideoCapture("path/to/video.mp4") - model = YOLO(model="yolo11s-seg.pt") # Model file i.e. yolo11s.pt or yolo11m-seg.pt + names = { + 0: "person", + 5: "bus", + 11: "stop sign", + } - if not cap.isOpened(): - print("Error: Could not open video.") - exit() + image = cv.imread("ultralytics/assets/bus.jpg") + ann = Annotator( + image, + line_width=None, # default auto-size + font_size=None, # default auto-size + font="Arial.ttf", # must be ImageFont compatible + pil=False, # use PIL, otherwise uses OpenCV + ) - # Initialize the video writer object. - w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) - video_writer = cv2.VideoWriter("ultralytics.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) + xyxy_boxes = np.array( + [ + [5, 22.878, 231.27, 804.98, 756.83], # class-idx x1 y1 x2 y2 + [0, 48.552, 398.56, 245.35, 902.71], + [0, 669.47, 392.19, 809.72, 877.04], + [0, 221.52, 405.8, 344.98, 857.54], + [0, 0, 550.53, 63.01, 873.44], + [11, 0.0584, 254.46, 32.561, 324.87], + ] + ) - masks = None # Initialize variable to store masks data - f = 0 # Initialize frame count variable for enabling mouse event. - line_x = w # Store width of line. - dragging = False # Initialize bool variable for line dragging. - classes = model.names # Store model classes names for plotting. - window_name = "Ultralytics Sweep Annotator" + for nb, box in enumerate(xyxy_boxes): + c_idx, *box = box + label = f"{str(nb).zfill(2)}:{names.get(int(c_idx))}" + ann.box_label(box, label, color=colors(c_idx, bgr=True)) - - def drag_line(event, x, _, flags, param): - """Mouse callback function to enable dragging a vertical sweep line across the video frame.""" - global line_x, dragging - if event == cv2.EVENT_LBUTTONDOWN or (flags & cv2.EVENT_FLAG_LBUTTON): - line_x = max(0, min(x, w)) - dragging = True - - - while cap.isOpened(): # Loop over the video capture object. - ret, im0 = cap.read() - if not ret: - break - f = f + 1 # Increment frame count. - count = 0 # Re-initialize count variable on every frame for precise counts. - results = model.track(im0, persist=True)[0] - - if f == 1: - cv2.namedWindow(window_name) - cv2.setMouseCallback(window_name, drag_line) - - annotator = SolutionAnnotator(im0) - - if results.boxes.is_track: - if results.masks is not None: - masks = [np.array(m, dtype=np.int32) for m in results.masks.xy] - - boxes = results.boxes.xyxy.tolist() - track_ids = results.boxes.id.int().cpu().tolist() - clss = results.boxes.cls.cpu().tolist() - - for mask, box, cls, t_id in zip(masks or [None] * len(boxes), boxes, clss, track_ids): - color = colors(t_id, True) # Assign different color to each tracked object. - label = f"{classes[cls]}:{t_id}" - if mask is not None and mask.size > 0: - if box[0] > line_x: - count += 1 - cv2.polylines(im0, [mask], True, color, 2) - x, y = mask.min(axis=0) - (w_m, _), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) - cv2.rectangle(im0, (x, y - 20), (x + w_m, y), color, -1) - cv2.putText(im0, label, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) - else: - if box[0] > line_x: - count += 1 - annotator.box_label(box=box, color=color, label=label) - - # Generate draggable sweep line - annotator.sweep_annotator(line_x=line_x, line_y=h, label=f"COUNT:{count}") - - cv2.imshow(window_name, im0) - video_writer.write(im0) - if cv2.waitKey(1) & 0xFF == ord("q"): - break - - # Release the resources - cap.release() - video_writer.release() - cv2.destroyAllWindows() + image_with_bboxes = ann.result() ``` + === "Oriented Bounding Boxes (OBB)" + + ```python + import cv2 as cv + import numpy as np + + from ultralytics.utils.plotting import Annotator, colors + + obb_names = {10: "small vehicle"} + obb_image = cv.imread("datasets/dota8/images/train/P1142__1024__0___824.jpg") + obb_boxes = np.array( + [ + [0, 635, 560, 919, 719, 1087, 420, 803, 261], # class-idx x1 y1 x2 y2 x3 y2 x4 y4 + [0, 331, 19, 493, 260, 776, 70, 613, -171], + [9, 869, 161, 886, 147, 851, 101, 833, 115], + ] + ) + ann = Annotator( + obb_image, + line_width=None, # default auto-size + font_size=None, # default auto-size + font="Arial.ttf", # must be ImageFont compatible + pil=False, # use PIL, otherwise uses OpenCV + ) + for obb in obb_boxes: + c_idx, *obb = obb + obb = np.array(obb).reshape(-1, 4, 2).squeeze() + label = f"{obb_names.get(int(c_idx))}" + ann.box_label( + obb, + label, + color=colors(c_idx, True), + rotated=True, + ) + + image_with_obb = ann.result() + ``` + +Names can be used from `model.names` when [working with detection results](../modes/predict.md#working-with-results). +Also see the [`Annotator` Reference Page](../reference/utils/plotting.md/#ultralytics.utils.plotting.Annotator) for additional insight. + +#### Ultralytics Sweep Annotation + +!!! example "Sweep Annotation using Ultralytics Utilities" + + ```python + import cv2 + import numpy as np + + from ultralytics import YOLO + from ultralytics.solutions.solutions import SolutionAnnotator + from ultralytics.utils.plotting import colors + + # User defined video path and model file + cap = cv2.VideoCapture("path/to/video.mp4") + model = YOLO(model="yolo11s-seg.pt") # Model file i.e. yolo11s.pt or yolo11m-seg.pt + + if not cap.isOpened(): + print("Error: Could not open video.") + exit() + + # Initialize the video writer object. + w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) + video_writer = cv2.VideoWriter("ultralytics.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) + + masks = None # Initialize variable to store masks data + f = 0 # Initialize frame count variable for enabling mouse event. + line_x = w # Store width of line. + dragging = False # Initialize bool variable for line dragging. + classes = model.names # Store model classes names for plotting. + window_name = "Ultralytics Sweep Annotator" + + + def drag_line(event, x, _, flags, param): + """Mouse callback function to enable dragging a vertical sweep line across the video frame.""" + global line_x, dragging + if event == cv2.EVENT_LBUTTONDOWN or (flags & cv2.EVENT_FLAG_LBUTTON): + line_x = max(0, min(x, w)) + dragging = True + + + while cap.isOpened(): # Loop over the video capture object. + ret, im0 = cap.read() + if not ret: + break + f = f + 1 # Increment frame count. + count = 0 # Re-initialize count variable on every frame for precise counts. + results = model.track(im0, persist=True)[0] + + if f == 1: + cv2.namedWindow(window_name) + cv2.setMouseCallback(window_name, drag_line) + + annotator = SolutionAnnotator(im0) + + if results.boxes.is_track: + if results.masks is not None: + masks = [np.array(m, dtype=np.int32) for m in results.masks.xy] + + boxes = results.boxes.xyxy.tolist() + track_ids = results.boxes.id.int().cpu().tolist() + clss = results.boxes.cls.cpu().tolist() + + for mask, box, cls, t_id in zip(masks or [None] * len(boxes), boxes, clss, track_ids): + color = colors(t_id, True) # Assign different color to each tracked object. + label = f"{classes[cls]}:{t_id}" + if mask is not None and mask.size > 0: + if box[0] > line_x: + count += 1 + cv2.polylines(im0, [mask], True, color, 2) + x, y = mask.min(axis=0) + (w_m, _), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) + cv2.rectangle(im0, (x, y - 20), (x + w_m, y), color, -1) + cv2.putText(im0, label, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) + else: + if box[0] > line_x: + count += 1 + annotator.box_label(box=box, color=color, label=label) + + # Generate draggable sweep line + annotator.sweep_annotator(line_x=line_x, line_y=h, label=f"COUNT:{count}") + + cv2.imshow(window_name, im0) + video_writer.write(im0) + if cv2.waitKey(1) & 0xFF == ord("q"): + break + + # Release the resources + cap.release() + video_writer.release() + cv2.destroyAllWindows() + ``` + Find additional details about the `sweep_annotator` method in our reference section [here](../reference/solutions/solutions.md/#ultralytics.solutions.solutions.SolutionAnnotator.sweep_annotator). -#### Horizontal Bounding Boxes +#### Adaptive label Annotation -```python -import cv2 as cv -import numpy as np +!!! warning -from ultralytics.utils.plotting import Annotator, colors + Starting from **Ultralytics v8.3.167**, `circle_label` and `text_label` have been replaced by a unified `adaptive_label` function. You can now specify the annotation type using the `shape` argument: -names = { - 0: "person", - 5: "bus", - 11: "stop sign", -} - -image = cv.imread("ultralytics/assets/bus.jpg") -ann = Annotator( - image, - line_width=None, # default auto-size - font_size=None, # default auto-size - font="Arial.ttf", # must be ImageFont compatible - pil=False, # use PIL, otherwise uses OpenCV -) - -xyxy_boxes = np.array( - [ - [5, 22.878, 231.27, 804.98, 756.83], # class-idx x1 y1 x2 y2 - [0, 48.552, 398.56, 245.35, 902.71], - [0, 669.47, 392.19, 809.72, 877.04], - [0, 221.52, 405.8, 344.98, 857.54], - [0, 0, 550.53, 63.01, 873.44], - [11, 0.0584, 254.46, 32.561, 324.87], - ] -) - -for nb, box in enumerate(xyxy_boxes): - c_idx, *box = box - label = f"{str(nb).zfill(2)}:{names.get(int(c_idx))}" - ann.box_label(box, label, color=colors(c_idx, bgr=True)) - -image_with_bboxes = ann.result() -``` - -Names can be used from `model.names` when [working with detection results](../modes/predict.md#working-with-results). - -#### Oriented Bounding Boxes (OBB) - -```python -import cv2 as cv -import numpy as np - -from ultralytics.utils.plotting import Annotator, colors - -obb_names = {10: "small vehicle"} -obb_image = cv.imread("datasets/dota8/images/train/P1142__1024__0___824.jpg") -obb_boxes = np.array( - [ - [0, 635, 560, 919, 719, 1087, 420, 803, 261], # class-idx x1 y1 x2 y2 x3 y2 x4 y4 - [0, 331, 19, 493, 260, 776, 70, 613, -171], - [9, 869, 161, 886, 147, 851, 101, 833, 115], - ] -) -ann = Annotator( - obb_image, - line_width=None, # default auto-size - font_size=None, # default auto-size - font="Arial.ttf", # must be ImageFont compatible - pil=False, # use PIL, otherwise uses OpenCV -) -for obb in obb_boxes: - c_idx, *obb = obb - obb = np.array(obb).reshape(-1, 4, 2).squeeze() - label = f"{obb_names.get(int(c_idx))}" - ann.box_label( - obb, - label, - color=colors(c_idx, True), - rotated=True, - ) - -image_with_obb = ann.result() -``` - -#### Bounding Boxes Circle Annotation [Circle Label](https://docs.ultralytics.com/reference/utils/plotting/#ultralytics.utils.plotting.Annotator.circle_label) + * **Rectangle**: `annotator.adaptive_label(box, label=names[int(cls)], color=colors(cls, True), shape="rect")` + * **Circle**: `annotator.adaptive_label(box, label=names[int(cls)], color=colors(cls, True), shape="circle")`


@@ -574,85 +584,88 @@ image_with_obb = ann.result() Watch: In-Depth Guide to Text & Circle Annotations with Python Live Demos | Ultralytics Annotations 🚀

-```python -import cv2 +!!! example "Adaptive label Annotation using Ultralytics Utilities" -from ultralytics import YOLO -from ultralytics.solutions.solutions import SolutionAnnotator -from ultralytics.utils.plotting import colors + === "[Circle Annotation](https://docs.ultralytics.com/reference/utils/plotting/#ultralytics.utils.plotting.Annotator.adaptive_label)" -model = YOLO("yolo11s.pt") -names = model.names -cap = cv2.VideoCapture("path/to/video.mp4") + ```python + import cv2 -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)) + from ultralytics import YOLO + from ultralytics.solutions.solutions import SolutionAnnotator + from ultralytics.utils.plotting import colors -while True: - ret, im0 = cap.read() - if not ret: - break + model = YOLO("yolo11s.pt") + names = model.names + cap = cv2.VideoCapture("path/to/video.mp4") - annotator = SolutionAnnotator(im0) - results = model.predict(im0) - boxes = results[0].boxes.xyxy.cpu() - clss = results[0].boxes.cls.cpu().tolist() + 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)) - for box, cls in zip(boxes, clss): - annotator.circle_label(box, label=names[int(cls)], color=colors(cls, True)) + while True: + ret, im0 = cap.read() + if not ret: + break - writer.write(im0) - cv2.imshow("Ultralytics circle annotation", im0) + annotator = SolutionAnnotator(im0) + results = model.predict(im0)[0] + boxes = results.boxes.xyxy.cpu() + clss = results.boxes.cls.cpu().tolist() - if cv2.waitKey(1) & 0xFF == ord("q"): - break + for box, cls in zip(boxes, clss): + annotator.adaptive_label(box, label=names[int(cls)], color=colors(cls, True), shape="circle") + writer.write(im0) + cv2.imshow("Ultralytics circle annotation", im0) -writer.release() -cap.release() -cv2.destroyAllWindows() -``` + if cv2.waitKey(1) & 0xFF == ord("q"): + break -#### Bounding Boxes Text Annotation [Text Label](https://docs.ultralytics.com/reference/utils/plotting/#ultralytics.utils.plotting.Annotator.text_label) + writer.release() + cap.release() + cv2.destroyAllWindows() + ``` -```python -import cv2 + === "[Text Annotation](https://docs.ultralytics.com/reference/utils/plotting/#ultralytics.utils.plotting.Annotator.adaptive_label)" -from ultralytics import YOLO -from ultralytics.solutions.solutions import SolutionAnnotator -from ultralytics.utils.plotting import colors + ```python + import cv2 -model = YOLO("yolo11s.pt") -names = model.names -cap = cv2.VideoCapture("path/to/video.mp4") + from ultralytics import YOLO + from ultralytics.solutions.solutions import SolutionAnnotator + from ultralytics.utils.plotting import colors -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)) + model = YOLO("yolo11s.pt") + names = model.names + cap = cv2.VideoCapture("path/to/video.mp4") -while True: - ret, im0 = cap.read() - if not ret: - break + 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)) - annotator = SolutionAnnotator(im0) - results = model.predict(im0) - boxes = results[0].boxes.xyxy.cpu() - clss = results[0].boxes.cls.cpu().tolist() + while True: + ret, im0 = cap.read() + if not ret: + break - for box, cls in zip(boxes, clss): - annotator.text_label(box, label=names[int(cls)], color=colors(cls, True)) + annotator = SolutionAnnotator(im0) + results = model.predict(im0)[0] + boxes = results.boxes.xyxy.cpu() + clss = results.boxes.cls.cpu().tolist() - writer.write(im0) - cv2.imshow("Ultralytics text annotation", im0) + for box, cls in zip(boxes, clss): + annotator.adaptive_label(box, label=names[int(cls)], color=colors(cls, True), shape="rect") - if cv2.waitKey(1) & 0xFF == ord("q"): - break + writer.write(im0) + cv2.imshow("Ultralytics text annotation", im0) -writer.release() -cap.release() -cv2.destroyAllWindows() -``` + if cv2.waitKey(1) & 0xFF == ord("q"): + break -See the [`Annotator` Reference Page](../reference/utils/plotting.md#ultralytics.utils.plotting.Annotator) for additional insight. + writer.release() + cap.release() + cv2.destroyAllWindows() + ``` + +See the [`SolutionAnnotator` Reference Page](../reference/solutions/solutions.md/#ultralytics.solutions.solutions.SolutionAnnotator.adaptive_label) for additional insight. ## Miscellaneous diff --git a/ultralytics/solutions/region_counter.py b/ultralytics/solutions/region_counter.py index 7239915d1c..72a89611b5 100644 --- a/ultralytics/solutions/region_counter.py +++ b/ultralytics/solutions/region_counter.py @@ -118,12 +118,13 @@ class RegionCounter(BaseSolution): x1, y1, x2, y2 = map(int, region["polygon"].bounds) pts = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)] annotator.draw_region(pts, region["region_color"], self.line_width * 2) - annotator.text_label( + annotator.adaptive_label( [x1, y1, x2, y2], label=str(region["counts"]), color=region["region_color"], txt_color=region["text_color"], margin=self.line_width * 4, + shape="rect", ) region["counts"] = 0 # Reset for next frame plot_im = annotator.result() diff --git a/ultralytics/solutions/solutions.py b/ultralytics/solutions/solutions.py index 3ce20154fa..f12658be8b 100644 --- a/ultralytics/solutions/solutions.py +++ b/ultralytics/solutions/solutions.py @@ -287,8 +287,7 @@ class SolutionAnnotator(Annotator): display_objects_labels: Annotate bounding boxes with object class labels. sweep_annotator: Visualize a vertical sweep line and optional label. visioneye: Map and connect object centroids to a visual "eye" point. - circle_label: Draw a circular label within a bounding box. - text_label: Draw a rectangular label within a bounding box. + adaptive_label: Draw a circular or rectangle background shape label in center of a bounding box. Examples: >>> annotator = SolutionAnnotator(image) @@ -695,90 +694,58 @@ class SolutionAnnotator(Annotator): cv2.circle(self.im, center_bbox, self.tf * 2, color, -1) cv2.line(self.im, center_point, center_bbox, color, self.tf) - def circle_label( - self, - box: Tuple[float, float, float, float], - label: str = "", - color: Tuple[int, int, int] = (128, 128, 128), - txt_color: Tuple[int, int, int] = (255, 255, 255), - margin: int = 2, - ): - """ - Draw a label with a background circle centered within a given bounding box. - - Args: - box (Tuple[float, float, float, float]): The bounding box coordinates (x1, y1, x2, y2). - label (str): The text label to be displayed. - color (Tuple[int, int, int]): The background color of the circle (B, G, R). - txt_color (Tuple[int, int, int]): The color of the text (R, G, B). - margin (int): The margin between the text and the circle border. - """ - if len(label) > 3: - LOGGER.warning(f"Length of label is {len(label)}, only first 3 letters will be used 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( + def adaptive_label( self, box: Tuple[float, float, float, float], label: str = "", color: Tuple[int, int, int] = (128, 128, 128), txt_color: Tuple[int, int, int] = (255, 255, 255), + shape: str = "rect", margin: int = 5, ): """ - Draw a label with a background rectangle centered within a given bounding box. + Draw a label with a background rectangle or circle centered within a given bounding box. Args: box (Tuple[float, float, float, float]): The bounding box coordinates (x1, y1, x2, y2). label (str): The text label to be displayed. color (Tuple[int, int, int]): The background color of the rectangle (B, G, R). txt_color (Tuple[int, int, int]): The color of the text (R, G, B). + shape (str): The shape of the label i.e "circle" or "rect" margin (int): 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) + if shape == "circle" and len(label) > 3: + LOGGER.warning(f"Length of label is {len(label)}, only first 3 letters will be used for circle annotation.") + label = label[:3] + + x_center, y_center = int((box[0] + box[2]) / 2), int((box[1] + box[3]) / 2) # Calculate center of the bbox + text_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, self.sf - 0.15, self.tf)[0] # Get size of the text + text_x, text_y = x_center - text_size[0] // 2, y_center + text_size[1] // 2 # Calculate top-left corner of text + + if shape == "circle": + cv2.circle( + self.im, + (x_center, y_center), + int(((text_size[0] ** 2 + text_size[1] ** 2) ** 0.5) / 2) + margin, # Calculate the radius + color, + -1, + ) + else: + cv2.rectangle( + self.im, + (text_x - margin, text_y - text_size[1] - margin), # Calculate coordinates of the rectangle + (text_x + text_size[0] + margin, text_y + margin), # Calculate coordinates of the rectangle + color, + -1, + ) + # Draw the text on top of the rectangle cv2.putText( self.im, label, - (text_x, text_y), + (text_x, text_y), # Calculate top-left corner of the text cv2.FONT_HERSHEY_SIMPLEX, - self.sf - 0.1, + self.sf - 0.15, self.get_txt_color(color, txt_color), self.tf, lineType=cv2.LINE_AA,