superiorvision 0.30.0.dev0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (96) hide show
  1. superiorvision/__init__.py +19 -0
  2. superiorvision-0.30.0.dev0.dist-info/METADATA +398 -0
  3. superiorvision-0.30.0.dev0.dist-info/RECORD +96 -0
  4. superiorvision-0.30.0.dev0.dist-info/WHEEL +5 -0
  5. superiorvision-0.30.0.dev0.dist-info/licenses/LICENSE.md +9 -0
  6. superiorvision-0.30.0.dev0.dist-info/top_level.txt +2 -0
  7. supervision/__init__.py +318 -0
  8. supervision/annotators/__init__.py +0 -0
  9. supervision/annotators/base.py +21 -0
  10. supervision/annotators/core.py +3551 -0
  11. supervision/annotators/utils.py +541 -0
  12. supervision/assets/__init__.py +4 -0
  13. supervision/assets/downloader.py +166 -0
  14. supervision/assets/list.py +84 -0
  15. supervision/classification/__init__.py +0 -0
  16. supervision/classification/core.py +216 -0
  17. supervision/config.py +13 -0
  18. supervision/dataset/__init__.py +0 -0
  19. supervision/dataset/core.py +1313 -0
  20. supervision/dataset/formats/__init__.py +0 -0
  21. supervision/dataset/formats/coco.py +708 -0
  22. supervision/dataset/formats/createml.py +331 -0
  23. supervision/dataset/formats/labelme.py +405 -0
  24. supervision/dataset/formats/pascal_voc.py +449 -0
  25. supervision/dataset/formats/yolo.py +502 -0
  26. supervision/dataset/utils.py +265 -0
  27. supervision/detection/__init__.py +0 -0
  28. supervision/detection/compact_mask.py +1927 -0
  29. supervision/detection/core.py +3864 -0
  30. supervision/detection/line_zone.py +924 -0
  31. supervision/detection/overlap_filter.py +426 -0
  32. supervision/detection/tensor_utils.py +1596 -0
  33. supervision/detection/tensor_views.py +48 -0
  34. supervision/detection/tools/__init__.py +0 -0
  35. supervision/detection/tools/csv_sink.py +242 -0
  36. supervision/detection/tools/inference_slicer.py +776 -0
  37. supervision/detection/tools/json_sink.py +215 -0
  38. supervision/detection/tools/polygon_zone.py +266 -0
  39. supervision/detection/tools/smoother.py +218 -0
  40. supervision/detection/tools/transformers.py +265 -0
  41. supervision/detection/utils/__init__.py +12 -0
  42. supervision/detection/utils/_typing.py +11 -0
  43. supervision/detection/utils/boxes.py +510 -0
  44. supervision/detection/utils/converters.py +830 -0
  45. supervision/detection/utils/internal.py +806 -0
  46. supervision/detection/utils/iou_and_nms.py +1606 -0
  47. supervision/detection/utils/masks.py +625 -0
  48. supervision/detection/utils/polygons.py +128 -0
  49. supervision/detection/utils/vlms.py +97 -0
  50. supervision/detection/vlm.py +945 -0
  51. supervision/draw/__init__.py +0 -0
  52. supervision/draw/base.py +14 -0
  53. supervision/draw/color.py +598 -0
  54. supervision/draw/utils.py +527 -0
  55. supervision/geometry/__init__.py +0 -0
  56. supervision/geometry/core.py +208 -0
  57. supervision/geometry/utils.py +54 -0
  58. supervision/key_points/__init__.py +0 -0
  59. supervision/key_points/annotators.py +997 -0
  60. supervision/key_points/core.py +1506 -0
  61. supervision/key_points/skeletons.py +2646 -0
  62. supervision/keypoint/__init__.py +13 -0
  63. supervision/keypoint/annotators.py +13 -0
  64. supervision/keypoint/core.py +8 -0
  65. supervision/metrics/__init__.py +40 -0
  66. supervision/metrics/core.py +74 -0
  67. supervision/metrics/detection.py +1632 -0
  68. supervision/metrics/f1_score.py +815 -0
  69. supervision/metrics/mean_average_precision.py +1723 -0
  70. supervision/metrics/mean_average_recall.py +734 -0
  71. supervision/metrics/precision.py +815 -0
  72. supervision/metrics/recall.py +774 -0
  73. supervision/metrics/utils/__init__.py +0 -0
  74. supervision/metrics/utils/matching.py +56 -0
  75. supervision/metrics/utils/object_size.py +256 -0
  76. supervision/metrics/utils/utils.py +9 -0
  77. supervision/py.typed +0 -0
  78. supervision/tracker/__init__.py +5 -0
  79. supervision/tracker/byte_tracker/__init__.py +0 -0
  80. supervision/tracker/byte_tracker/core.py +448 -0
  81. supervision/tracker/byte_tracker/kalman_filter.py +186 -0
  82. supervision/tracker/byte_tracker/matching.py +211 -0
  83. supervision/tracker/byte_tracker/single_object_track.py +194 -0
  84. supervision/tracker/byte_tracker/utils.py +34 -0
  85. supervision/utils/__init__.py +0 -0
  86. supervision/utils/conversion.py +223 -0
  87. supervision/utils/deprecate.py +54 -0
  88. supervision/utils/file.py +209 -0
  89. supervision/utils/image.py +988 -0
  90. supervision/utils/internal.py +209 -0
  91. supervision/utils/iterables.py +103 -0
  92. supervision/utils/logger.py +61 -0
  93. supervision/utils/notebook.py +124 -0
  94. supervision/utils/tensor.py +362 -0
  95. supervision/utils/video.py +533 -0
  96. supervision/validators/__init__.py +379 -0
@@ -0,0 +1,3551 @@
1
+ from collections.abc import Iterator
2
+ from functools import lru_cache
3
+ from math import sqrt
4
+ from typing import Any, ClassVar, cast
5
+
6
+ import cv2
7
+ import numpy as np
8
+ import numpy.typing as npt
9
+ import torch
10
+ from PIL import Image, ImageDraw, ImageFont
11
+ from scipy.interpolate import splev, splprep
12
+
13
+ from supervision.annotators.base import BaseAnnotator
14
+ from supervision.annotators.utils import (
15
+ PENDING_TRACK_ID,
16
+ ColorLookup,
17
+ Trace,
18
+ _validate_labels,
19
+ calculate_dynamic_kernel_size,
20
+ calculate_dynamic_pixel_size,
21
+ get_labels_text,
22
+ hex_to_rgba,
23
+ resolve_color,
24
+ resolve_text_background_xyxy,
25
+ snap_boxes,
26
+ wrap_text,
27
+ )
28
+ from supervision.config import ORIENTED_BOX_COORDINATES
29
+ from supervision.detection.compact_mask import CompactMask
30
+ from supervision.detection.core import Detections
31
+ from supervision.detection.utils.boxes import clip_boxes, spread_out_boxes
32
+ from supervision.detection.utils.converters import (
33
+ mask_to_polygons,
34
+ polygon_to_mask,
35
+ xyxy_to_polygons,
36
+ )
37
+ from supervision.detection.utils.masks import _masks_to_roi
38
+ from supervision.draw.base import ImageType
39
+ from supervision.draw.color import Color, ColorPalette
40
+ from supervision.draw.utils import draw_polygon, draw_rounded_rectangle, draw_text
41
+ from supervision.geometry.core import Point, Position, Rect
42
+ from supervision.utils.conversion import (
43
+ ensure_cv2_image_for_class_method,
44
+ ensure_pil_image_for_class_method,
45
+ )
46
+ from supervision.utils.deprecate import deprecated, void
47
+ from supervision.utils.image import (
48
+ _overlay_image,
49
+ crop_image,
50
+ letterbox_image,
51
+ scale_image,
52
+ )
53
+ from supervision.utils.logger import _get_logger
54
+ from supervision.utils.tensor import to_numpy
55
+
56
+ logger = _get_logger(__name__)
57
+
58
+
59
+ @lru_cache
60
+ def _load_icon_from_path(
61
+ icon_path: str, icon_resolution_wh: tuple[int, int]
62
+ ) -> npt.NDArray[np.uint8]:
63
+ """Load and resize an icon image through a cache shared by annotators."""
64
+ icon = cv2.imread(icon_path, cv2.IMREAD_UNCHANGED)
65
+ if icon is None:
66
+ raise FileNotFoundError(f"Error: Couldn't load the icon image from {icon_path}")
67
+ icon_array = cast(npt.NDArray[np.uint8], icon)
68
+ result: npt.NDArray[np.uint8] = letterbox_image(
69
+ image=icon_array, resolution_wh=icon_resolution_wh
70
+ )
71
+ return result
72
+
73
+
74
+ def _normalize_color_input(color: Color | ColorPalette | str) -> Color | ColorPalette:
75
+ """Normalize accepted color inputs to internal color objects.
76
+
77
+ Accepts `Color`, `ColorPalette`, or hex string input. Hex strings are parsed via
78
+ `hex_to_rgba` and converted to `Color` (alpha channel is ignored because annotator
79
+ drawing uses RGB/BGR colors).
80
+ """
81
+ if isinstance(color, str):
82
+ r, g, b, _ = hex_to_rgba(color)
83
+ return Color.from_rgb_tuple((r, g, b))
84
+ return color
85
+
86
+
87
+ CV2_FONT = cv2.FONT_HERSHEY_SIMPLEX
88
+
89
+
90
+ class _BaseLabelAnnotator(BaseAnnotator):
91
+ """
92
+ Base class for annotators that add labels to detections.
93
+
94
+ Attributes:
95
+ color: The color to use for the label background.
96
+ color_lookup: The method used to determine the color of the label.
97
+ text_color: The color to use for the label text.
98
+ text_padding: The padding around the label text, in pixels.
99
+ text_anchor: The position of the text relative to the detection
100
+ bounding box.
101
+ text_offset: A tuple of 2D coordinates `(x, y)` to
102
+ offset the text position from the anchor point, in pixels.
103
+ border_radius: The radius of the label background corners, in pixels.
104
+ smart_position: Whether to intelligently adjust the label position to
105
+ avoid overlapping with other elements.
106
+ max_line_length: Maximum number of characters per line before
107
+ wrapping the text. None means no wrapping.
108
+ """
109
+
110
+ def __init__(
111
+ self,
112
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
113
+ color_lookup: ColorLookup = ColorLookup.CLASS,
114
+ text_color: Color | ColorPalette | str = Color.WHITE,
115
+ text_padding: int = 10,
116
+ text_position: Position = Position.TOP_LEFT,
117
+ text_offset: tuple[int, int] = (0, 0),
118
+ border_radius: int = 0,
119
+ smart_position: bool = False,
120
+ max_line_length: int | None = None,
121
+ ):
122
+ """
123
+ Initializes the _BaseLabelAnnotator.
124
+
125
+ Args:
126
+ color: The color to use for the label
127
+ background.
128
+ color_lookup: The method used to determine the color
129
+ of the label
130
+ text_color: The color to use for the
131
+ label text.
132
+ text_padding: The padding around the label text, in pixels.
133
+ text_position: The position of the text relative to the
134
+ detection bounding box.
135
+ text_offset: A tuple of 2D coordinates
136
+ `(x, y)` to offset the text position from the anchor point, in pixels.
137
+ border_radius: The radius of the label background corners,
138
+ in pixels.
139
+ smart_position: Whether to intelligently adjust the label
140
+ position to avoid overlapping with other elements.
141
+ max_line_length: Maximum number of characters per
142
+ line before wrapping the text. None means no wrapping.
143
+ """
144
+ self.color: Color | ColorPalette = _normalize_color_input(color)
145
+ self.color_lookup: ColorLookup = color_lookup
146
+ self.text_color: Color | ColorPalette = _normalize_color_input(text_color)
147
+ self.text_padding: int = text_padding
148
+ self.text_anchor: Position = text_position
149
+ self.text_offset: tuple[int, int] = text_offset
150
+ self.border_radius: int = border_radius
151
+ self.smart_position = smart_position
152
+ self.max_line_length: int | None = max_line_length
153
+
154
+ def _adjust_labels_in_frame(
155
+ self,
156
+ resolution_wh: tuple[int, int],
157
+ labels: list[str],
158
+ label_properties: npt.NDArray[np.float32],
159
+ ) -> npt.NDArray[np.float32]:
160
+ """
161
+ Adjusts the position of labels to ensure they stay within the frame boundaries.
162
+
163
+ Args:
164
+ resolution_wh: The width and height of the frame.
165
+ labels: The list of text labels.
166
+ label_properties: An array of label properties, where each row
167
+ contains [x1, y1, x2, y2, text_height, ...].
168
+
169
+ Returns:
170
+ The adjusted label properties.
171
+ """
172
+ adjusted_properties = label_properties.copy()
173
+
174
+ # First, make sure the boxes don't go outside the frame
175
+ adjusted_properties[:, :4] = snap_boxes(
176
+ adjusted_properties[:, :4],
177
+ resolution_wh,
178
+ )
179
+
180
+ # Apply the spread out algorithm to avoid box overlaps
181
+ if len(labels) > 1:
182
+ # Extract the box coordinates
183
+ boxes = adjusted_properties[:, :4]
184
+ # Use the spread_out_boxes function to adjust overlapping boxes
185
+ spread_boxes = spread_out_boxes(boxes)
186
+ # Update the properties with the spread out boxes
187
+ adjusted_properties[:, :4] = spread_boxes
188
+
189
+ # Additional check to ensure boxes are still within frame after spreading
190
+ adjusted_properties[:, :4] = snap_boxes(
191
+ adjusted_properties[:, :4], resolution_wh
192
+ )
193
+
194
+ return cast(
195
+ npt.NDArray[np.float32],
196
+ np.asarray(adjusted_properties, dtype=np.float32),
197
+ )
198
+
199
+
200
+ class BoxAnnotator(BaseAnnotator):
201
+ """
202
+ A class for drawing bounding boxes on an image using provided detections.
203
+ """
204
+
205
+ def __init__(
206
+ self,
207
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
208
+ thickness: int = 2,
209
+ color_lookup: ColorLookup = ColorLookup.CLASS,
210
+ ):
211
+ """
212
+ Args:
213
+ color: The color or color palette to use for
214
+ annotating detections.
215
+ thickness: Thickness of the bounding box lines.
216
+ color_lookup: Strategy for mapping colors to annotations.
217
+ Options are `INDEX`, `CLASS`, `TRACK`.
218
+ """
219
+ self.color: Color | ColorPalette = _normalize_color_input(color)
220
+ self.thickness: int = thickness
221
+ self.color_lookup: ColorLookup = color_lookup
222
+
223
+ @ensure_cv2_image_for_class_method
224
+ def annotate(
225
+ self,
226
+ scene: ImageType,
227
+ detections: Detections,
228
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
229
+ ) -> ImageType:
230
+ """
231
+ Annotates the given scene with bounding boxes based on the provided detections.
232
+
233
+ Args:
234
+ scene: The image where bounding boxes will be drawn. `ImageType`
235
+ is a flexible type, accepting either `numpy.ndarray` or
236
+ `PIL.Image.Image`.
237
+ detections: Object detections to annotate.
238
+ custom_color_lookup: Custom color lookup array.
239
+ Allows to override the default color mapping strategy.
240
+
241
+ Returns:
242
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
243
+ or `PIL.Image.Image`)
244
+
245
+ Examples:
246
+ ```pycon
247
+ >>> import numpy as np
248
+ >>> import supervision as sv
249
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
250
+ >>> detections = sv.Detections(
251
+ ... xyxy=np.array([[20, 20, 80, 80]]),
252
+ ... class_id=np.array([0])
253
+ ... )
254
+ >>> box_annotator = sv.BoxAnnotator()
255
+ >>> annotated_frame = box_annotator.annotate(
256
+ ... scene=image.copy(),
257
+ ... detections=detections
258
+ ... )
259
+
260
+ ```
261
+
262
+ ![bounding-box-annotator-example](https://media.roboflow.com/
263
+ supervision-annotator-examples/bounding-box-annotator-example-purple.png)
264
+ """
265
+ if not isinstance(scene, np.ndarray):
266
+ return scene
267
+ xyxy = to_numpy(detections.xyxy).astype(int)
268
+ for detection_idx, (x1, y1, x2, y2) in enumerate(xyxy):
269
+ color = resolve_color(
270
+ color=self.color,
271
+ detections=detections,
272
+ detection_idx=detection_idx,
273
+ color_lookup=self.color_lookup
274
+ if custom_color_lookup is None
275
+ else custom_color_lookup,
276
+ )
277
+ cv2.rectangle(
278
+ img=scene,
279
+ pt1=(x1, y1),
280
+ pt2=(x2, y2),
281
+ color=color.as_bgr(),
282
+ thickness=self.thickness,
283
+ )
284
+ return scene
285
+
286
+
287
+ class OrientedBoxAnnotator(BaseAnnotator):
288
+ """
289
+ A class for drawing oriented bounding boxes on an image using provided detections.
290
+ """
291
+
292
+ def __init__(
293
+ self,
294
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
295
+ thickness: int = 2,
296
+ color_lookup: ColorLookup = ColorLookup.CLASS,
297
+ ):
298
+ """
299
+ Args:
300
+ color: The color or color palette to use for
301
+ annotating detections.
302
+ thickness: Thickness of the bounding box lines.
303
+ color_lookup: Strategy for mapping colors to annotations.
304
+ Options are `INDEX`, `CLASS`, `TRACK`.
305
+ """
306
+ self.color: Color | ColorPalette = _normalize_color_input(color)
307
+ self.thickness: int = thickness
308
+ self.color_lookup: ColorLookup = color_lookup
309
+
310
+ @ensure_cv2_image_for_class_method
311
+ def annotate(
312
+ self,
313
+ scene: ImageType,
314
+ detections: Detections,
315
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
316
+ ) -> ImageType:
317
+ """
318
+ Annotates the given scene with oriented bounding boxes based on the
319
+ provided detections.
320
+
321
+ Args:
322
+ scene: The image where bounding boxes will be drawn.
323
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
324
+ or `PIL.Image.Image`.
325
+ detections: Object detections to annotate.
326
+ custom_color_lookup: Custom color lookup array.
327
+ Allows to override the default color mapping strategy.
328
+
329
+ Returns:
330
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
331
+ or `PIL.Image.Image`)
332
+
333
+ Example:
334
+ ```python
335
+ import cv2
336
+ import supervision as sv
337
+ from ultralytics import YOLO
338
+
339
+ image = cv2.imread("<SOURCE_IMAGE_PATH>")
340
+ model = YOLO("yolov8n-obb.pt")
341
+
342
+ result = model(image)[0]
343
+ detections = sv.Detections.from_ultralytics(result)
344
+
345
+ oriented_box_annotator = sv.OrientedBoxAnnotator()
346
+ annotated_frame = oriented_box_annotator.annotate(
347
+ scene=image.copy(),
348
+ detections=detections
349
+ )
350
+ ```
351
+ """
352
+ if not isinstance(scene, np.ndarray):
353
+ return scene
354
+ if detections.data is None or ORIENTED_BOX_COORDINATES not in detections.data:
355
+ return scene
356
+ obb_boxes = to_numpy(detections.data[ORIENTED_BOX_COORDINATES]).astype(int)
357
+
358
+ for detection_idx in range(len(detections)):
359
+ obb = obb_boxes[detection_idx]
360
+ color = resolve_color(
361
+ color=self.color,
362
+ detections=detections,
363
+ detection_idx=detection_idx,
364
+ color_lookup=self.color_lookup
365
+ if custom_color_lookup is None
366
+ else custom_color_lookup,
367
+ )
368
+
369
+ cv2.drawContours(scene, [obb], 0, color.as_bgr(), self.thickness)
370
+
371
+ return scene
372
+
373
+
374
+ # --- Shared mask-painting utilities ---
375
+ def _iter_mask_crops(
376
+ detections: Detections,
377
+ ) -> Iterator[tuple[int, npt.NDArray[np.bool_], npt.NDArray[np.int32] | None]]:
378
+ """Yield ``(detection_idx, mask_or_crop, offset_or_None)`` for each mask.
379
+
380
+ Encapsulates the ``CompactMask`` vs dense dispatch so individual annotators
381
+ do not need inline ``isinstance`` checks. For ``CompactMask`` inputs yields
382
+ the bbox crop and its ``(x1, y1)`` image-space origin; for dense masks
383
+ yields the full-frame boolean slice with ``offset=None``.
384
+
385
+ Args:
386
+ detections: Object detections whose masks to iterate.
387
+
388
+ Yields:
389
+ Tuple of ``(detection_idx, mask_or_crop, offset_or_None)``.
390
+ ``mask_or_crop`` is boolean (crop-sized for ``CompactMask``, full-frame
391
+ for dense). ``offset_or_None`` is an int32 ``(x1, y1)`` array for
392
+ crop→image translation, or ``None`` for dense masks.
393
+ """
394
+ masks = detections.mask
395
+ if masks is None:
396
+ return
397
+ # TODO: replace isinstance dispatch with a MaskLike Protocol (separate PR)
398
+ compact_mask = masks if isinstance(masks, CompactMask) else None
399
+ for detection_idx in range(len(detections)):
400
+ if compact_mask is None:
401
+ yield (
402
+ detection_idx,
403
+ cast(npt.NDArray[np.bool_], masks[detection_idx]),
404
+ None,
405
+ )
406
+ else:
407
+ yield (
408
+ detection_idx,
409
+ compact_mask.crop(detection_idx),
410
+ compact_mask.offsets[detection_idx],
411
+ )
412
+
413
+
414
+ def _paint_masks_by_area(
415
+ canvas: npt.NDArray[np.uint8],
416
+ detections: Detections,
417
+ color: Color | ColorPalette,
418
+ color_lookup: ColorLookup | npt.NDArray[np.int_],
419
+ collect_union: bool = False,
420
+ canvas_origin: tuple[int, int] = (0, 0),
421
+ ) -> npt.NDArray[np.bool_] | None:
422
+ """Paint each detection's mask into `canvas` in descending-area order.
423
+
424
+ Smaller masks are drawn on top of larger ones. `CompactMask` detections
425
+ are painted into their bounding-box crop only, avoiding a full `(H, W)`
426
+ allocation per mask; dense masks fall back to full-frame boolean indexing.
427
+
428
+ Args:
429
+ canvas: BGR image array painted in place. Shape ``(H, W, 3)``.
430
+ detections: Detections whose masks to paint. Returns immediately
431
+ without modifying `canvas` when ``detections.mask`` is ``None``.
432
+ color: Single color or palette used to resolve each detection's color.
433
+ color_lookup: Strategy for mapping colors to detection indices.
434
+ collect_union: When ``True``, allocate and return a ``(H, W)``
435
+ boolean array that accumulates the union of all painted masks
436
+ (useful for callers like `HaloAnnotator` that need the combined
437
+ mask footprint). When ``False`` (default), returns ``None``.
438
+ canvas_origin: Absolute ``(x, y)`` origin of `canvas` within the source
439
+ image. Use the default for full-frame painting.
440
+
441
+ Returns:
442
+ A boolean array matching the canvas dimensions when
443
+ ``collect_union=True``, otherwise ``None``. When called with an
444
+ ROI sub-canvas, dimensions are the ROI size, not the full image.
445
+ """
446
+ masks = detections.mask
447
+ if masks is None:
448
+ return None
449
+ union: npt.NDArray[np.bool_] | None = (
450
+ np.zeros(canvas.shape[:2], dtype=bool) if collect_union else None
451
+ )
452
+ compact_mask = masks if isinstance(masks, CompactMask) else None
453
+ origin_x, origin_y = canvas_origin
454
+ canvas_h, canvas_w = canvas.shape[:2]
455
+ sorted_indices = torch.argsort(detections.area, descending=True).cpu().tolist()
456
+ for detection_idx in sorted_indices:
457
+ color_bgr = resolve_color(
458
+ color=color,
459
+ detections=detections,
460
+ detection_idx=detection_idx,
461
+ color_lookup=color_lookup,
462
+ ).as_bgr()
463
+ if compact_mask is not None:
464
+ x1 = int(compact_mask.offsets[detection_idx, 0])
465
+ y1 = int(compact_mask.offsets[detection_idx, 1])
466
+ crop_m = compact_mask.crop(detection_idx)
467
+ crop_h, crop_w = crop_m.shape
468
+ crop_x1 = max(0, origin_x - x1)
469
+ crop_y1 = max(0, origin_y - y1)
470
+ canvas_x1 = max(0, x1 - origin_x)
471
+ canvas_y1 = max(0, y1 - origin_y)
472
+ paint_w = min(crop_w - crop_x1, canvas_w - canvas_x1)
473
+ paint_h = min(crop_h - crop_y1, canvas_h - canvas_y1)
474
+ if paint_w <= 0 or paint_h <= 0:
475
+ continue
476
+ crop_slice = crop_m[
477
+ crop_y1 : crop_y1 + paint_h, crop_x1 : crop_x1 + paint_w
478
+ ]
479
+ crop_slice = to_numpy(crop_slice).astype(bool, copy=False)
480
+ canvas_slice = canvas[
481
+ canvas_y1 : canvas_y1 + paint_h,
482
+ canvas_x1 : canvas_x1 + paint_w,
483
+ ]
484
+ canvas_slice[crop_slice] = color_bgr
485
+ if union is not None:
486
+ union[
487
+ canvas_y1 : canvas_y1 + paint_h,
488
+ canvas_x1 : canvas_x1 + paint_w,
489
+ ] |= crop_slice
490
+ else:
491
+ mask = to_numpy(masks[detection_idx]).astype(bool, copy=False)
492
+ mask = mask[origin_y : origin_y + canvas_h, origin_x : origin_x + canvas_w]
493
+ canvas[mask] = color_bgr
494
+ if union is not None:
495
+ union |= mask
496
+ return union
497
+
498
+
499
+ class MaskAnnotator(BaseAnnotator):
500
+ """
501
+ A class for drawing masks on an image using provided detections.
502
+
503
+ !!! warning
504
+
505
+ This annotator uses `sv.Detections.mask`.
506
+ """
507
+
508
+ requires_mask = True
509
+
510
+ def __init__(
511
+ self,
512
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
513
+ opacity: float = 0.5,
514
+ color_lookup: ColorLookup = ColorLookup.CLASS,
515
+ ):
516
+ """
517
+ Args:
518
+ color: The color or color palette to use for
519
+ annotating detections.
520
+ opacity: Opacity of the overlay mask. Must be between `0` and `1`.
521
+ color_lookup: Strategy for mapping colors to annotations.
522
+ Options are `INDEX`, `CLASS`, `TRACK`.
523
+ """
524
+ self.color: Color | ColorPalette = _normalize_color_input(color)
525
+ self.opacity = opacity
526
+ self.color_lookup: ColorLookup = color_lookup
527
+
528
+ @ensure_cv2_image_for_class_method
529
+ def annotate(
530
+ self,
531
+ scene: ImageType,
532
+ detections: Detections,
533
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
534
+ ) -> ImageType:
535
+ """
536
+ Annotates the given scene with masks based on the provided detections.
537
+
538
+ Args:
539
+ scene: The image where masks will be drawn.
540
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
541
+ or `PIL.Image.Image`.
542
+ detections: Object detections to annotate.
543
+ custom_color_lookup: Custom color lookup array.
544
+ Allows to override the default color mapping strategy.
545
+
546
+ Returns:
547
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
548
+ or `PIL.Image.Image`)
549
+
550
+ Examples:
551
+ ```pycon
552
+ >>> import numpy as np
553
+ >>> import supervision as sv
554
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
555
+ >>> detections = sv.Detections(
556
+ ... xyxy=np.array([[20, 20, 80, 80]]),
557
+ ... mask=np.zeros((1, 100, 100), dtype=bool),
558
+ ... class_id=np.array([0])
559
+ ... )
560
+ >>> mask_annotator = sv.MaskAnnotator()
561
+ >>> annotated_frame = mask_annotator.annotate(
562
+ ... scene=image.copy(),
563
+ ... detections=detections
564
+ ... )
565
+
566
+ ```
567
+
568
+ ![mask-annotator-example](https://media.roboflow.com/
569
+ supervision-annotator-examples/mask-annotator-example-purple.png)
570
+ """
571
+ if not isinstance(scene, np.ndarray):
572
+ return scene
573
+ if detections.mask is None:
574
+ return scene
575
+
576
+ image_shape = (int(scene.shape[0]), int(scene.shape[1]))
577
+ effective_lookup = (
578
+ self.color_lookup if custom_color_lookup is None else custom_color_lookup
579
+ )
580
+ if len(detections) > 0:
581
+ resolve_color(
582
+ color=self.color,
583
+ detections=detections,
584
+ detection_idx=0,
585
+ color_lookup=effective_lookup,
586
+ )
587
+ roi = _masks_to_roi(detections.mask, image_shape, detections.xyxy)
588
+ if roi is None:
589
+ return scene
590
+
591
+ x1, y1, x2, y2 = roi
592
+ scene_roi = scene[y1:y2, x1:x2]
593
+ colored_mask = np.array(scene_roi, copy=True, dtype=np.uint8)
594
+ _paint_masks_by_area(
595
+ colored_mask,
596
+ detections,
597
+ self.color,
598
+ effective_lookup,
599
+ canvas_origin=(x1, y1),
600
+ )
601
+ tmp = cv2.addWeighted(
602
+ colored_mask, self.opacity, scene_roi.copy(), 1 - self.opacity, 0
603
+ )
604
+ scene_roi[:] = tmp
605
+ return scene
606
+
607
+
608
+ class PolygonAnnotator(BaseAnnotator):
609
+ """
610
+ A class for drawing polygons on an image using provided detections.
611
+
612
+ !!! warning
613
+
614
+ This annotator uses `sv.Detections.mask`.
615
+ """
616
+
617
+ requires_mask = True
618
+
619
+ def __init__(
620
+ self,
621
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
622
+ thickness: int = 2,
623
+ color_lookup: ColorLookup = ColorLookup.CLASS,
624
+ ):
625
+ """
626
+ Args:
627
+ color: The color or color palette to use for
628
+ annotating detections.
629
+ thickness: Thickness of the polygon lines.
630
+ color_lookup: Strategy for mapping colors to annotations.
631
+ Options are `INDEX`, `CLASS`, `TRACK`.
632
+ """
633
+ self.color: Color | ColorPalette = _normalize_color_input(color)
634
+ self.thickness: int = thickness
635
+ self.color_lookup: ColorLookup = color_lookup
636
+
637
+ @ensure_cv2_image_for_class_method
638
+ def annotate(
639
+ self,
640
+ scene: ImageType,
641
+ detections: Detections,
642
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
643
+ ) -> ImageType:
644
+ """
645
+ Annotates the given scene with polygons based on the provided detections.
646
+
647
+ Args:
648
+ scene: The image where polygons will be drawn.
649
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
650
+ or `PIL.Image.Image`.
651
+ detections: Object detections to annotate.
652
+ custom_color_lookup: Custom color lookup array.
653
+ Allows to override the default color mapping strategy.
654
+
655
+ Returns:
656
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
657
+ or `PIL.Image.Image`)
658
+
659
+ Examples:
660
+ ```pycon
661
+ >>> import numpy as np
662
+ >>> import supervision as sv
663
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
664
+ >>> detections = sv.Detections(
665
+ ... xyxy=np.array([[20, 20, 80, 80]]),
666
+ ... class_id=np.array([0])
667
+ ... )
668
+ >>> polygon_annotator = sv.PolygonAnnotator()
669
+ >>> annotated_frame = polygon_annotator.annotate(
670
+ ... scene=image.copy(),
671
+ ... detections=detections
672
+ ... )
673
+
674
+ ```
675
+
676
+ Note:
677
+ When `detections.mask` is a `CompactMask`, each detection's polygon
678
+ is decoded from a bbox-sized crop (O(crop_area)) rather than a
679
+ full-frame ``(H, W)`` allocation (O(H·W)). Polygon coordinates are
680
+ shifted from crop-local space to image space via the stored
681
+ ``(x1, y1)`` bbox origin. Pixels outside the declared ``xyxy`` box
682
+ are not represented in compact storage and will not be drawn.
683
+
684
+ ![polygon-annotator-example](https://media.roboflow.com/
685
+ supervision-annotator-examples/polygon-annotator-example-purple.png)
686
+ """
687
+ if not isinstance(scene, np.ndarray):
688
+ return scene
689
+ if detections.mask is None:
690
+ return scene
691
+
692
+ for detection_idx, mask, offset in _iter_mask_crops(detections):
693
+ color = resolve_color(
694
+ color=self.color,
695
+ detections=detections,
696
+ detection_idx=detection_idx,
697
+ color_lookup=self.color_lookup
698
+ if custom_color_lookup is None
699
+ else custom_color_lookup,
700
+ )
701
+ for polygon in mask_to_polygons(mask=mask):
702
+ if offset is not None:
703
+ # translate crop-local polygon to image space via (x1, y1) origin
704
+ polygon = polygon + to_numpy(offset)
705
+ scene = draw_polygon(
706
+ scene=scene,
707
+ polygon=cast(npt.NDArray[np.int_], polygon),
708
+ color=color,
709
+ thickness=self.thickness,
710
+ )
711
+
712
+ return scene
713
+
714
+
715
+ class ColorAnnotator(BaseAnnotator):
716
+ """
717
+ A class for drawing box masks on an image using provided detections.
718
+ """
719
+
720
+ def __init__(
721
+ self,
722
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
723
+ opacity: float = 0.5,
724
+ color_lookup: ColorLookup = ColorLookup.CLASS,
725
+ ):
726
+ """
727
+ Args:
728
+ color: The color or color palette to use for
729
+ annotating detections.
730
+ opacity: Opacity of the overlay mask. Must be between `0` and `1`.
731
+ color_lookup: Strategy for mapping colors to annotations.
732
+ Options are `INDEX`, `CLASS`, `TRACK`.
733
+ """
734
+ self.color: Color | ColorPalette = _normalize_color_input(color)
735
+ self.color_lookup: ColorLookup = color_lookup
736
+ self.opacity = opacity
737
+
738
+ @ensure_cv2_image_for_class_method
739
+ def annotate(
740
+ self,
741
+ scene: ImageType,
742
+ detections: Detections,
743
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
744
+ ) -> ImageType:
745
+ """
746
+ Annotates the given scene with box masks based on the provided detections.
747
+
748
+ Args:
749
+ scene: The image where bounding boxes will be drawn.
750
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
751
+ or `PIL.Image.Image`.
752
+ detections: Object detections to annotate.
753
+ custom_color_lookup: Custom color lookup array.
754
+ Allows to override the default color mapping strategy.
755
+
756
+ Returns:
757
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
758
+ or `PIL.Image.Image`)
759
+
760
+ Examples:
761
+ ```pycon
762
+ >>> import numpy as np
763
+ >>> import supervision as sv
764
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
765
+ >>> detections = sv.Detections(
766
+ ... xyxy=np.array([[20, 20, 80, 80]]),
767
+ ... class_id=np.array([0])
768
+ ... )
769
+ >>> color_annotator = sv.ColorAnnotator()
770
+ >>> annotated_frame = color_annotator.annotate(
771
+ ... scene=image.copy(),
772
+ ... detections=detections
773
+ ... )
774
+
775
+ ```
776
+
777
+ ![box-mask-annotator-example](https://media.roboflow.com/
778
+ supervision-annotator-examples/box-mask-annotator-example-purple.png)
779
+ """
780
+ if not isinstance(scene, np.ndarray):
781
+ return scene
782
+ scene_with_boxes = scene.copy()
783
+ xyxy = to_numpy(detections.xyxy).astype(int)
784
+ for detection_idx, (x1, y1, x2, y2) in enumerate(xyxy):
785
+ color = resolve_color(
786
+ color=self.color,
787
+ detections=detections,
788
+ detection_idx=detection_idx,
789
+ color_lookup=self.color_lookup
790
+ if custom_color_lookup is None
791
+ else custom_color_lookup,
792
+ )
793
+ cv2.rectangle(
794
+ img=scene_with_boxes,
795
+ pt1=(x1, y1),
796
+ pt2=(x2, y2),
797
+ color=color.as_bgr(),
798
+ thickness=-1,
799
+ )
800
+
801
+ cv2.addWeighted(
802
+ scene_with_boxes, self.opacity, scene, 1 - self.opacity, gamma=0, dst=scene
803
+ )
804
+ return scene
805
+
806
+
807
+ class HaloAnnotator(BaseAnnotator):
808
+ """
809
+ A class for drawing Halos on an image using provided detections.
810
+
811
+ !!! warning
812
+
813
+ This annotator uses `sv.Detections.mask`.
814
+ """
815
+
816
+ requires_mask = True
817
+
818
+ def __init__(
819
+ self,
820
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
821
+ opacity: float = 0.8,
822
+ kernel_size: int = 40,
823
+ color_lookup: ColorLookup = ColorLookup.CLASS,
824
+ ):
825
+ """
826
+ Args:
827
+ color: The color or color palette to use for
828
+ annotating detections.
829
+ opacity: Opacity of the overlay mask. Must be between `0` and `1`.
830
+ kernel_size: The size of the average pooling kernel used for creating
831
+ the halo.
832
+ color_lookup: Strategy for mapping colors to annotations.
833
+ Options are `INDEX`, `CLASS`, `TRACK`.
834
+ """
835
+ self.color: Color | ColorPalette = _normalize_color_input(color)
836
+ self.opacity = opacity
837
+ self.color_lookup: ColorLookup = color_lookup
838
+ self.kernel_size: int = kernel_size
839
+
840
+ @ensure_cv2_image_for_class_method
841
+ def annotate(
842
+ self,
843
+ scene: ImageType,
844
+ detections: Detections,
845
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
846
+ ) -> ImageType:
847
+ """
848
+ Annotates the given scene with halos based on the provided detections.
849
+
850
+ Args:
851
+ scene: The image where the halo effect will be applied.
852
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
853
+ or `PIL.Image.Image`.
854
+ detections: Object detections to annotate.
855
+ custom_color_lookup: Custom color lookup array.
856
+ Allows to override the default color mapping strategy.
857
+
858
+ Returns:
859
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
860
+ or `PIL.Image.Image`)
861
+
862
+ Examples:
863
+ ```pycon
864
+ >>> import numpy as np
865
+ >>> import supervision as sv
866
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
867
+ >>> detections = sv.Detections(
868
+ ... xyxy=np.array([[20, 20, 80, 80]]),
869
+ ... mask=np.zeros((1, 100, 100), dtype=bool),
870
+ ... class_id=np.array([0])
871
+ ... )
872
+ >>> halo_annotator = sv.HaloAnnotator()
873
+ >>> annotated_frame = halo_annotator.annotate(
874
+ ... scene=image.copy(),
875
+ ... detections=detections
876
+ ... )
877
+
878
+ ```
879
+
880
+ ![halo-annotator-example](https://media.roboflow.com/
881
+ supervision-annotator-examples/halo-annotator-example-purple.png)
882
+ """
883
+ if not isinstance(scene, np.ndarray):
884
+ return scene
885
+ if detections.mask is None:
886
+ return scene
887
+ colored_mask = np.zeros_like(scene, dtype=np.uint8)
888
+ fmask = _paint_masks_by_area(
889
+ colored_mask,
890
+ detections,
891
+ self.color,
892
+ self.color_lookup if custom_color_lookup is None else custom_color_lookup,
893
+ collect_union=True,
894
+ )
895
+ assert fmask is not None # collect_union=True always returns an array
896
+
897
+ colored_mask = cast(
898
+ npt.NDArray[np.uint8],
899
+ cv2.blur(colored_mask, (self.kernel_size, self.kernel_size)),
900
+ )
901
+ colored_mask[fmask] = [0, 0, 0]
902
+ gray = cv2.cvtColor(colored_mask, cv2.COLOR_BGR2GRAY)
903
+ gray_max = gray.max()
904
+ if gray_max == 0:
905
+ # no halo to draw (e.g. empty masks); leave the scene untouched
906
+ return scene
907
+ alpha = self.opacity * gray / gray_max
908
+ alpha_mask = alpha[:, :, np.newaxis]
909
+ # Blend in float space so halo opacity cannot wrap around uint8 boundaries.
910
+ blended_scene = np.clip(
911
+ scene.astype(np.float32) * (1 - alpha_mask)
912
+ + colored_mask.astype(np.float32) * alpha_mask,
913
+ 0,
914
+ 255,
915
+ ).astype(np.uint8)
916
+ np.copyto(scene, blended_scene)
917
+ return scene
918
+
919
+
920
+ class EllipseAnnotator(BaseAnnotator):
921
+ """
922
+ A class for drawing ellipses on an image using provided detections.
923
+ """
924
+
925
+ def __init__(
926
+ self,
927
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
928
+ thickness: int = 2,
929
+ start_angle: int = -45,
930
+ end_angle: int = 235,
931
+ color_lookup: ColorLookup = ColorLookup.CLASS,
932
+ ):
933
+ """
934
+ Args:
935
+ color: The color or color palette to use for
936
+ annotating detections.
937
+ thickness: Thickness of the ellipse lines.
938
+ start_angle: Starting angle of the ellipse.
939
+ end_angle: Ending angle of the ellipse.
940
+ color_lookup: Strategy for mapping colors to annotations.
941
+ Options are `INDEX`, `CLASS`, `TRACK`.
942
+ """
943
+ self.color: Color | ColorPalette = _normalize_color_input(color)
944
+ self.thickness: int = thickness
945
+ self.start_angle: int = start_angle
946
+ self.end_angle: int = end_angle
947
+ self.color_lookup: ColorLookup = color_lookup
948
+
949
+ @ensure_cv2_image_for_class_method
950
+ def annotate(
951
+ self,
952
+ scene: ImageType,
953
+ detections: Detections,
954
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
955
+ ) -> ImageType:
956
+ """
957
+ Annotates the given scene with ellipses based on the provided detections.
958
+
959
+ Args:
960
+ scene: The image where ellipses will be drawn.
961
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
962
+ or `PIL.Image.Image`.
963
+ detections: Object detections to annotate.
964
+ custom_color_lookup: Custom color lookup array.
965
+ Allows to override the default color mapping strategy.
966
+
967
+ Returns:
968
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
969
+ or `PIL.Image.Image`)
970
+
971
+ Examples:
972
+ ```pycon
973
+ >>> import numpy as np
974
+ >>> import supervision as sv
975
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
976
+ >>> detections = sv.Detections(
977
+ ... xyxy=np.array([[20, 20, 80, 80]]),
978
+ ... class_id=np.array([0])
979
+ ... )
980
+ >>> ellipse_annotator = sv.EllipseAnnotator()
981
+ >>> annotated_frame = ellipse_annotator.annotate(
982
+ ... scene=image.copy(),
983
+ ... detections=detections
984
+ ... )
985
+
986
+ ```
987
+
988
+ ![ellipse-annotator-example](https://media.roboflow.com/
989
+ supervision-annotator-examples/ellipse-annotator-example-purple.png)
990
+ """
991
+ if not isinstance(scene, np.ndarray):
992
+ return scene
993
+ xyxy = to_numpy(detections.xyxy).astype(int)
994
+ for detection_idx, (x1, _y1, x2, y2) in enumerate(xyxy):
995
+ color = resolve_color(
996
+ color=self.color,
997
+ detections=detections,
998
+ detection_idx=detection_idx,
999
+ color_lookup=self.color_lookup
1000
+ if custom_color_lookup is None
1001
+ else custom_color_lookup,
1002
+ )
1003
+ center = (int((x1 + x2) / 2), y2)
1004
+ width = x2 - x1
1005
+ cv2.ellipse(
1006
+ scene,
1007
+ center=center,
1008
+ axes=(int(width), int(0.35 * width)),
1009
+ angle=0.0,
1010
+ startAngle=self.start_angle,
1011
+ endAngle=self.end_angle,
1012
+ color=color.as_bgr(),
1013
+ thickness=self.thickness,
1014
+ lineType=cv2.LINE_4,
1015
+ )
1016
+ return scene
1017
+
1018
+
1019
+ class BoxCornerAnnotator(BaseAnnotator):
1020
+ """
1021
+ A class for drawing box corners on an image using provided detections.
1022
+ """
1023
+
1024
+ def __init__(
1025
+ self,
1026
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
1027
+ thickness: int = 4,
1028
+ corner_length: int = 15,
1029
+ color_lookup: ColorLookup = ColorLookup.CLASS,
1030
+ ):
1031
+ """
1032
+ Args:
1033
+ color: The color or color palette to use for
1034
+ annotating detections.
1035
+ thickness: Thickness of the corner lines.
1036
+ corner_length: Length of each corner line.
1037
+ color_lookup: Strategy for mapping colors to annotations.
1038
+ Options are `INDEX`, `CLASS`, `TRACK`.
1039
+ """
1040
+ self.color: Color | ColorPalette = _normalize_color_input(color)
1041
+ self.thickness: int = thickness
1042
+ self.corner_length: int = corner_length
1043
+ self.color_lookup: ColorLookup = color_lookup
1044
+
1045
+ @ensure_cv2_image_for_class_method
1046
+ def annotate(
1047
+ self,
1048
+ scene: ImageType,
1049
+ detections: Detections,
1050
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
1051
+ ) -> ImageType:
1052
+ """
1053
+ Annotates the given scene with box corners based on the provided detections.
1054
+
1055
+ Args:
1056
+ scene: The image where box corners will be drawn.
1057
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
1058
+ or `PIL.Image.Image`.
1059
+ detections: Object detections to annotate.
1060
+ custom_color_lookup: Custom color lookup array.
1061
+ Allows to override the default color mapping strategy.
1062
+
1063
+ Returns:
1064
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
1065
+ or `PIL.Image.Image`)
1066
+
1067
+ Examples:
1068
+ ```pycon
1069
+ >>> import numpy as np
1070
+ >>> import supervision as sv
1071
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
1072
+ >>> detections = sv.Detections(
1073
+ ... xyxy=np.array([[20, 20, 80, 80]]),
1074
+ ... class_id=np.array([0])
1075
+ ... )
1076
+ >>> corner_annotator = sv.BoxCornerAnnotator()
1077
+ >>> annotated_frame = corner_annotator.annotate(
1078
+ ... scene=image.copy(),
1079
+ ... detections=detections
1080
+ ... )
1081
+
1082
+ ```
1083
+
1084
+ ![box-corner-annotator-example](https://media.roboflow.com/
1085
+ supervision-annotator-examples/box-corner-annotator-example-purple.png)
1086
+ """
1087
+ if not isinstance(scene, np.ndarray):
1088
+ return scene
1089
+ xyxy = to_numpy(detections.xyxy).astype(int)
1090
+ for detection_idx, (x1, y1, x2, y2) in enumerate(xyxy):
1091
+ color = resolve_color(
1092
+ color=self.color,
1093
+ detections=detections,
1094
+ detection_idx=detection_idx,
1095
+ color_lookup=self.color_lookup
1096
+ if custom_color_lookup is None
1097
+ else custom_color_lookup,
1098
+ )
1099
+ corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
1100
+
1101
+ for x, y in corners:
1102
+ x_end = x + self.corner_length if x == x1 else x - self.corner_length
1103
+ cv2.line(
1104
+ scene, (x, y), (x_end, y), color.as_bgr(), thickness=self.thickness
1105
+ )
1106
+
1107
+ y_end = y + self.corner_length if y == y1 else y - self.corner_length
1108
+ cv2.line(
1109
+ scene, (x, y), (x, y_end), color.as_bgr(), thickness=self.thickness
1110
+ )
1111
+ return scene
1112
+
1113
+
1114
+ class CircleAnnotator(BaseAnnotator):
1115
+ """
1116
+ A class for drawing circle on an image using provided detections.
1117
+ """
1118
+
1119
+ def __init__(
1120
+ self,
1121
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
1122
+ thickness: int = 2,
1123
+ color_lookup: ColorLookup = ColorLookup.CLASS,
1124
+ ):
1125
+ """
1126
+ Args:
1127
+ color: The color or color palette to use for
1128
+ annotating detections.
1129
+ thickness: Thickness of the circle line.
1130
+ color_lookup: Strategy for mapping colors to annotations.
1131
+ Options are `INDEX`, `CLASS`, `TRACK`.
1132
+ """
1133
+
1134
+ self.color: Color | ColorPalette = _normalize_color_input(color)
1135
+ self.thickness: int = thickness
1136
+ self.color_lookup: ColorLookup = color_lookup
1137
+
1138
+ @ensure_cv2_image_for_class_method
1139
+ def annotate(
1140
+ self,
1141
+ scene: ImageType,
1142
+ detections: Detections,
1143
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
1144
+ ) -> ImageType:
1145
+ """
1146
+ Annotates the given scene with circles based on the provided detections.
1147
+
1148
+ Args:
1149
+ scene: The image where box corners will be drawn.
1150
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
1151
+ or `PIL.Image.Image`.
1152
+ detections: Object detections to annotate.
1153
+ custom_color_lookup: Custom color lookup array.
1154
+ Allows to override the default color mapping strategy.
1155
+
1156
+ Returns:
1157
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
1158
+ or `PIL.Image.Image`)
1159
+
1160
+ Examples:
1161
+ ```pycon
1162
+ >>> import numpy as np
1163
+ >>> import supervision as sv
1164
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
1165
+ >>> detections = sv.Detections(
1166
+ ... xyxy=np.array([[20, 20, 80, 80]]),
1167
+ ... class_id=np.array([0])
1168
+ ... )
1169
+ >>> circle_annotator = sv.CircleAnnotator()
1170
+ >>> annotated_frame = circle_annotator.annotate(
1171
+ ... scene=image.copy(),
1172
+ ... detections=detections
1173
+ ... )
1174
+
1175
+ ```
1176
+
1177
+
1178
+ ![circle-annotator-example](https://media.roboflow.com/
1179
+ supervision-annotator-examples/circle-annotator-example-purple.png)
1180
+ """
1181
+ if not isinstance(scene, np.ndarray):
1182
+ return scene
1183
+ xyxy = to_numpy(detections.xyxy).astype(int)
1184
+ for detection_idx, (x1, y1, x2, y2) in enumerate(xyxy):
1185
+ center = ((x1 + x2) // 2, (y1 + y2) // 2)
1186
+ distance = sqrt((x1 - center[0]) ** 2 + (y1 - center[1]) ** 2)
1187
+ color = resolve_color(
1188
+ color=self.color,
1189
+ detections=detections,
1190
+ detection_idx=detection_idx,
1191
+ color_lookup=self.color_lookup
1192
+ if custom_color_lookup is None
1193
+ else custom_color_lookup,
1194
+ )
1195
+ cv2.circle(
1196
+ img=scene,
1197
+ center=center,
1198
+ radius=int(distance),
1199
+ color=color.as_bgr(),
1200
+ thickness=self.thickness,
1201
+ )
1202
+
1203
+ return scene
1204
+
1205
+
1206
+ class DotAnnotator(BaseAnnotator):
1207
+ """
1208
+ A class for drawing dots on an image at specific coordinates based on provided
1209
+ detections.
1210
+ """
1211
+
1212
+ def __init__(
1213
+ self,
1214
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
1215
+ radius: int = 4,
1216
+ position: Position = Position.CENTER,
1217
+ color_lookup: ColorLookup = ColorLookup.CLASS,
1218
+ outline_thickness: int = 0,
1219
+ outline_color: Color | ColorPalette | str = Color.BLACK,
1220
+ ):
1221
+ """
1222
+ Args:
1223
+ color: The color or color palette to use for
1224
+ annotating detections.
1225
+ radius: Radius of the drawn dots.
1226
+ position: The anchor position for placing the dot.
1227
+ color_lookup: Strategy for mapping colors to annotations.
1228
+ Options are `INDEX`, `CLASS`, `TRACK`.
1229
+ outline_thickness: Thickness of the outline of the dot.
1230
+ outline_color: The color or color palette to
1231
+ use for outline. It is activated by setting outline_thickness to a value
1232
+ greater than 0.
1233
+ """
1234
+ self.color: Color | ColorPalette = _normalize_color_input(color)
1235
+ self.radius: int = radius
1236
+ self.position: Position = position
1237
+ self.color_lookup: ColorLookup = color_lookup
1238
+ self.outline_thickness = outline_thickness
1239
+ self.outline_color: Color | ColorPalette = _normalize_color_input(outline_color)
1240
+
1241
+ @ensure_cv2_image_for_class_method
1242
+ def annotate(
1243
+ self,
1244
+ scene: ImageType,
1245
+ detections: Detections,
1246
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
1247
+ ) -> ImageType:
1248
+ """
1249
+ Annotates the given scene with dots based on the provided detections.
1250
+
1251
+ Args:
1252
+ scene: The image where dots will be drawn.
1253
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
1254
+ or `PIL.Image.Image`.
1255
+ detections: Object detections to annotate.
1256
+ custom_color_lookup: Custom color lookup array.
1257
+ Allows to override the default color mapping strategy.
1258
+
1259
+ Returns:
1260
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
1261
+ or `PIL.Image.Image`)
1262
+
1263
+ Examples:
1264
+ ```pycon
1265
+ >>> import numpy as np
1266
+ >>> import supervision as sv
1267
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
1268
+ >>> detections = sv.Detections(
1269
+ ... xyxy=np.array([[20, 20, 80, 80]]),
1270
+ ... class_id=np.array([0])
1271
+ ... )
1272
+ >>> dot_annotator = sv.DotAnnotator()
1273
+ >>> annotated_frame = dot_annotator.annotate(
1274
+ ... scene=image.copy(),
1275
+ ... detections=detections
1276
+ ... )
1277
+
1278
+ ```
1279
+
1280
+ ![dot-annotator-example](https://media.roboflow.com/
1281
+ supervision-annotator-examples/dot-annotator-example-purple.png)
1282
+ """
1283
+ if not isinstance(scene, np.ndarray):
1284
+ return scene
1285
+ xy = detections.get_anchors_coordinates(anchor=self.position)
1286
+ for detection_idx in range(len(detections)):
1287
+ color = resolve_color(
1288
+ color=self.color,
1289
+ detections=detections,
1290
+ detection_idx=detection_idx,
1291
+ color_lookup=self.color_lookup
1292
+ if custom_color_lookup is None
1293
+ else custom_color_lookup,
1294
+ )
1295
+ center = (int(xy[detection_idx, 0]), int(xy[detection_idx, 1]))
1296
+
1297
+ cv2.circle(scene, center, self.radius, color.as_bgr(), -1)
1298
+ if self.outline_thickness:
1299
+ outline_color = resolve_color(
1300
+ color=self.outline_color,
1301
+ detections=detections,
1302
+ detection_idx=detection_idx,
1303
+ color_lookup=self.color_lookup
1304
+ if custom_color_lookup is None
1305
+ else custom_color_lookup,
1306
+ )
1307
+ cv2.circle(
1308
+ scene,
1309
+ center,
1310
+ self.radius,
1311
+ outline_color.as_bgr(),
1312
+ self.outline_thickness,
1313
+ )
1314
+ return scene
1315
+
1316
+
1317
+ class LabelAnnotator(_BaseLabelAnnotator):
1318
+ """
1319
+ A class for annotating labels on an image using provided detections.
1320
+ """
1321
+
1322
+ def __init__(
1323
+ self,
1324
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
1325
+ color_lookup: ColorLookup = ColorLookup.CLASS,
1326
+ text_color: Color | ColorPalette | str = Color.WHITE,
1327
+ text_scale: float = 0.5,
1328
+ text_thickness: int = 1,
1329
+ text_padding: int = 10,
1330
+ text_position: Position = Position.TOP_LEFT,
1331
+ text_offset: tuple[int, int] = (0, 0),
1332
+ border_radius: int = 0,
1333
+ smart_position: bool = False,
1334
+ max_line_length: int | None = None,
1335
+ ):
1336
+ """
1337
+ Args:
1338
+ color: The color or color palette to use for
1339
+ annotating the text background.
1340
+ color_lookup: Strategy for mapping colors to annotations.
1341
+ Options are `INDEX`, `CLASS`, `TRACK`.
1342
+ text_color: The color or color palette to use
1343
+ for the text.
1344
+ text_scale: Font scale for the text.
1345
+ text_thickness: Thickness of the text characters.
1346
+ text_padding: Padding around the text within its background box.
1347
+ text_position: Position of the text relative to the detection.
1348
+ Possible values are defined in the `Position` enum.
1349
+ text_offset: A tuple of 2D coordinates `(x, y)` to
1350
+ offset the text position from the anchor point, in pixels.
1351
+ border_radius: The radius to apply round edges. If the selected
1352
+ value is higher than the lower dimension, width or height, is clipped.
1353
+ smart_position: Spread out the labels to avoid overlapping.
1354
+ max_line_length: Maximum number of characters per line
1355
+ before wrapping the text. None means no wrapping.
1356
+ """
1357
+ self.text_scale: float = text_scale
1358
+ self.text_thickness: int = text_thickness
1359
+ super().__init__(
1360
+ color=color,
1361
+ color_lookup=color_lookup,
1362
+ text_color=text_color,
1363
+ text_padding=text_padding,
1364
+ text_position=text_position,
1365
+ text_offset=text_offset,
1366
+ border_radius=border_radius,
1367
+ smart_position=smart_position,
1368
+ max_line_length=max_line_length,
1369
+ )
1370
+
1371
+ @ensure_cv2_image_for_class_method
1372
+ def annotate(
1373
+ self,
1374
+ scene: ImageType,
1375
+ detections: Detections,
1376
+ labels: list[str] | None = None,
1377
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
1378
+ ) -> ImageType:
1379
+ """
1380
+ Annotates the given scene with labels based on the provided detections.
1381
+
1382
+ Args:
1383
+ scene: The image where labels will be drawn.
1384
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
1385
+ or `PIL.Image.Image`.
1386
+ detections: Object detections to annotate.
1387
+ labels: Custom labels for each detection.
1388
+ custom_color_lookup: Custom color lookup array.
1389
+ Allows to override the default color mapping strategy.
1390
+
1391
+ Returns:
1392
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
1393
+ or `PIL.Image.Image`)
1394
+
1395
+ Examples:
1396
+ ```pycon
1397
+ >>> import numpy as np
1398
+ >>> import supervision as sv
1399
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
1400
+ >>> detections = sv.Detections(
1401
+ ... xyxy=np.array([[20, 20, 80, 80]]),
1402
+ ... confidence=np.array([0.9]),
1403
+ ... class_id=np.array([0]),
1404
+ ... data={'class_name': np.array(['person'])}
1405
+ ... )
1406
+ >>> labels = [
1407
+ ... f"{class_name} {confidence:.2f}"
1408
+ ... for class_name, confidence
1409
+ ... in zip(detections['class_name'], detections.confidence)
1410
+ ... ]
1411
+ >>> label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
1412
+ >>> annotated_frame = label_annotator.annotate(
1413
+ ... scene=image.copy(),
1414
+ ... detections=detections,
1415
+ ... labels=labels
1416
+ ... )
1417
+
1418
+ ```
1419
+
1420
+ ![label-annotator-example](https://media.roboflow.com/
1421
+ supervision-annotator-examples/label-annotator-example-purple.png)
1422
+ """
1423
+ if not isinstance(scene, np.ndarray):
1424
+ return scene
1425
+ _validate_labels(labels, detections)
1426
+
1427
+ labels = get_labels_text(detections, labels)
1428
+ label_properties: npt.NDArray[np.float32] = self._get_label_properties(
1429
+ detections, labels
1430
+ )
1431
+
1432
+ if self.smart_position:
1433
+ label_properties = self._adjust_labels_in_frame(
1434
+ (scene.shape[1], scene.shape[0]),
1435
+ labels,
1436
+ label_properties,
1437
+ )
1438
+
1439
+ self._draw_labels(
1440
+ scene=scene,
1441
+ labels=labels,
1442
+ label_properties=label_properties,
1443
+ detections=detections,
1444
+ custom_color_lookup=custom_color_lookup,
1445
+ )
1446
+
1447
+ return scene
1448
+
1449
+ def _get_label_properties(
1450
+ self,
1451
+ detections: Detections,
1452
+ labels: list[str],
1453
+ ) -> Any:
1454
+ label_properties = []
1455
+ anchors_coordinates: npt.NDArray[np.int32] = detections.get_anchors_coordinates(
1456
+ anchor=self.text_anchor
1457
+ ).astype(int)
1458
+
1459
+ for label, center_coordinates in zip(labels, anchors_coordinates):
1460
+ center_coordinates = (
1461
+ center_coordinates[0] + self.text_offset[0],
1462
+ center_coordinates[1] + self.text_offset[1],
1463
+ )
1464
+
1465
+ wrapped_lines = wrap_text(label, self.max_line_length)
1466
+ line_heights = []
1467
+ line_widths = []
1468
+
1469
+ for line in wrapped_lines:
1470
+ (text_w, text_h) = cv2.getTextSize(
1471
+ text=line,
1472
+ fontFace=CV2_FONT,
1473
+ fontScale=self.text_scale,
1474
+ thickness=self.text_thickness,
1475
+ )[0]
1476
+ line_heights.append(text_h)
1477
+ line_widths.append(text_w)
1478
+
1479
+ # Get the maximum width and total height
1480
+ max_width = max(line_widths) if line_widths else 0
1481
+ total_height = (
1482
+ sum(line_heights) + (len(line_heights) - 1) * self.text_padding
1483
+ )
1484
+
1485
+ # Add padding around all sides
1486
+ width_padded = max_width + 2 * self.text_padding
1487
+ height_padded = total_height + 2 * self.text_padding
1488
+
1489
+ text_background_xyxy = resolve_text_background_xyxy(
1490
+ center_coordinates=center_coordinates,
1491
+ text_wh=(width_padded, height_padded),
1492
+ position=self.text_anchor,
1493
+ )
1494
+
1495
+ label_properties.append(
1496
+ [
1497
+ *text_background_xyxy,
1498
+ total_height,
1499
+ ]
1500
+ )
1501
+ return np.array(label_properties, dtype=np.float32).reshape(-1, 5)
1502
+
1503
+ def _draw_labels(
1504
+ self,
1505
+ scene: npt.NDArray[np.uint8],
1506
+ labels: list[str],
1507
+ label_properties: npt.NDArray[np.float32],
1508
+ detections: Detections,
1509
+ custom_color_lookup: npt.NDArray[np.int_] | None,
1510
+ ) -> None:
1511
+ assert len(labels) == len(label_properties) == len(detections), (
1512
+ f"Number of label properties ({len(label_properties)}), "
1513
+ f"labels ({len(labels)}) and detections ({len(detections)}) "
1514
+ "do not match."
1515
+ )
1516
+
1517
+ color_lookup = (
1518
+ custom_color_lookup
1519
+ if custom_color_lookup is not None
1520
+ else self.color_lookup
1521
+ )
1522
+
1523
+ for idx, label_property in enumerate(label_properties):
1524
+ background_color = resolve_color(
1525
+ color=self.color,
1526
+ detections=detections,
1527
+ detection_idx=idx,
1528
+ color_lookup=color_lookup,
1529
+ )
1530
+ text_color = resolve_color(
1531
+ color=self.text_color,
1532
+ detections=detections,
1533
+ detection_idx=idx,
1534
+ color_lookup=color_lookup,
1535
+ )
1536
+
1537
+ box_xyxy = label_property[:4].astype(int)
1538
+
1539
+ self.draw_rounded_rectangle(
1540
+ scene=scene,
1541
+ xyxy=box_xyxy,
1542
+ color=background_color.as_bgr(),
1543
+ border_radius=self.border_radius,
1544
+ )
1545
+
1546
+ # Handle multiline text
1547
+ wrapped_lines = wrap_text(labels[idx], self.max_line_length)
1548
+ current_y = box_xyxy[1] + self.text_padding # Start y position
1549
+
1550
+ for line in wrapped_lines:
1551
+ if not line:
1552
+ # Use a character with ascenders and descenders as height reference
1553
+ (_, text_h) = cv2.getTextSize(
1554
+ text="Tg",
1555
+ fontFace=CV2_FONT,
1556
+ fontScale=self.text_scale,
1557
+ thickness=self.text_thickness,
1558
+ )[0]
1559
+ current_y += text_h + self.text_padding
1560
+ continue
1561
+
1562
+ (_, text_h) = cv2.getTextSize(
1563
+ text=line,
1564
+ fontFace=CV2_FONT,
1565
+ fontScale=self.text_scale,
1566
+ thickness=self.text_thickness,
1567
+ )[0]
1568
+
1569
+ text_x = box_xyxy[0] + self.text_padding
1570
+ text_y = current_y + text_h # Add height to get to text baseline
1571
+
1572
+ cv2.putText(
1573
+ img=scene,
1574
+ text=line,
1575
+ org=(text_x, text_y),
1576
+ fontFace=CV2_FONT,
1577
+ fontScale=self.text_scale,
1578
+ color=text_color.as_bgr(),
1579
+ thickness=self.text_thickness,
1580
+ lineType=cv2.LINE_AA,
1581
+ )
1582
+
1583
+ current_y += text_h + self.text_padding # Move to next line position
1584
+
1585
+ @staticmethod
1586
+ def draw_rounded_rectangle(
1587
+ scene: npt.NDArray[np.uint8],
1588
+ xyxy: tuple[int, int, int, int],
1589
+ color: tuple[int, int, int],
1590
+ border_radius: int,
1591
+ ) -> npt.NDArray[np.uint8]:
1592
+ """Draw a filled rectangle with optional rounded corners on an image.
1593
+
1594
+ Args:
1595
+ scene: BGR image array to draw on; modified in-place and returned.
1596
+ xyxy: Bounding box as (x1, y1, x2, y2) pixel coordinates.
1597
+ color: Fill color as a BGR tuple (e.g. ``(0, 0, 255)`` for red).
1598
+ border_radius: Corner rounding radius in pixels. Values <= 0
1599
+ (including values clamped to 0 by a degenerate box) draw a
1600
+ plain filled rectangle with square corners.
1601
+
1602
+ Returns:
1603
+ The annotated ``scene`` array.
1604
+
1605
+ Example:
1606
+ ```python
1607
+ import numpy as np
1608
+ import supervision as sv
1609
+
1610
+ scene = np.zeros((200, 200, 3), dtype=np.uint8)
1611
+ scene = sv.LabelAnnotator.draw_rounded_rectangle(
1612
+ scene=scene,
1613
+ xyxy=(10, 10, 100, 50),
1614
+ color=(0, 255, 0),
1615
+ border_radius=0,
1616
+ )
1617
+ ```
1618
+ """
1619
+ x1, y1, x2, y2 = xyxy
1620
+ width = x2 - x1
1621
+ height = y2 - y1
1622
+
1623
+ border_radius = min(border_radius, min(width, height) // 2)
1624
+
1625
+ if border_radius <= 0:
1626
+ # square corners: a single fill rectangle (the common default), rather
1627
+ # than two rectangles plus four zero-radius corner circles
1628
+ cv2.rectangle(
1629
+ img=scene,
1630
+ pt1=(x1, y1),
1631
+ pt2=(x2, y2),
1632
+ color=color,
1633
+ thickness=-1,
1634
+ )
1635
+ return scene
1636
+
1637
+ rectangle_coordinates = [
1638
+ ((x1 + border_radius, y1), (x2 - border_radius, y2)),
1639
+ ((x1, y1 + border_radius), (x2, y2 - border_radius)),
1640
+ ]
1641
+ circle_centers = [
1642
+ (x1 + border_radius, y1 + border_radius),
1643
+ (x2 - border_radius, y1 + border_radius),
1644
+ (x1 + border_radius, y2 - border_radius),
1645
+ (x2 - border_radius, y2 - border_radius),
1646
+ ]
1647
+
1648
+ for coordinates in rectangle_coordinates:
1649
+ cv2.rectangle(
1650
+ img=scene,
1651
+ pt1=coordinates[0],
1652
+ pt2=coordinates[1],
1653
+ color=color,
1654
+ thickness=-1,
1655
+ )
1656
+ for center in circle_centers:
1657
+ cv2.circle(
1658
+ img=scene,
1659
+ center=center,
1660
+ radius=border_radius,
1661
+ color=color,
1662
+ thickness=-1,
1663
+ )
1664
+ return scene
1665
+
1666
+
1667
+ class RichLabelAnnotator(_BaseLabelAnnotator):
1668
+ """
1669
+ A class for annotating labels on an image using provided detections,
1670
+ with support for Unicode characters by using a custom font.
1671
+ """
1672
+
1673
+ def __init__(
1674
+ self,
1675
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
1676
+ color_lookup: ColorLookup = ColorLookup.CLASS,
1677
+ text_color: Color | ColorPalette | str = Color.WHITE,
1678
+ font_path: str | None = None,
1679
+ font_size: int = 10,
1680
+ text_padding: int = 10,
1681
+ text_position: Position = Position.TOP_LEFT,
1682
+ text_offset: tuple[int, int] = (0, 0),
1683
+ border_radius: int = 0,
1684
+ smart_position: bool = False,
1685
+ max_line_length: int | None = None,
1686
+ ):
1687
+ """
1688
+ Args:
1689
+ color: The color or color palette to use for
1690
+ annotating the text background.
1691
+ color_lookup: Strategy for mapping colors to annotations.
1692
+ Options are `INDEX`, `CLASS`, `TRACK`.
1693
+ text_color: The color to use for the text.
1694
+ font_path: Path to the font file (e.g., ".ttf" or ".otf")
1695
+ to use for rendering text. If `None`, the default PIL font will be used.
1696
+ font_size: Font size for the text.
1697
+ text_padding: Padding around the text within its background box.
1698
+ text_position: Position of the text relative to the detection.
1699
+ Possible values are defined in the `Position` enum.
1700
+ text_offset: A tuple of 2D coordinates `(x, y)` to
1701
+ offset the text position from the anchor point, in pixels.
1702
+ border_radius: The radius to apply round edges. If the selected
1703
+ value is higher than the lower dimension, width or height, is clipped.
1704
+ smart_position: Spread out the labels to avoid overlapping.
1705
+ max_line_length: Maximum number of characters per line
1706
+ before wrapping the text. None means no wrapping.
1707
+ """
1708
+ self.font_path = font_path
1709
+ self.font_size = font_size
1710
+ self.font = self._load_font(font_size, font_path)
1711
+ super().__init__(
1712
+ color=color,
1713
+ color_lookup=color_lookup,
1714
+ text_color=text_color,
1715
+ text_padding=text_padding,
1716
+ text_position=text_position,
1717
+ text_offset=text_offset,
1718
+ border_radius=border_radius,
1719
+ smart_position=smart_position,
1720
+ max_line_length=max_line_length,
1721
+ )
1722
+
1723
+ @ensure_pil_image_for_class_method
1724
+ def annotate(
1725
+ self,
1726
+ scene: ImageType,
1727
+ detections: Detections,
1728
+ labels: list[str] | None = None,
1729
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
1730
+ ) -> ImageType:
1731
+ """
1732
+ Annotates the given scene with labels based on the provided
1733
+ detections, with support for Unicode characters.
1734
+
1735
+ Args:
1736
+ scene: The image where labels will be drawn.
1737
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
1738
+ or `PIL.Image.Image`.
1739
+ detections: Object detections to annotate.
1740
+ labels: Custom labels for each detection.
1741
+ custom_color_lookup: Custom color lookup array.
1742
+ Allows to override the default color mapping strategy.
1743
+
1744
+ Returns:
1745
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
1746
+ or `PIL.Image.Image`)
1747
+
1748
+ Examples:
1749
+ ```pycon
1750
+ >>> import numpy as np
1751
+ >>> import supervision as sv
1752
+ >>> from PIL import Image
1753
+ >>> image = Image.fromarray(np.zeros((100, 100, 3), dtype=np.uint8))
1754
+ >>> detections = sv.Detections(
1755
+ ... xyxy=np.array([[20, 20, 80, 80]]),
1756
+ ... confidence=np.array([0.9]),
1757
+ ... class_id=np.array([0]),
1758
+ ... data={'class_name': np.array(['person'])}
1759
+ ... )
1760
+ >>> labels = [
1761
+ ... f"{class_name} {confidence:.2f}"
1762
+ ... for class_name, confidence
1763
+ ... in zip(detections['class_name'], detections.confidence)
1764
+ ... ]
1765
+ >>> rich_label_annotator = sv.RichLabelAnnotator()
1766
+ >>> annotated_frame = rich_label_annotator.annotate(
1767
+ ... scene=image.copy(),
1768
+ ... detections=detections,
1769
+ ... labels=labels
1770
+ ... )
1771
+
1772
+ ```
1773
+ """
1774
+ _validate_labels(labels, detections)
1775
+
1776
+ draw = ImageDraw.Draw(scene)
1777
+ labels = get_labels_text(detections, labels)
1778
+ label_properties: npt.NDArray[np.float32] = self._get_label_properties(
1779
+ draw, detections, labels
1780
+ )
1781
+
1782
+ if self.smart_position:
1783
+ scene_pil = cast(Image.Image, scene)
1784
+ label_properties = self._adjust_labels_in_frame(
1785
+ (scene_pil.width, scene_pil.height),
1786
+ labels,
1787
+ label_properties,
1788
+ )
1789
+
1790
+ self._draw_labels(
1791
+ draw=draw,
1792
+ labels=labels,
1793
+ label_properties=label_properties,
1794
+ detections=detections,
1795
+ custom_color_lookup=custom_color_lookup,
1796
+ )
1797
+
1798
+ return scene
1799
+
1800
+ def _get_label_properties(
1801
+ self, draw: ImageDraw.ImageDraw, detections: Detections, labels: list[str]
1802
+ ) -> Any:
1803
+ label_properties = []
1804
+
1805
+ anchor_coordinates: npt.NDArray[np.int32] = detections.get_anchors_coordinates(
1806
+ anchor=self.text_anchor
1807
+ ).astype(int)
1808
+
1809
+ for label, center_coordinates in zip(labels, anchor_coordinates):
1810
+ center_coordinates = (
1811
+ center_coordinates[0] + self.text_offset[0],
1812
+ center_coordinates[1] + self.text_offset[1],
1813
+ )
1814
+
1815
+ wrapped_lines = wrap_text(label, self.max_line_length)
1816
+
1817
+ # Calculate the total text height and maximum width
1818
+ max_width = 0.0
1819
+ total_height = 0.0
1820
+
1821
+ for line in wrapped_lines:
1822
+ left, top, right, bottom = draw.textbbox((0, 0), line, font=self.font)
1823
+ line_width = right - left
1824
+ line_height = bottom - top
1825
+
1826
+ max_width = max(max_width, line_width)
1827
+ total_height += line_height
1828
+
1829
+ # Add inter-line spacing
1830
+ if len(wrapped_lines) > 1:
1831
+ total_height += (len(wrapped_lines) - 1) * self.text_padding
1832
+
1833
+ width_padded = int(max_width + 2 * self.text_padding)
1834
+ height_padded = int(total_height + 2 * self.text_padding)
1835
+
1836
+ text_background_xyxy = resolve_text_background_xyxy(
1837
+ center_coordinates=center_coordinates,
1838
+ text_wh=(width_padded, height_padded),
1839
+ position=self.text_anchor,
1840
+ )
1841
+
1842
+ # Get the text origin offsets
1843
+ text_left, text_top, _, _ = draw.textbbox((0, 0), "Tg", font=self.font)
1844
+
1845
+ label_properties.append([*text_background_xyxy, text_left, text_top])
1846
+
1847
+ result: npt.NDArray[np.float32] = np.array(
1848
+ label_properties, dtype=np.float32
1849
+ ).reshape(-1, 6)
1850
+ return result
1851
+
1852
+ def _draw_labels(
1853
+ self,
1854
+ draw: ImageDraw.ImageDraw,
1855
+ labels: list[str],
1856
+ label_properties: npt.NDArray[np.float32],
1857
+ detections: Detections,
1858
+ custom_color_lookup: npt.NDArray[np.int_] | None,
1859
+ ) -> None:
1860
+ assert len(labels) == len(label_properties) == len(detections), (
1861
+ f"Number of label properties ({len(label_properties)}), "
1862
+ f"labels ({len(labels)}) and detections ({len(detections)}) "
1863
+ "do not match."
1864
+ )
1865
+ color_lookup = (
1866
+ custom_color_lookup
1867
+ if custom_color_lookup is not None
1868
+ else self.color_lookup
1869
+ )
1870
+
1871
+ for idx, label_property in enumerate(label_properties):
1872
+ background_color = resolve_color(
1873
+ color=self.color,
1874
+ detections=detections,
1875
+ detection_idx=idx,
1876
+ color_lookup=color_lookup,
1877
+ )
1878
+ text_color = resolve_color(
1879
+ color=self.text_color,
1880
+ detections=detections,
1881
+ detection_idx=idx,
1882
+ color_lookup=color_lookup,
1883
+ )
1884
+
1885
+ box_xyxy = label_property[:4].astype(int)
1886
+ text_left = label_property[4]
1887
+ text_top = label_property[5]
1888
+
1889
+ # Draw the rounded rectangle background
1890
+ draw.rounded_rectangle(
1891
+ tuple(box_xyxy),
1892
+ radius=self.border_radius,
1893
+ fill=background_color.as_rgb(),
1894
+ outline=None,
1895
+ )
1896
+
1897
+ # Draw each line of text
1898
+ wrapped_lines = wrap_text(labels[idx], self.max_line_length)
1899
+ x_position = box_xyxy[0] + self.text_padding - text_left
1900
+ y_position = box_xyxy[1] + self.text_padding - text_top
1901
+
1902
+ for line in wrapped_lines:
1903
+ draw.text(
1904
+ xy=(x_position, y_position),
1905
+ text=line,
1906
+ font=self.font,
1907
+ fill=text_color.as_rgb(),
1908
+ )
1909
+
1910
+ # Move to the next line position
1911
+ _left, top, _right, bottom = draw.textbbox((0, 0), line, font=self.font)
1912
+ line_height = bottom - top
1913
+ y_position += line_height + self.text_padding
1914
+
1915
+ @staticmethod
1916
+ def _load_font(
1917
+ font_size: int, font_path: str | None
1918
+ ) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
1919
+ def load_default_font(
1920
+ size: int,
1921
+ ) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
1922
+ try:
1923
+ return ImageFont.load_default(size)
1924
+ except TypeError:
1925
+ return ImageFont.load_default()
1926
+
1927
+ if font_path is None:
1928
+ return load_default_font(font_size)
1929
+
1930
+ try:
1931
+ return ImageFont.truetype(font_path, font_size)
1932
+ except OSError:
1933
+ logger.warning(
1934
+ "Font path '%s' not found. Using PIL's default font.", font_path
1935
+ )
1936
+ return load_default_font(font_size)
1937
+
1938
+
1939
+ class IconAnnotator(BaseAnnotator):
1940
+ """
1941
+ A class for drawing an icon on an image, using provided detections.
1942
+ """
1943
+
1944
+ def __init__(
1945
+ self,
1946
+ icon_resolution_wh: tuple[int, int] = (64, 64),
1947
+ icon_position: Position = Position.TOP_CENTER,
1948
+ offset_xy: tuple[int, int] = (0, 0),
1949
+ ):
1950
+ """
1951
+ Args:
1952
+ icon_resolution_wh: The size of drawn icons.
1953
+ All icons will be resized to this resolution, keeping the aspect ratio.
1954
+ icon_position: The position of the icon.
1955
+ offset_xy: The offset to apply to the icon position,
1956
+ in pixels. Can be both positive and negative.
1957
+ """
1958
+ self.icon_resolution_wh = icon_resolution_wh
1959
+ self.position = icon_position
1960
+ self.offset_xy = offset_xy
1961
+
1962
+ @ensure_cv2_image_for_class_method
1963
+ def annotate(
1964
+ self,
1965
+ scene: ImageType,
1966
+ detections: Detections,
1967
+ icon_path: str | list[str] = "",
1968
+ ) -> ImageType:
1969
+ """
1970
+ Annotates the given scene with given icons.
1971
+
1972
+ Args:
1973
+ scene: The image where labels will be drawn.
1974
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
1975
+ or `PIL.Image.Image`.
1976
+ detections: Object detections to annotate.
1977
+ icon_path: The path to the PNG image to use as an
1978
+ icon. Must be a single path or a list of paths, one for each detection.
1979
+ Pass an empty string `""` to draw nothing.
1980
+
1981
+ Returns:
1982
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
1983
+ or `PIL.Image.Image`)
1984
+
1985
+ Example:
1986
+ ```python
1987
+ import supervision as sv
1988
+
1989
+ image = ...
1990
+ detections = sv.Detections(...)
1991
+
1992
+ available_icons = ["roboflow.png", "lenny.png"]
1993
+ icon_paths = [np.random.choice(available_icons) for _ in detections]
1994
+
1995
+ icon_annotator = sv.IconAnnotator()
1996
+ annotated_frame = icon_annotator.annotate(
1997
+ scene=image.copy(),
1998
+ detections=detections,
1999
+ icon_path=icon_paths
2000
+ )
2001
+ ```
2002
+
2003
+ ![icon-annotator-example](https://media.roboflow.com/
2004
+ supervision-annotator-examples/icon-annotator-example.png)
2005
+ """
2006
+ if not isinstance(scene, np.ndarray):
2007
+ return scene
2008
+ if isinstance(icon_path, list) and len(icon_path) != len(detections):
2009
+ raise ValueError(
2010
+ f"The number of icon paths provided ({len(icon_path)}) does not match "
2011
+ f"the number of detections ({len(detections)}). Either provide a single"
2012
+ f" icon path or one for each detection."
2013
+ )
2014
+
2015
+ xy: npt.NDArray[np.int32] = detections.get_anchors_coordinates(
2016
+ anchor=self.position
2017
+ ).astype(int)
2018
+
2019
+ for detection_idx in range(len(detections)):
2020
+ current_path = (
2021
+ icon_path if isinstance(icon_path, str) else icon_path[detection_idx]
2022
+ )
2023
+ if current_path == "":
2024
+ continue
2025
+ icon = self._load_icon(current_path)
2026
+ icon_h, icon_w = icon.shape[:2]
2027
+
2028
+ x = int(xy[detection_idx, 0] - icon_w / 2 + self.offset_xy[0])
2029
+ y = int(xy[detection_idx, 1] - icon_h / 2 + self.offset_xy[1])
2030
+
2031
+ scene[:] = _overlay_image(scene, icon, (x, y))
2032
+ return scene
2033
+
2034
+ def _load_icon(self, icon_path: str) -> npt.NDArray[np.uint8]:
2035
+ """Load an icon through the module-level cache shared by annotators."""
2036
+ return _load_icon_from_path(
2037
+ icon_path=icon_path, icon_resolution_wh=self.icon_resolution_wh
2038
+ )
2039
+
2040
+
2041
+ class BlurAnnotator(BaseAnnotator):
2042
+ """
2043
+ A class for blurring regions in an image using provided detections.
2044
+ """
2045
+
2046
+ def __init__(self, kernel_size: int | None = None):
2047
+ """
2048
+ Args:
2049
+ kernel_size: The size of the average pooling kernel used for blurring.
2050
+ If not set, a dynamic size is computed as one-third of the shorter
2051
+ bounding-box dimension. Must be >= 1 when provided.
2052
+ """
2053
+ if kernel_size is not None and kernel_size < 1:
2054
+ raise ValueError(f"kernel_size must be >= 1, got {kernel_size}.")
2055
+ self.kernel_size: int | None = kernel_size
2056
+
2057
+ @ensure_cv2_image_for_class_method
2058
+ def annotate(
2059
+ self,
2060
+ scene: ImageType,
2061
+ detections: Detections,
2062
+ ) -> ImageType:
2063
+ """
2064
+ Annotates the given scene by blurring regions based on the provided detections.
2065
+
2066
+ Args:
2067
+ scene: The image where blurring will be applied.
2068
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
2069
+ or `PIL.Image.Image`.
2070
+ detections: Object detections to annotate.
2071
+
2072
+ Returns:
2073
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
2074
+ or `PIL.Image.Image`)
2075
+
2076
+ Examples:
2077
+ ```pycon
2078
+ >>> import numpy as np
2079
+ >>> import supervision as sv
2080
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
2081
+ >>> detections = sv.Detections(
2082
+ ... xyxy=np.array([[20, 20, 80, 80]]),
2083
+ ... class_id=np.array([0])
2084
+ ... )
2085
+ >>> blur_annotator = sv.BlurAnnotator()
2086
+ >>> annotated_frame = blur_annotator.annotate(
2087
+ ... scene=image.copy(),
2088
+ ... detections=detections
2089
+ ... )
2090
+
2091
+ ```
2092
+
2093
+ ![blur-annotator-example](https://media.roboflow.com/
2094
+ supervision-annotator-examples/blur-annotator-example-purple.png)
2095
+ """
2096
+ if not isinstance(scene, np.ndarray):
2097
+ return scene
2098
+ image_height, image_width = scene.shape[:2]
2099
+ clipped_xyxy: npt.NDArray[np.int32] = clip_boxes(
2100
+ xyxy=to_numpy(detections.xyxy),
2101
+ resolution_wh=(image_width, image_height),
2102
+ ).astype(int)
2103
+
2104
+ for x1, y1, x2, y2 in clipped_xyxy:
2105
+ if x2 <= x1 or y2 <= y1:
2106
+ continue
2107
+ roi = scene[y1:y2, x1:x2]
2108
+ kernel_size = (
2109
+ self.kernel_size
2110
+ if self.kernel_size is not None
2111
+ else calculate_dynamic_kernel_size(x1, y1, x2, y2)
2112
+ )
2113
+ roi = cast(npt.NDArray[np.uint8], cv2.blur(roi, (kernel_size, kernel_size)))
2114
+ scene[y1:y2, x1:x2] = roi
2115
+
2116
+ return scene
2117
+
2118
+
2119
+ class TraceAnnotator(BaseAnnotator):
2120
+ """
2121
+ A class for drawing trace paths on an image based on detection coordinates.
2122
+
2123
+ !!! warning
2124
+
2125
+ This annotator uses the `sv.Detections.tracker_id`. Read
2126
+ [here](/latest/trackers/) to learn how to plug
2127
+ tracking into your inference pipeline.
2128
+ """
2129
+
2130
+ def __init__(
2131
+ self,
2132
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
2133
+ position: Position = Position.CENTER,
2134
+ trace_length: int = 30,
2135
+ thickness: int = 2,
2136
+ smooth: bool = False,
2137
+ color_lookup: ColorLookup = ColorLookup.CLASS,
2138
+ ):
2139
+ """
2140
+ Args:
2141
+ color: The color to draw the trace, can be
2142
+ a single color or a color palette.
2143
+ position: The position of the trace.
2144
+ Defaults to `CENTER`.
2145
+ trace_length: The maximum length of the trace in terms of historical
2146
+ points. Defaults to `30`.
2147
+ thickness: The thickness of the trace lines. Defaults to `2`.
2148
+ smooth: If `True`, applies spline smoothing to trace lines using
2149
+ consecutive unique anchor points. Falls back to a raw polyline
2150
+ when fewer than 4 unique points are available (e.g. when a
2151
+ tracker is stationary).
2152
+ color_lookup: Strategy for mapping colors to annotations.
2153
+ Options are `INDEX`, `CLASS`, `TRACK`.
2154
+ """
2155
+ self.color: Color | ColorPalette = _normalize_color_input(color)
2156
+ self.trace = Trace(max_size=trace_length, anchor=position)
2157
+ self.thickness = thickness
2158
+ self.smooth = smooth
2159
+ self.color_lookup: ColorLookup = color_lookup
2160
+
2161
+ def reset(self) -> None:
2162
+ """
2163
+ Clears the accumulated trace history so the annotator can be reused
2164
+ across independent streams without carrying over points from a
2165
+ previous stream.
2166
+
2167
+ Examples:
2168
+ ```pycon
2169
+ >>> import numpy as np
2170
+ >>> import supervision as sv
2171
+ >>> image = np.zeros((20, 20, 3), dtype=np.uint8)
2172
+ >>> detections = sv.Detections(
2173
+ ... xyxy=np.array([[1, 1, 10, 10]]),
2174
+ ... class_id=np.array([0]),
2175
+ ... tracker_id=np.array([1])
2176
+ ... )
2177
+ >>> trace_annotator = sv.TraceAnnotator()
2178
+ >>> _ = trace_annotator.annotate(scene=image.copy(), detections=detections)
2179
+ >>> trace_annotator.trace.xy.shape
2180
+ (1, 2)
2181
+ >>> trace_annotator.reset()
2182
+ >>> trace_annotator.trace.current_frame_id
2183
+ 0
2184
+ >>> trace_annotator.trace.xy.shape
2185
+ (0, 2)
2186
+
2187
+ ```
2188
+ """
2189
+ self.trace.reset()
2190
+
2191
+ @ensure_cv2_image_for_class_method
2192
+ def annotate(
2193
+ self,
2194
+ scene: ImageType,
2195
+ detections: Detections,
2196
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
2197
+ ) -> ImageType:
2198
+ """
2199
+ Draws trace paths on the frame based on the detection coordinates provided.
2200
+
2201
+ Args:
2202
+ scene: The image on which the traces will be drawn.
2203
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
2204
+ or `PIL.Image.Image`.
2205
+ detections: The detections which include coordinates for
2206
+ which the traces will be drawn.
2207
+ custom_color_lookup: Custom color lookup array.
2208
+ Allows to override the default color mapping strategy.
2209
+
2210
+ Returns:
2211
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
2212
+ or `PIL.Image.Image`)
2213
+
2214
+ Example:
2215
+ ```python
2216
+ import supervision as sv
2217
+ from ultralytics import YOLO
2218
+
2219
+ model = YOLO('yolov8x.pt')
2220
+ trace_annotator = sv.TraceAnnotator()
2221
+
2222
+ video_info = sv.VideoInfo.from_video_path(video_path='...')
2223
+ frames_generator = sv.get_video_frames_generator(source_path='...')
2224
+ tracker = sv.ByteTrack()
2225
+
2226
+ with sv.VideoSink(target_path='...', video_info=video_info) as sink:
2227
+ for frame in frames_generator:
2228
+ result = model(frame)[0]
2229
+ detections = sv.Detections.from_ultralytics(result)
2230
+ detections = tracker.update_with_detections(detections)
2231
+ annotated_frame = trace_annotator.annotate(
2232
+ scene=frame.copy(),
2233
+ detections=detections)
2234
+ sink.write_frame(frame=annotated_frame)
2235
+ ```
2236
+
2237
+ ![trace-annotator-example](https://media.roboflow.com/
2238
+ supervision-annotator-examples/trace-annotator-example-purple.png)
2239
+ """
2240
+ if not isinstance(scene, np.ndarray):
2241
+ return scene
2242
+ if detections.tracker_id is None:
2243
+ raise ValueError(
2244
+ "The `tracker_id` field is missing in the provided detections."
2245
+ " See more: https://supervision.roboflow.com/latest/how_to/track_objects"
2246
+ )
2247
+ filtered_detections: Detections = detections[
2248
+ detections.tracker_id != PENDING_TRACK_ID
2249
+ ] # type: ignore
2250
+
2251
+ self.trace.put(filtered_detections)
2252
+ for detection_idx in range(len(filtered_detections)):
2253
+ tracker_id_val = filtered_detections.tracker_id[detection_idx] # type: ignore
2254
+ if tracker_id_val is None:
2255
+ continue
2256
+ tracker_id = int(tracker_id_val)
2257
+ color = resolve_color(
2258
+ color=self.color,
2259
+ detections=filtered_detections,
2260
+ detection_idx=detection_idx,
2261
+ color_lookup=self.color_lookup
2262
+ if custom_color_lookup is None
2263
+ else custom_color_lookup,
2264
+ )
2265
+ xy = self.trace.get(tracker_id=tracker_id)
2266
+ spline_points: npt.NDArray[np.int32] = xy.astype(np.int32)
2267
+
2268
+ if self.smooth:
2269
+ unique_xy = xy[
2270
+ np.concatenate(([True], np.any(np.diff(xy, axis=0) != 0, axis=1)))
2271
+ ]
2272
+ if len(unique_xy) > 3:
2273
+ try:
2274
+ x, y = unique_xy[:, 0], unique_xy[:, 1]
2275
+ tck, _u = splprep([x, y], s=20)
2276
+ x_new, y_new = splev(
2277
+ np.linspace(0, 1, 100),
2278
+ cast(
2279
+ tuple[
2280
+ npt.NDArray[np.float64],
2281
+ npt.NDArray[np.float64],
2282
+ int,
2283
+ ],
2284
+ tck,
2285
+ ),
2286
+ )
2287
+ spline_points = np.stack((x_new, y_new), axis=1).astype(
2288
+ np.int32
2289
+ )
2290
+ except ValueError:
2291
+ spline_points = unique_xy.astype(np.int32)
2292
+ else:
2293
+ spline_points = unique_xy.astype(np.int32)
2294
+
2295
+ if len(xy) > 1:
2296
+ cv2.polylines(
2297
+ scene,
2298
+ [spline_points],
2299
+ False,
2300
+ color=color.as_bgr(),
2301
+ thickness=self.thickness,
2302
+ )
2303
+ return scene
2304
+
2305
+
2306
+ class HeatMapAnnotator(BaseAnnotator):
2307
+ """
2308
+ A class for drawing heatmaps on an image based on provided detections.
2309
+ Heat accumulates over time and is drawn as a semi-transparent overlay
2310
+ of blurred circles.
2311
+ """
2312
+
2313
+ def __init__(
2314
+ self,
2315
+ position: Position = Position.BOTTOM_CENTER,
2316
+ opacity: float = 0.2,
2317
+ radius: int = 40,
2318
+ kernel_size: int | None = 25,
2319
+ top_hue: int = 0,
2320
+ low_hue: int = 125,
2321
+ ):
2322
+ """
2323
+ Args:
2324
+ position: The position of the heatmap. Defaults to
2325
+ `BOTTOM_CENTER`.
2326
+ opacity: Opacity of the overlay mask. Must be between `0` and `1`.
2327
+ radius: Radius of the heat circle.
2328
+ kernel_size: Kernel size for blurring the heatmap. Pass `None`
2329
+ to disable blurring entirely.
2330
+ top_hue: Hue at the top of the heatmap. Defaults to 0 (red).
2331
+ low_hue: Hue at the bottom of the heatmap. Defaults to 125 (blue).
2332
+ """
2333
+ self.position = position
2334
+ self.opacity = opacity
2335
+ self.radius = radius
2336
+ self.kernel_size = kernel_size
2337
+ self.top_hue = top_hue
2338
+ self.low_hue = low_hue
2339
+ self.heat_mask: npt.NDArray[np.float32] | None = None
2340
+
2341
+ def reset(self) -> None:
2342
+ """
2343
+ Clears the accumulated heat so the annotator can be reused across
2344
+ independent streams. `annotate` already reinitializes the heat mask
2345
+ when the scene resolution changes; call this to discard heat from a
2346
+ previous stream that shares the same resolution.
2347
+
2348
+ Examples:
2349
+ ```pycon
2350
+ >>> import numpy as np
2351
+ >>> import supervision as sv
2352
+ >>> image = np.zeros((40, 40, 3), dtype=np.uint8)
2353
+ >>> detections = sv.Detections(xyxy=np.array([[10, 10, 20, 20]]))
2354
+ >>> heat_map_annotator = sv.HeatMapAnnotator()
2355
+ >>> _ = heat_map_annotator.annotate(
2356
+ ... scene=image.copy(),
2357
+ ... detections=detections
2358
+ ... )
2359
+ >>> bool(heat_map_annotator.heat_mask.sum() > 0)
2360
+ True
2361
+ >>> heat_map_annotator.reset()
2362
+ >>> heat_map_annotator.heat_mask is None
2363
+ True
2364
+
2365
+ ```
2366
+ """
2367
+ self.heat_mask = None
2368
+
2369
+ @ensure_cv2_image_for_class_method
2370
+ def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
2371
+ """
2372
+ Annotates the scene with a heatmap based on the provided detections.
2373
+
2374
+ Args:
2375
+ scene: The image where the heatmap will be drawn.
2376
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
2377
+ or `PIL.Image.Image`.
2378
+ detections: Object detections to annotate.
2379
+
2380
+ Returns:
2381
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
2382
+ or `PIL.Image.Image`)
2383
+
2384
+ Note:
2385
+ When `detections` is empty or no heat has accumulated yet, the
2386
+ scene is returned unchanged without raising a ``RuntimeWarning``.
2387
+
2388
+ Example:
2389
+ ```python
2390
+ import supervision as sv
2391
+ from ultralytics import YOLO
2392
+
2393
+ model = YOLO('yolov8x.pt')
2394
+
2395
+ heat_map_annotator = sv.HeatMapAnnotator()
2396
+
2397
+ video_info = sv.VideoInfo.from_video_path(video_path='...')
2398
+ frames_generator = sv.get_video_frames_generator(source_path='...')
2399
+
2400
+ with sv.VideoSink(target_path='...', video_info=video_info) as sink:
2401
+ for frame in frames_generator:
2402
+ result = model(frame)[0]
2403
+ detections = sv.Detections.from_ultralytics(result)
2404
+ annotated_frame = heat_map_annotator.annotate(
2405
+ scene=frame.copy(),
2406
+ detections=detections)
2407
+ sink.write_frame(frame=annotated_frame)
2408
+ ```
2409
+
2410
+ ![heatmap-annotator-example](https://media.roboflow.com/
2411
+ supervision-annotator-examples/heat-map-annotator-example-purple.png)
2412
+ """
2413
+ if not isinstance(scene, np.ndarray):
2414
+ return scene
2415
+ if self.heat_mask is None or self.heat_mask.shape != scene.shape[:2]:
2416
+ self.heat_mask = np.zeros(scene.shape[:2], dtype=np.float32)
2417
+
2418
+ mask: npt.NDArray[np.float32] = np.zeros(scene.shape[:2], dtype=np.float32)
2419
+ anchors = to_numpy(
2420
+ detections.get_anchors_coordinates(self.position)
2421
+ ).astype(int)
2422
+ for x, y in anchors:
2423
+ cv2.circle(
2424
+ img=mask,
2425
+ center=(x, y),
2426
+ radius=self.radius,
2427
+ color=(1,),
2428
+ thickness=-1, # fill
2429
+ )
2430
+ self.heat_mask = mask + self.heat_mask
2431
+ heat_mask = self.heat_mask
2432
+ heat_values = heat_mask.copy()
2433
+ max_val = heat_values.max()
2434
+ if max_val > 0:
2435
+ heat_values = self.low_hue - heat_values / max_val * (
2436
+ self.low_hue - self.top_hue
2437
+ )
2438
+ heat_hue = heat_values.astype(np.uint8)
2439
+ if self.kernel_size is not None:
2440
+ heat_hue = cast(
2441
+ npt.NDArray[np.uint8],
2442
+ cv2.blur(heat_hue, (self.kernel_size, self.kernel_size)),
2443
+ )
2444
+ hsv = np.full(scene.shape, 255, dtype=np.uint8)
2445
+ hsv[..., 0] = heat_hue
2446
+ heat_bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
2447
+ mask2d = heat_mask > 0
2448
+ blended = cv2.addWeighted(heat_bgr, self.opacity, scene, 1 - self.opacity, 0)
2449
+ scene[mask2d] = blended[mask2d]
2450
+ return scene
2451
+
2452
+
2453
+ class PixelateAnnotator(BaseAnnotator):
2454
+ """
2455
+ A class for pixelating regions in an image using provided detections.
2456
+ """
2457
+
2458
+ def __init__(self, pixel_size: int | None = None):
2459
+ """
2460
+ Args:
2461
+ pixel_size: The size of the pixelation. If not set, a dynamic size is
2462
+ computed as one-half of the shorter bounding-box dimension. When set
2463
+ and the detection area is smaller than `pixel_size`, the region is
2464
+ filled with its average colour instead to avoid an OpenCV crash.
2465
+ Must be >= 1 when provided.
2466
+ """
2467
+ if pixel_size is not None and pixel_size < 1:
2468
+ raise ValueError(f"pixel_size must be >= 1, got {pixel_size}.")
2469
+ self.pixel_size: int | None = pixel_size
2470
+
2471
+ @ensure_cv2_image_for_class_method
2472
+ def annotate(
2473
+ self,
2474
+ scene: ImageType,
2475
+ detections: Detections,
2476
+ ) -> ImageType:
2477
+ """
2478
+ Annotates the given scene by pixelating regions based on the provided
2479
+ detections.
2480
+
2481
+ Args:
2482
+ scene: The image where pixelating will be applied.
2483
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
2484
+ or `PIL.Image.Image`.
2485
+ detections: Object detections to annotate.
2486
+
2487
+ Returns:
2488
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
2489
+ or `PIL.Image.Image`)
2490
+
2491
+ Example:
2492
+ ```python
2493
+ import supervision as sv
2494
+
2495
+ image = ...
2496
+ detections = sv.Detections(...)
2497
+
2498
+ pixelate_annotator = sv.PixelateAnnotator()
2499
+ annotated_frame = pixelate_annotator.annotate(
2500
+ scene=image.copy(),
2501
+ detections=detections
2502
+ )
2503
+ ```
2504
+
2505
+ ![pixelate-annotator-example](https://media.roboflow.com/
2506
+ supervision-annotator-examples/pixelate-annotator-example-10.png)
2507
+ """
2508
+ if not isinstance(scene, np.ndarray):
2509
+ return scene
2510
+ image_height, image_width = scene.shape[:2]
2511
+ clipped_xyxy: npt.NDArray[np.int32] = clip_boxes(
2512
+ xyxy=to_numpy(detections.xyxy),
2513
+ resolution_wh=(image_width, image_height),
2514
+ ).astype(int)
2515
+
2516
+ for x1, y1, x2, y2 in clipped_xyxy:
2517
+ if x2 <= x1 or y2 <= y1:
2518
+ continue
2519
+ roi = scene[y1:y2, x1:x2]
2520
+
2521
+ pixel_size = (
2522
+ self.pixel_size
2523
+ if self.pixel_size is not None
2524
+ else calculate_dynamic_pixel_size(x1, y1, x2, y2)
2525
+ )
2526
+ if min(y2 - y1, x2 - x1) < pixel_size:
2527
+ if roi.ndim == 2 or (roi.ndim == 3 and roi.shape[2] == 1):
2528
+ scene[y1:y2, x1:x2] = cv2.mean(roi)[0]
2529
+ else:
2530
+ num_channels = scene.shape[2]
2531
+ scene[y1:y2, x1:x2] = cv2.mean(roi)[:num_channels]
2532
+ continue
2533
+
2534
+ scaled_up_roi = cv2.resize(
2535
+ src=roi, dsize=None, fx=1 / pixel_size, fy=1 / pixel_size
2536
+ )
2537
+ scaled_down_roi = cv2.resize(
2538
+ src=scaled_up_roi,
2539
+ dsize=(roi.shape[1], roi.shape[0]),
2540
+ interpolation=cv2.INTER_NEAREST,
2541
+ )
2542
+
2543
+ scene[y1:y2, x1:x2] = scaled_down_roi
2544
+
2545
+ return scene
2546
+
2547
+
2548
+ class TriangleAnnotator(BaseAnnotator):
2549
+ """
2550
+ A class for drawing triangle markers on an image at specific coordinates based on
2551
+ provided detections.
2552
+ """
2553
+
2554
+ def __init__(
2555
+ self,
2556
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
2557
+ base: int = 10,
2558
+ height: int = 10,
2559
+ position: Position = Position.TOP_CENTER,
2560
+ color_lookup: ColorLookup = ColorLookup.CLASS,
2561
+ outline_thickness: int = 0,
2562
+ outline_color: Color | ColorPalette | str = Color.BLACK,
2563
+ ):
2564
+ """
2565
+ Args:
2566
+ color: The color or color palette to use for
2567
+ annotating detections.
2568
+ base: The base width of the triangle.
2569
+ height: The height of the triangle.
2570
+ position: The anchor position for placing the triangle.
2571
+ color_lookup: Strategy for mapping colors to annotations.
2572
+ Options are `INDEX`, `CLASS`, `TRACK`.
2573
+ outline_thickness: Thickness of the outline of the triangle.
2574
+ outline_color: The color or color palette to
2575
+ use for outline. It is activated by setting outline_thickness to a value
2576
+ greater than 0.
2577
+ """
2578
+ self.color: Color | ColorPalette = _normalize_color_input(color)
2579
+ self.base: int = base
2580
+ self.height: int = height
2581
+ self.position: Position = position
2582
+ self.color_lookup: ColorLookup = color_lookup
2583
+ self.outline_thickness: int = outline_thickness
2584
+ self.outline_color: Color | ColorPalette = _normalize_color_input(outline_color)
2585
+
2586
+ @ensure_cv2_image_for_class_method
2587
+ def annotate(
2588
+ self,
2589
+ scene: ImageType,
2590
+ detections: Detections,
2591
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
2592
+ ) -> ImageType:
2593
+ """
2594
+ Annotates the given scene with triangles based on the provided detections.
2595
+
2596
+ Args:
2597
+ scene: The image where triangles will be drawn.
2598
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
2599
+ or `PIL.Image.Image`.
2600
+ detections: Object detections to annotate.
2601
+ custom_color_lookup: Custom color lookup array.
2602
+ Allows to override the default color mapping strategy.
2603
+
2604
+ Returns:
2605
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
2606
+ or `PIL.Image.Image`)
2607
+
2608
+ Examples:
2609
+ ```pycon
2610
+ >>> import numpy as np
2611
+ >>> import supervision as sv
2612
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
2613
+ >>> detections = sv.Detections(
2614
+ ... xyxy=np.array([[20, 20, 80, 80]]),
2615
+ ... class_id=np.array([0])
2616
+ ... )
2617
+ >>> triangle_annotator = sv.TriangleAnnotator()
2618
+ >>> annotated_frame = triangle_annotator.annotate(
2619
+ ... scene=image.copy(),
2620
+ ... detections=detections
2621
+ ... )
2622
+
2623
+ ```
2624
+
2625
+ ![triangle-annotator-example](https://media.roboflow.com/
2626
+ supervision-annotator-examples/triangle-annotator-example.png)
2627
+ """
2628
+ if not isinstance(scene, np.ndarray):
2629
+ return scene
2630
+ xy = to_numpy(
2631
+ detections.get_anchors_coordinates(anchor=self.position)
2632
+ ).astype(int)
2633
+ for detection_idx in range(len(detections)):
2634
+ color = resolve_color(
2635
+ color=self.color,
2636
+ detections=detections,
2637
+ detection_idx=detection_idx,
2638
+ color_lookup=self.color_lookup
2639
+ if custom_color_lookup is None
2640
+ else custom_color_lookup,
2641
+ )
2642
+ tip_x, tip_y = int(xy[detection_idx, 0]), int(xy[detection_idx, 1])
2643
+ vertices = np.array(
2644
+ [
2645
+ [tip_x - self.base // 2, tip_y - self.height],
2646
+ [tip_x + self.base // 2, tip_y - self.height],
2647
+ [tip_x, tip_y],
2648
+ ],
2649
+ np.int32,
2650
+ )
2651
+
2652
+ cv2.fillPoly(scene, [vertices], color.as_bgr())
2653
+ if self.outline_thickness:
2654
+ outline_color = resolve_color(
2655
+ color=self.outline_color,
2656
+ detections=detections,
2657
+ detection_idx=detection_idx,
2658
+ color_lookup=self.color_lookup
2659
+ if custom_color_lookup is None
2660
+ else custom_color_lookup,
2661
+ )
2662
+ cv2.polylines(
2663
+ scene,
2664
+ [vertices],
2665
+ True,
2666
+ outline_color.as_bgr(),
2667
+ thickness=self.outline_thickness,
2668
+ )
2669
+ return scene
2670
+
2671
+
2672
+ class RoundBoxAnnotator(BaseAnnotator):
2673
+ """
2674
+ A class for drawing bounding boxes with round edges on an image
2675
+ using provided detections.
2676
+ """
2677
+
2678
+ def __init__(
2679
+ self,
2680
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
2681
+ thickness: int = 2,
2682
+ color_lookup: ColorLookup = ColorLookup.CLASS,
2683
+ roundness: float = 0.6,
2684
+ ):
2685
+ """
2686
+ Args:
2687
+ color: The color or color palette to use for
2688
+ annotating detections.
2689
+ thickness: Thickness of the bounding box lines.
2690
+ color_lookup: Strategy for mapping colors to annotations.
2691
+ Options are `INDEX`, `CLASS`, `TRACK`.
2692
+ roundness: Percent of roundness for edges of bounding box.
2693
+ Value must be float 0 < roundness <= 1.0
2694
+ By default roundness percent is calculated based on smaller side
2695
+ length (width or height).
2696
+ """
2697
+ self.color: Color | ColorPalette = _normalize_color_input(color)
2698
+ self.thickness: int = thickness
2699
+ self.color_lookup: ColorLookup = color_lookup
2700
+ if not 0 < roundness <= 1.0:
2701
+ raise ValueError("roundness attribute must be float between (0, 1.0]")
2702
+ self.roundness: float = roundness
2703
+
2704
+ @ensure_cv2_image_for_class_method
2705
+ def annotate(
2706
+ self,
2707
+ scene: ImageType,
2708
+ detections: Detections,
2709
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
2710
+ ) -> ImageType:
2711
+ """
2712
+ Annotates the given scene with bounding boxes with rounded edges
2713
+ based on the provided detections.
2714
+
2715
+ Args:
2716
+ scene: The image where rounded bounding boxes will be drawn.
2717
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
2718
+ or `PIL.Image.Image`.
2719
+ detections: Object detections to annotate.
2720
+ custom_color_lookup: Custom color lookup array.
2721
+ Allows to override the default color mapping strategy.
2722
+
2723
+ Returns:
2724
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
2725
+ or `PIL.Image.Image`)
2726
+
2727
+ Examples:
2728
+ ```pycon
2729
+ >>> import numpy as np
2730
+ >>> import supervision as sv
2731
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
2732
+ >>> detections = sv.Detections(
2733
+ ... xyxy=np.array([[20, 20, 80, 80]]),
2734
+ ... class_id=np.array([0])
2735
+ ... )
2736
+ >>> round_box_annotator = sv.RoundBoxAnnotator()
2737
+ >>> annotated_frame = round_box_annotator.annotate(
2738
+ ... scene=image.copy(),
2739
+ ... detections=detections
2740
+ ... )
2741
+
2742
+ ```
2743
+
2744
+ ![round-box-annotator-example](https://media.roboflow.com/
2745
+ supervision-annotator-examples/round-box-annotator-example-purple.png)
2746
+ """
2747
+ if not isinstance(scene, np.ndarray):
2748
+ return scene
2749
+ xyxy = to_numpy(detections.xyxy).astype(int)
2750
+ for detection_idx, (x1, y1, x2, y2) in enumerate(xyxy):
2751
+ color = resolve_color(
2752
+ color=self.color,
2753
+ detections=detections,
2754
+ detection_idx=detection_idx,
2755
+ color_lookup=self.color_lookup
2756
+ if custom_color_lookup is None
2757
+ else custom_color_lookup,
2758
+ )
2759
+
2760
+ radius = (
2761
+ int((x2 - x1) // 2 * self.roundness)
2762
+ if abs(x1 - x2) < abs(y1 - y2)
2763
+ else int((y2 - y1) // 2 * self.roundness)
2764
+ )
2765
+
2766
+ circle_coordinates = [
2767
+ ((x1 + radius), (y1 + radius)),
2768
+ ((x2 - radius), (y1 + radius)),
2769
+ ((x2 - radius), (y2 - radius)),
2770
+ ((x1 + radius), (y2 - radius)),
2771
+ ]
2772
+
2773
+ line_coordinates = [
2774
+ ((x1 + radius, y1), (x2 - radius, y1)),
2775
+ ((x2, y1 + radius), (x2, y2 - radius)),
2776
+ ((x1 + radius, y2), (x2 - radius, y2)),
2777
+ ((x1, y1 + radius), (x1, y2 - radius)),
2778
+ ]
2779
+
2780
+ start_angles = (180, 270, 0, 90)
2781
+ end_angles = (270, 360, 90, 180)
2782
+
2783
+ for center_coordinates, line, start_angle, end_angle in zip(
2784
+ circle_coordinates, line_coordinates, start_angles, end_angles
2785
+ ):
2786
+ cv2.ellipse(
2787
+ img=scene,
2788
+ center=center_coordinates,
2789
+ axes=(radius, radius),
2790
+ angle=0,
2791
+ startAngle=start_angle,
2792
+ endAngle=end_angle,
2793
+ color=color.as_bgr(),
2794
+ thickness=self.thickness,
2795
+ )
2796
+
2797
+ cv2.line(
2798
+ img=scene,
2799
+ pt1=line[0],
2800
+ pt2=line[1],
2801
+ color=color.as_bgr(),
2802
+ thickness=self.thickness,
2803
+ )
2804
+
2805
+ return scene
2806
+
2807
+
2808
+ class PercentageBarAnnotator(BaseAnnotator):
2809
+ """
2810
+ A class for drawing percentage bars on an image using provided detections.
2811
+ """
2812
+
2813
+ def __init__(
2814
+ self,
2815
+ height: int = 16,
2816
+ width: int = 80,
2817
+ color: Color | ColorPalette | str = ColorPalette.DEFAULT,
2818
+ border_color: Color | str = Color.BLACK,
2819
+ position: Position = Position.TOP_CENTER,
2820
+ color_lookup: ColorLookup = ColorLookup.CLASS,
2821
+ border_thickness: int | None = None,
2822
+ ):
2823
+ """
2824
+ Args:
2825
+ height: The height in pixels of the percentage bar.
2826
+ width: The width in pixels of the percentage bar.
2827
+ color: The color or color palette to use for
2828
+ annotating detections.
2829
+ border_color: The color of the border lines.
2830
+ position: The anchor position of drawing the percentage bar.
2831
+ color_lookup: Strategy for mapping colors to annotations.
2832
+ Options are `INDEX`, `CLASS`, `TRACK`.
2833
+ border_thickness: The thickness of the border lines.
2834
+ """
2835
+ self.height: int = height
2836
+ self.width: int = width
2837
+ self.color: Color | ColorPalette = _normalize_color_input(color)
2838
+ self.border_color = cast(Color, _normalize_color_input(border_color))
2839
+ self.position: Position = position
2840
+ self.color_lookup: ColorLookup = color_lookup
2841
+
2842
+ self.border_thickness = (
2843
+ border_thickness
2844
+ if border_thickness is not None
2845
+ else int(0.15 * self.height)
2846
+ )
2847
+
2848
+ @ensure_cv2_image_for_class_method
2849
+ def annotate(
2850
+ self,
2851
+ scene: ImageType,
2852
+ detections: Detections,
2853
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
2854
+ custom_values: npt.NDArray[np.float64] | None = None,
2855
+ ) -> ImageType:
2856
+ """
2857
+ Annotates the given scene with percentage bars based on the provided
2858
+ detections. The percentage bars visually represent the confidence or custom
2859
+ values associated with each detection.
2860
+
2861
+ Args:
2862
+ scene: The image where percentage bars will be drawn.
2863
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
2864
+ or `PIL.Image.Image`.
2865
+ detections: Object detections to annotate.
2866
+ custom_color_lookup: Custom color lookup array.
2867
+ Allows to override the default color mapping strategy.
2868
+ custom_values: Custom values array to use instead
2869
+ of the default detection confidences. This array should have the
2870
+ same length as the number of detections and contain a value between
2871
+ 0 and 1 (inclusive) for each detection, representing the percentage
2872
+ to be displayed.
2873
+
2874
+ Returns:
2875
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
2876
+ or `PIL.Image.Image`)
2877
+
2878
+ Examples:
2879
+ ```pycon
2880
+ >>> import numpy as np
2881
+ >>> import supervision as sv
2882
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
2883
+ >>> detections = sv.Detections(
2884
+ ... xyxy=np.array([[20, 20, 80, 80]]),
2885
+ ... confidence=np.array([0.9]),
2886
+ ... class_id=np.array([0])
2887
+ ... )
2888
+ >>> percentage_bar_annotator = sv.PercentageBarAnnotator()
2889
+ >>> annotated_frame = percentage_bar_annotator.annotate(
2890
+ ... scene=image.copy(),
2891
+ ... detections=detections
2892
+ ... )
2893
+
2894
+ ```
2895
+
2896
+ ![percentage-bar-example](https://media.roboflow.com/
2897
+ supervision-annotator-examples/percentage-bar-annotator-example-purple.png)
2898
+ """
2899
+ if not isinstance(scene, np.ndarray):
2900
+ return scene
2901
+ self._validate_custom_values(custom_values=custom_values, detections=detections)
2902
+
2903
+ anchors = to_numpy(
2904
+ detections.get_anchors_coordinates(anchor=self.position)
2905
+ ).astype(int)
2906
+ confidence = (
2907
+ to_numpy(detections.confidence)
2908
+ if detections.confidence is not None
2909
+ else None
2910
+ )
2911
+ for detection_idx in range(len(detections)):
2912
+ anchor = anchors[detection_idx]
2913
+ border_coordinates = self.calculate_border_coordinates(
2914
+ anchor_xy=(int(anchor[0]), int(anchor[1])),
2915
+ border_wh=(self.width, self.height),
2916
+ position=self.position,
2917
+ )
2918
+ border_width = border_coordinates[1][0] - border_coordinates[0][0]
2919
+
2920
+ if custom_values is not None:
2921
+ value = custom_values[detection_idx]
2922
+ else:
2923
+ assert confidence is not None # MyPy type hint
2924
+ value = confidence[detection_idx]
2925
+
2926
+ color = resolve_color(
2927
+ color=self.color,
2928
+ detections=detections,
2929
+ detection_idx=detection_idx,
2930
+ color_lookup=self.color_lookup
2931
+ if custom_color_lookup is None
2932
+ else custom_color_lookup,
2933
+ )
2934
+ cv2.rectangle(
2935
+ img=scene,
2936
+ pt1=border_coordinates[0],
2937
+ pt2=(
2938
+ border_coordinates[0][0] + int(border_width * value),
2939
+ border_coordinates[1][1],
2940
+ ),
2941
+ color=color.as_bgr(),
2942
+ thickness=-1,
2943
+ )
2944
+ cv2.rectangle(
2945
+ img=scene,
2946
+ pt1=border_coordinates[0],
2947
+ pt2=border_coordinates[1],
2948
+ color=self.border_color.as_bgr(),
2949
+ thickness=self.border_thickness,
2950
+ )
2951
+ return scene
2952
+
2953
+ @staticmethod
2954
+ def calculate_border_coordinates(
2955
+ anchor_xy: tuple[int, int], border_wh: tuple[int, int], position: Position
2956
+ ) -> tuple[tuple[int, int], tuple[int, int]]:
2957
+ """Compute the border corner coordinates for a given anchor position."""
2958
+ cx, cy = anchor_xy
2959
+ width, height = border_wh
2960
+
2961
+ if position == Position.TOP_LEFT:
2962
+ return (cx - width, cy - height), (cx, cy)
2963
+ elif position == Position.TOP_CENTER:
2964
+ return (cx - width // 2, cy), (cx + width // 2, cy - height)
2965
+ elif position == Position.TOP_RIGHT:
2966
+ return (cx, cy), (cx + width, cy - height)
2967
+ elif position == Position.CENTER_LEFT:
2968
+ return (cx - width, cy - height // 2), (cx, cy + height // 2)
2969
+ elif position == Position.CENTER or position == Position.CENTER_OF_MASS:
2970
+ return (
2971
+ (cx - width // 2, cy - height // 2),
2972
+ (cx + width // 2, cy + height // 2),
2973
+ )
2974
+ elif position == Position.CENTER_RIGHT:
2975
+ return (cx, cy - height // 2), (cx + width, cy + height // 2)
2976
+ elif position == Position.BOTTOM_LEFT:
2977
+ return (cx - width, cy), (cx, cy + height)
2978
+ elif position == Position.BOTTOM_CENTER:
2979
+ return (cx - width // 2, cy), (cx + width // 2, cy + height)
2980
+ elif position == Position.BOTTOM_RIGHT:
2981
+ return (cx, cy), (cx + width, cy + height)
2982
+ raise ValueError(f"Unsupported position: {position}")
2983
+
2984
+ @staticmethod
2985
+ def _validate_custom_values(
2986
+ custom_values: npt.NDArray[np.float64] | list[float] | None,
2987
+ detections: Detections,
2988
+ ) -> None:
2989
+ if custom_values is None:
2990
+ if detections.confidence is None:
2991
+ raise ValueError(
2992
+ "The provided detections do not contain confidence values. "
2993
+ "Please provide `custom_values` or ensure that the detections "
2994
+ "contain confidence values (e.g. by using a different model)."
2995
+ )
2996
+
2997
+ else:
2998
+ if not isinstance(custom_values, (np.ndarray, list)):
2999
+ raise TypeError(
3000
+ "custom_values must be either a numpy array or a list of floats."
3001
+ )
3002
+
3003
+ if len(custom_values) != len(detections):
3004
+ raise ValueError(
3005
+ "The length of custom_values must match the number of detections."
3006
+ )
3007
+
3008
+ if not all(0 <= value <= 1 for value in custom_values):
3009
+ raise ValueError("All values in custom_values must be between 0 and 1.")
3010
+
3011
+ @staticmethod
3012
+ @deprecated( # type: ignore[untyped-decorator]
3013
+ target=_validate_custom_values.__func__, # type: ignore[attr-defined]
3014
+ deprecated_in="0.29.0",
3015
+ remove_in="0.32.0",
3016
+ )
3017
+ def validate_custom_values(
3018
+ custom_values: npt.NDArray[np.float64] | list[float] | None,
3019
+ detections: Detections,
3020
+ ) -> None:
3021
+ void(custom_values, detections)
3022
+
3023
+
3024
+ class CropAnnotator(BaseAnnotator):
3025
+ """
3026
+ A class for drawing scaled up crops of detections on the scene.
3027
+ """
3028
+
3029
+ def __init__(
3030
+ self,
3031
+ position: Position = Position.TOP_CENTER,
3032
+ scale_factor: float = 2.0,
3033
+ border_color: Color | ColorPalette | str = ColorPalette.DEFAULT,
3034
+ border_thickness: int = 2,
3035
+ border_color_lookup: ColorLookup = ColorLookup.CLASS,
3036
+ ):
3037
+ """
3038
+ Args:
3039
+ position: The anchor position for placing the cropped and scaled
3040
+ part of the detection in the scene.
3041
+ scale_factor: The factor by which to scale the cropped image part. A
3042
+ factor of 2, for example, would double the size of the cropped area,
3043
+ allowing for a closer view of the detection.
3044
+ border_color: The color or color palette to
3045
+ use for annotating border around the cropped area.
3046
+ border_thickness: The thickness of the border around the cropped area.
3047
+ border_color_lookup: Strategy for mapping colors to
3048
+ annotations. Options are `INDEX`, `CLASS`, `TRACK`.
3049
+ """
3050
+ self.position: Position = position
3051
+ self.scale_factor: float = scale_factor
3052
+ self.border_color: Color | ColorPalette = _normalize_color_input(border_color)
3053
+ self.border_thickness: int = border_thickness
3054
+ self.border_color_lookup: ColorLookup = border_color_lookup
3055
+
3056
+ @ensure_cv2_image_for_class_method
3057
+ def annotate(
3058
+ self,
3059
+ scene: ImageType,
3060
+ detections: Detections,
3061
+ custom_color_lookup: npt.NDArray[np.int_] | None = None,
3062
+ ) -> ImageType:
3063
+ """
3064
+ Annotates the provided scene with scaled and cropped parts of the image based
3065
+ on the provided detections. Each detection is cropped from the original scene
3066
+ and scaled according to the annotator's scale factor before being placed back
3067
+ onto the scene at the specified position.
3068
+
3069
+
3070
+ Args:
3071
+ scene: The image where cropped detection will be placed.
3072
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
3073
+ or `PIL.Image.Image`.
3074
+ detections: Object detections to annotate.
3075
+ custom_color_lookup: Custom color lookup array.
3076
+ Allows to override the default color mapping strategy.
3077
+
3078
+ Returns:
3079
+ The annotated image.
3080
+
3081
+ Note:
3082
+ Detections whose bounding boxes extend partially outside `scene` are
3083
+ clipped to scene bounds before cropping. Detections fully outside the
3084
+ scene collapse to zero area after clipping and are skipped without
3085
+ raising an error.
3086
+
3087
+ Examples:
3088
+ ```pycon
3089
+ >>> import numpy as np
3090
+ >>> import supervision as sv
3091
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
3092
+ >>> detections = sv.Detections(
3093
+ ... xyxy=np.array([[20, 20, 80, 80]]),
3094
+ ... class_id=np.array([0])
3095
+ ... )
3096
+ >>> crop_annotator = sv.CropAnnotator()
3097
+ >>> annotated_frame = crop_annotator.annotate(
3098
+ ... scene=image.copy(),
3099
+ ... detections=detections
3100
+ ... )
3101
+
3102
+ ```
3103
+
3104
+ ![crop-annotator-example](https://media.roboflow.com/
3105
+ supervision-annotator-examples/crop-annotator-example.png)
3106
+ """
3107
+ if not isinstance(scene, np.ndarray):
3108
+ return scene
3109
+ image_height, image_width = scene.shape[:2]
3110
+ clipped_xyxy: npt.NDArray[np.int32] = clip_boxes(
3111
+ xyxy=to_numpy(detections.xyxy),
3112
+ resolution_wh=(image_width, image_height),
3113
+ ).astype(np.int32)
3114
+ anchors: npt.NDArray[np.int32] = detections.get_anchors_coordinates(
3115
+ anchor=self.position
3116
+ ).astype(np.int32)
3117
+ # Snapshot before the loop so later crops are taken from the original image,
3118
+ # not a scene already annotated by earlier iterations (overlapping-box case).
3119
+ source_scene = scene.copy()
3120
+
3121
+ for idx, (xyxy, anchor) in enumerate(zip(clipped_xyxy, anchors)):
3122
+ crop_x1, crop_y1, crop_x2, crop_y2 = xyxy
3123
+ if crop_x2 <= crop_x1 or crop_y2 <= crop_y1:
3124
+ continue
3125
+ crop = crop_image(image=source_scene, xyxy=xyxy)
3126
+ resized_crop = scale_image(image=crop, scale_factor=self.scale_factor)
3127
+ crop_wh = resized_crop.shape[1], resized_crop.shape[0]
3128
+ (x1, y1), (x2, y2) = self.calculate_crop_coordinates(
3129
+ anchor=anchor, crop_wh=crop_wh, position=self.position
3130
+ )
3131
+ scene = _overlay_image(image=scene, overlay=resized_crop, anchor=(x1, y1))
3132
+ color = resolve_color(
3133
+ color=self.border_color,
3134
+ detections=detections,
3135
+ detection_idx=idx,
3136
+ color_lookup=self.border_color_lookup
3137
+ if custom_color_lookup is None
3138
+ else custom_color_lookup,
3139
+ )
3140
+ cv2.rectangle(
3141
+ img=scene,
3142
+ pt1=(x1, y1),
3143
+ pt2=(x2, y2),
3144
+ color=color.as_bgr(),
3145
+ thickness=self.border_thickness,
3146
+ )
3147
+
3148
+ return scene
3149
+
3150
+ @staticmethod
3151
+ def calculate_crop_coordinates(
3152
+ anchor: tuple[int, int], crop_wh: tuple[int, int], position: Position
3153
+ ) -> tuple[tuple[int, int], tuple[int, int]]:
3154
+ """Compute the crop coordinates for a given anchor position."""
3155
+ anchor_x, anchor_y = anchor
3156
+ width, height = crop_wh
3157
+
3158
+ if position == Position.TOP_LEFT:
3159
+ return (anchor_x - width, anchor_y - height), (anchor_x, anchor_y)
3160
+ elif position == Position.TOP_CENTER:
3161
+ return (
3162
+ (anchor_x - width // 2, anchor_y - height),
3163
+ (anchor_x + width // 2, anchor_y),
3164
+ )
3165
+ elif position == Position.TOP_RIGHT:
3166
+ return (anchor_x, anchor_y - height), (anchor_x + width, anchor_y)
3167
+ elif position == Position.CENTER_LEFT:
3168
+ return (
3169
+ (anchor_x - width, anchor_y - height // 2),
3170
+ (anchor_x, anchor_y + height // 2),
3171
+ )
3172
+ elif position == Position.CENTER or position == Position.CENTER_OF_MASS:
3173
+ return (
3174
+ (anchor_x - width // 2, anchor_y - height // 2),
3175
+ (anchor_x + width // 2, anchor_y + height // 2),
3176
+ )
3177
+ elif position == Position.CENTER_RIGHT:
3178
+ return (
3179
+ (anchor_x, anchor_y - height // 2),
3180
+ (anchor_x + width, anchor_y + height // 2),
3181
+ )
3182
+ elif position == Position.BOTTOM_LEFT:
3183
+ return (anchor_x - width, anchor_y), (anchor_x, anchor_y + height)
3184
+ elif position == Position.BOTTOM_CENTER:
3185
+ return (
3186
+ (anchor_x - width // 2, anchor_y),
3187
+ (anchor_x + width // 2, anchor_y + height),
3188
+ )
3189
+ elif position == Position.BOTTOM_RIGHT:
3190
+ return (anchor_x, anchor_y), (anchor_x + width, anchor_y + height)
3191
+ raise ValueError(f"Unsupported position: {position}")
3192
+
3193
+
3194
+ class BackgroundOverlayAnnotator(BaseAnnotator):
3195
+ """
3196
+ A class for drawing a colored overlay on the background of an image outside
3197
+ the region of detections.
3198
+
3199
+ If masks are provided, the background is colored outside the masks.
3200
+ If masks are not provided, the background is colored outside the bounding boxes.
3201
+
3202
+ You can use the `force_box` parameter to force the annotator to use bounding boxes.
3203
+
3204
+ !!! warning
3205
+
3206
+ This annotator uses `sv.Detections.mask`.
3207
+ """
3208
+
3209
+ def __init__(
3210
+ self,
3211
+ color: Color = Color.BLACK,
3212
+ opacity: float = 0.5,
3213
+ force_box: bool = False,
3214
+ ):
3215
+ """
3216
+ Args:
3217
+ color: The color to use for annotating detections.
3218
+ opacity: Opacity of the overlay mask. Must be between `0` and `1`.
3219
+ force_box: If `True`, forces the annotator to use bounding boxes when
3220
+ masks are provided in the supplied sv.Detections.
3221
+ """
3222
+ self.color: Color = color
3223
+ self.opacity = opacity
3224
+ self.force_box = force_box
3225
+
3226
+ @ensure_cv2_image_for_class_method
3227
+ def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
3228
+ """
3229
+ Applies a colored overlay to the scene outside of the detected regions.
3230
+
3231
+ Args:
3232
+ scene: The image where masks will be drawn.
3233
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
3234
+ or `PIL.Image.Image`.
3235
+ detections: Object detections to annotate.
3236
+
3237
+ Returns:
3238
+ The annotated image, matching the type of `scene` (`numpy.ndarray`
3239
+ or `PIL.Image.Image`)
3240
+
3241
+ Examples:
3242
+ ```pycon
3243
+ >>> import numpy as np
3244
+ >>> import supervision as sv
3245
+ >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
3246
+ >>> detections = sv.Detections(
3247
+ ... xyxy=np.array([[20, 20, 80, 80]]),
3248
+ ... class_id=np.array([0])
3249
+ ... )
3250
+ >>> background_overlay_annotator = sv.BackgroundOverlayAnnotator()
3251
+ >>> annotated_frame = background_overlay_annotator.annotate(
3252
+ ... scene=image.copy(),
3253
+ ... detections=detections
3254
+ ... )
3255
+
3256
+ ```
3257
+
3258
+ ![background-overlay-annotator-example](https://media.roboflow.com/
3259
+ supervision-annotator-examples/background-color-annotator-example-purple.png)
3260
+ """
3261
+ if not isinstance(scene, np.ndarray):
3262
+ return scene
3263
+ colored_mask = np.full_like(scene, self.color.as_bgr(), dtype=np.uint8)
3264
+
3265
+ cv2.addWeighted(
3266
+ scene, 1 - self.opacity, colored_mask, self.opacity, 0, dst=colored_mask
3267
+ )
3268
+
3269
+ if detections.mask is None or self.force_box:
3270
+ image_height, image_width = scene.shape[:2]
3271
+ clipped_xyxy: npt.NDArray[np.int32] = clip_boxes(
3272
+ xyxy=to_numpy(detections.xyxy),
3273
+ resolution_wh=(image_width, image_height),
3274
+ ).astype(np.int32)
3275
+ for x1, y1, x2, y2 in clipped_xyxy:
3276
+ colored_mask[y1:y2, x1:x2] = scene[y1:y2, x1:x2]
3277
+ else:
3278
+ masks = to_numpy(detections.mask).astype(bool, copy=False)
3279
+ for mask_bool in masks:
3280
+ colored_mask[mask_bool] = scene[mask_bool]
3281
+
3282
+ np.copyto(scene, colored_mask)
3283
+ return scene
3284
+
3285
+
3286
+ class ComparisonAnnotator:
3287
+ """
3288
+ Highlights the differences between two sets of detections.
3289
+ Useful for comparing results from two different models, or the difference
3290
+ between a ground truth and a prediction.
3291
+
3292
+ If present, uses the oriented bounding box data.
3293
+ Otherwise, if present, uses a mask.
3294
+ Otherwise, uses the bounding box data.
3295
+ """
3296
+
3297
+ # Not a BaseAnnotator subclass — duck-typing callers can still check requires_mask
3298
+ requires_mask: ClassVar[bool] = False
3299
+
3300
+ def __init__(
3301
+ self,
3302
+ color_1: Color = Color.RED,
3303
+ color_2: Color = Color.GREEN,
3304
+ color_overlap: Color = Color.BLUE,
3305
+ *,
3306
+ opacity: float = 0.75,
3307
+ label_1: str = "",
3308
+ label_2: str = "",
3309
+ label_overlap: str = "",
3310
+ label_scale: float = 1.0,
3311
+ ):
3312
+ """
3313
+ Args:
3314
+ color_1: Color of areas only present in the first set of
3315
+ detections.
3316
+ color_2: Color of areas only present in the second set of
3317
+ detections.
3318
+ color_overlap: Color of areas present in both sets of detections.
3319
+ opacity: Opacity of the overlay mask. Must be between `0` and `1`.
3320
+ label_1: Label for the first set of detections.
3321
+ label_2: Label for the second set of detections.
3322
+ label_overlap: Label for areas present in both sets of detections.
3323
+ label_scale: Controls how large the labels are.
3324
+ """
3325
+
3326
+ self.color_1 = color_1
3327
+ self.color_2 = color_2
3328
+ self.color_overlap = color_overlap
3329
+
3330
+ self.opacity = opacity
3331
+ self.label_1 = label_1
3332
+ self.label_2 = label_2
3333
+ self.label_overlap = label_overlap
3334
+ self.label_scale = label_scale
3335
+ self.text_thickness = int(self.label_scale + 1.2)
3336
+
3337
+ @ensure_cv2_image_for_class_method
3338
+ def annotate(
3339
+ self, scene: ImageType, detections_1: Detections, detections_2: Detections
3340
+ ) -> ImageType:
3341
+ """
3342
+ Highlights the differences between two sets of detections.
3343
+
3344
+ Args:
3345
+ scene: The image where detections will be drawn.
3346
+ `ImageType` is a flexible type, accepting either `numpy.ndarray`
3347
+ or `PIL.Image.Image`.
3348
+ detections_1: The first set of detections or predictions.
3349
+ detections_2: The second set of detections to compare or
3350
+ ground truth.
3351
+
3352
+ Returns:
3353
+ The annotated image.
3354
+
3355
+ Example:
3356
+ ```python
3357
+ import supervision as sv
3358
+
3359
+ image = ...
3360
+ detections_1 = sv.Detections(...)
3361
+ detections_2 = sv.Detections(...)
3362
+
3363
+ comparison_annotator = sv.ComparisonAnnotator()
3364
+ annotated_frame = comparison_annotator.annotate(
3365
+ scene=image.copy(),
3366
+ detections_1=detections_1,
3367
+ detections_2=detections_2
3368
+ )
3369
+ ```
3370
+
3371
+ ![comparison-annotator-example](https://media.roboflow.com/
3372
+ supervision-annotator-examples/comparison-annotator-example.png)
3373
+ """
3374
+ if not isinstance(scene, np.ndarray):
3375
+ return scene
3376
+ if detections_1.is_empty() and detections_2.is_empty():
3377
+ return scene
3378
+
3379
+ use_obb = self._use_obb(detections_1, detections_2)
3380
+ use_mask = self._use_mask(detections_1, detections_2)
3381
+
3382
+ if use_obb:
3383
+ mask_1 = self._mask_from_obb(scene, detections_1)
3384
+ mask_2 = self._mask_from_obb(scene, detections_2)
3385
+
3386
+ elif use_mask:
3387
+ mask_1 = self._mask_from_mask(scene, detections_1)
3388
+ mask_2 = self._mask_from_mask(scene, detections_2)
3389
+
3390
+ else:
3391
+ mask_1 = self._mask_from_xyxy(scene, detections_1)
3392
+ mask_2 = self._mask_from_xyxy(scene, detections_2)
3393
+
3394
+ mask_overlap = mask_1 & mask_2
3395
+ mask_1 = mask_1 & ~mask_overlap
3396
+ mask_2 = mask_2 & ~mask_overlap
3397
+
3398
+ color_layer = np.zeros_like(scene, dtype=np.uint8)
3399
+ color_layer[mask_overlap] = self.color_overlap.as_bgr()
3400
+ color_layer[mask_1] = self.color_1.as_bgr()
3401
+ color_layer[mask_2] = self.color_2.as_bgr()
3402
+
3403
+ scene[mask_overlap] = (1 - self.opacity) * scene[
3404
+ mask_overlap
3405
+ ] + self.opacity * color_layer[mask_overlap]
3406
+ scene[mask_1] = (1 - self.opacity) * scene[mask_1] + self.opacity * color_layer[
3407
+ mask_1
3408
+ ]
3409
+ scene[mask_2] = (1 - self.opacity) * scene[mask_2] + self.opacity * color_layer[
3410
+ mask_2
3411
+ ]
3412
+
3413
+ self._draw_labels(scene)
3414
+
3415
+ return scene
3416
+
3417
+ @staticmethod
3418
+ def _use_obb(detections_1: Detections, detections_2: Detections) -> bool:
3419
+ assert not detections_1.is_empty() or not detections_2.is_empty()
3420
+ is_obb_1 = ORIENTED_BOX_COORDINATES in detections_1.data
3421
+ is_obb_2 = ORIENTED_BOX_COORDINATES in detections_2.data
3422
+ return (
3423
+ (is_obb_1 and is_obb_2)
3424
+ or (is_obb_1 and detections_2.is_empty())
3425
+ or (detections_1.is_empty() and is_obb_2)
3426
+ )
3427
+
3428
+ @staticmethod
3429
+ def _use_mask(detections_1: Detections, detections_2: Detections) -> bool:
3430
+ assert not detections_1.is_empty() or not detections_2.is_empty()
3431
+ is_mask_1 = detections_1.mask is not None
3432
+ is_mask_2 = detections_2.mask is not None
3433
+ return (
3434
+ (is_mask_1 and is_mask_2)
3435
+ or (is_mask_1 and detections_2.is_empty())
3436
+ or (detections_1.is_empty() and is_mask_2)
3437
+ )
3438
+
3439
+ @staticmethod
3440
+ def _mask_from_xyxy(
3441
+ scene: npt.NDArray[np.uint8], detections: Detections
3442
+ ) -> npt.NDArray[np.bool_]:
3443
+ mask = np.zeros(scene.shape[:2], dtype=np.bool_)
3444
+ if detections.is_empty():
3445
+ return mask
3446
+
3447
+ resolution_wh = scene.shape[1], scene.shape[0]
3448
+ polygons = xyxy_to_polygons(to_numpy(detections.xyxy))
3449
+
3450
+ for polygon in polygons:
3451
+ polygon_mask = polygon_to_mask(polygon, resolution_wh=resolution_wh)
3452
+ mask |= polygon_mask.astype(np.bool_)
3453
+ return mask
3454
+
3455
+ @staticmethod
3456
+ def _mask_from_obb(
3457
+ scene: npt.NDArray[np.uint8], detections: Detections
3458
+ ) -> npt.NDArray[np.bool_]:
3459
+ mask = np.zeros(scene.shape[:2], dtype=np.bool_)
3460
+ if detections.is_empty():
3461
+ return mask
3462
+
3463
+ resolution_wh = scene.shape[1], scene.shape[0]
3464
+
3465
+ for polygon in to_numpy(detections.data[ORIENTED_BOX_COORDINATES]):
3466
+ polygon_mask = polygon_to_mask(polygon, resolution_wh=resolution_wh)
3467
+ mask |= polygon_mask.astype(np.bool_)
3468
+ return mask
3469
+
3470
+ @staticmethod
3471
+ def _mask_from_mask(
3472
+ scene: npt.NDArray[np.uint8], detections: Detections
3473
+ ) -> npt.NDArray[np.bool_]:
3474
+ mask = np.zeros(scene.shape[:2], dtype=np.bool_)
3475
+ if detections.is_empty():
3476
+ return mask
3477
+ assert detections.mask is not None
3478
+
3479
+ for detections_mask in to_numpy(detections.mask):
3480
+ mask |= detections_mask.astype(np.bool_)
3481
+ return mask
3482
+
3483
+ def _draw_labels(self, scene: npt.NDArray[np.uint8]) -> None:
3484
+ """
3485
+ Draw the labels, explaining what each color represents, with automatically
3486
+ computed positions.
3487
+
3488
+ Args:
3489
+ scene: The image where the labels will be drawn.
3490
+ """
3491
+ margin = int(50 * self.label_scale)
3492
+ gap = int(40 * self.label_scale)
3493
+ y0 = int(50 * self.label_scale)
3494
+ height = int(50 * self.label_scale)
3495
+
3496
+ marker_size = int(20 * self.label_scale)
3497
+ padding = int(10 * self.label_scale)
3498
+ text_box_corner_radius = int(10 * self.label_scale)
3499
+ marker_corner_radius = int(4 * self.label_scale)
3500
+ text_scale = self.label_scale
3501
+
3502
+ label_color_pairs = [
3503
+ (self.label_1, self.color_1),
3504
+ (self.label_2, self.color_2),
3505
+ (self.label_overlap, self.color_overlap),
3506
+ ]
3507
+
3508
+ x0 = margin
3509
+ for text, color in label_color_pairs:
3510
+ if not text:
3511
+ continue
3512
+
3513
+ (text_w, _) = cv2.getTextSize(
3514
+ text=text,
3515
+ fontFace=CV2_FONT,
3516
+ fontScale=self.label_scale,
3517
+ thickness=self.text_thickness,
3518
+ )[0]
3519
+
3520
+ width = text_w + marker_size + padding * 4
3521
+ center_x = x0 + width // 2
3522
+ center_y = y0 + height // 2
3523
+
3524
+ draw_rounded_rectangle(
3525
+ scene=scene,
3526
+ rect=Rect(x=x0, y=y0, width=width, height=height),
3527
+ color=Color.WHITE,
3528
+ border_radius=text_box_corner_radius,
3529
+ )
3530
+
3531
+ draw_rounded_rectangle(
3532
+ scene=scene,
3533
+ rect=Rect(
3534
+ x=x0 + padding,
3535
+ y=center_y - marker_size / 2,
3536
+ width=marker_size,
3537
+ height=marker_size,
3538
+ ),
3539
+ color=color,
3540
+ border_radius=marker_corner_radius,
3541
+ )
3542
+
3543
+ draw_text(
3544
+ scene,
3545
+ text,
3546
+ text_anchor=Point(x=center_x + marker_size, y=center_y),
3547
+ text_scale=text_scale,
3548
+ text_thickness=self.text_thickness,
3549
+ )
3550
+
3551
+ x0 += width + gap