vidformer.supervision

  1"""
  2vidformer.supervision is the [supervision](https://supervision.roboflow.com/) frontend for [vidformer](https://github.com/ixlab/vidformer).
  3"""
  4
  5from math import sqrt
  6
  7import numpy as np
  8import supervision as _sv
  9from supervision import Color, ColorLookup, ColorPalette, Detections
 10from supervision.annotators.utils import resolve_color, resolve_text_background_xyxy
 11from supervision.config import CLASS_NAME_DATA_FIELD
 12
 13# supervision moved this between two versions, so we need to handle both cases
 14try:
 15    from supervision.detection.utils import spread_out_boxes
 16except ImportError:
 17    from supervision.detection.utils.boxes import spread_out_boxes
 18
 19from supervision.geometry.core import Position
 20
 21import vidformer.cv2 as vf_cv2
 22
 23try:
 24    import cv2 as ocv_cv2
 25except ImportError:
 26    ocv_cv2 = None
 27
 28CV2_FONT = vf_cv2.FONT_HERSHEY_SIMPLEX
 29
 30
 31class BoxAnnotator:
 32    def __init__(
 33        self,
 34        color=ColorPalette.DEFAULT,
 35        thickness=2,
 36        color_lookup=ColorLookup.CLASS,
 37    ):
 38        self.color = color
 39        self.thickness = thickness
 40        self.color_lookup = color_lookup
 41
 42    def annotate(
 43        self,
 44        scene: vf_cv2.Frame,
 45        detections: Detections,
 46        custom_color_lookup=None,
 47    ):
 48        for detection_idx in range(len(detections)):
 49            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
 50            color = resolve_color(
 51                color=self.color,
 52                detections=detections,
 53                detection_idx=detection_idx,
 54                color_lookup=(
 55                    self.color_lookup
 56                    if custom_color_lookup is None
 57                    else custom_color_lookup
 58                ),
 59            )
 60            vf_cv2.rectangle(
 61                img=scene,
 62                pt1=(x1, y1),
 63                pt2=(x2, y2),
 64                color=color.as_rgb(),
 65                thickness=self.thickness,
 66            )
 67        return scene
 68
 69
 70class RoundBoxAnnotator:
 71    def __init__(
 72        self,
 73        color=ColorPalette.DEFAULT,
 74        thickness: int = 2,
 75        color_lookup: ColorLookup = ColorLookup.CLASS,
 76        roundness: float = 0.6,
 77    ):
 78        self.color = color
 79        self.thickness = thickness
 80        self.color_lookup = color_lookup
 81        if not 0 < roundness <= 1.0:
 82            raise ValueError("roundness attribute must be float between (0, 1.0]")
 83        self.roundness = roundness
 84
 85    def annotate(
 86        self,
 87        scene: vf_cv2.Frame,
 88        detections: _sv.Detections,
 89        custom_color_lookup=None,
 90    ):
 91        for detection_idx in range(len(detections)):
 92            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
 93            color = resolve_color(
 94                color=self.color,
 95                detections=detections,
 96                detection_idx=detection_idx,
 97                color_lookup=(
 98                    self.color_lookup
 99                    if custom_color_lookup is None
100                    else custom_color_lookup
101                ),
102            )
103            radius = (
104                int((x2 - x1) // 2 * self.roundness)
105                if abs(x1 - x2) < abs(y1 - y2)
106                else int((y2 - y1) // 2 * self.roundness)
107            )
108            circle_coordinates = [
109                ((x1 + radius), (y1 + radius)),
110                ((x2 - radius), (y1 + radius)),
111                ((x2 - radius), (y2 - radius)),
112                ((x1 + radius), (y2 - radius)),
113            ]
114            line_coordinates = [
115                ((x1 + radius, y1), (x2 - radius, y1)),
116                ((x2, y1 + radius), (x2, y2 - radius)),
117                ((x1 + radius, y2), (x2 - radius, y2)),
118                ((x1, y1 + radius), (x1, y2 - radius)),
119            ]
120            start_angles = (180, 270, 0, 90)
121            end_angles = (270, 360, 90, 180)
122            for center_coordinates, line, start_angle, end_angle in zip(
123                circle_coordinates, line_coordinates, start_angles, end_angles
124            ):
125                vf_cv2.ellipse(
126                    img=scene,
127                    center=center_coordinates,
128                    axes=(radius, radius),
129                    angle=0,
130                    startAngle=start_angle,
131                    endAngle=end_angle,
132                    color=color.as_rgb(),
133                    thickness=self.thickness,
134                )
135                vf_cv2.line(
136                    img=scene,
137                    pt1=line[0],
138                    pt2=line[1],
139                    color=color.as_rgb(),
140                    thickness=self.thickness,
141                )
142        return scene
143
144
145class BoxCornerAnnotator:
146    def __init__(
147        self,
148        color=ColorPalette.DEFAULT,
149        thickness=4,
150        corner_length=15,
151        color_lookup=ColorLookup.CLASS,
152    ):
153        self.color = color
154        self.thickness: int = thickness
155        self.corner_length: int = corner_length
156        self.color_lookup: ColorLookup = color_lookup
157
158    def annotate(
159        self,
160        scene: vf_cv2.Frame,
161        detections: Detections,
162        custom_color_lookup=None,
163    ):
164        for detection_idx in range(len(detections)):
165            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
166            color = resolve_color(
167                color=self.color,
168                detections=detections,
169                detection_idx=detection_idx,
170                color_lookup=(
171                    self.color_lookup
172                    if custom_color_lookup is None
173                    else custom_color_lookup
174                ),
175            )
176            corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
177            for x, y in corners:
178                x_end = x + self.corner_length if x == x1 else x - self.corner_length
179                vf_cv2.line(
180                    scene, (x, y), (x_end, y), color.as_rgb(), thickness=self.thickness
181                )
182
183                y_end = y + self.corner_length if y == y1 else y - self.corner_length
184                vf_cv2.line(
185                    scene, (x, y), (x, y_end), color.as_rgb(), thickness=self.thickness
186                )
187        return scene
188
189
190class ColorAnnotator:
191    def __init__(
192        self,
193        color=ColorPalette.DEFAULT,
194        opacity: float = 0.5,
195        color_lookup: ColorLookup = ColorLookup.CLASS,
196    ):
197        self.color = color
198        self.color_lookup: ColorLookup = color_lookup
199        self.opacity = opacity
200
201    def annotate(
202        self,
203        scene: vf_cv2.Frame,
204        detections: Detections,
205        custom_color_lookup=None,
206    ):
207        scene_with_boxes = scene.copy()
208        for detection_idx in range(len(detections)):
209            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
210            color = resolve_color(
211                color=self.color,
212                detections=detections,
213                detection_idx=detection_idx,
214                color_lookup=(
215                    self.color_lookup
216                    if custom_color_lookup is None
217                    else custom_color_lookup
218                ),
219            )
220            vf_cv2.rectangle(
221                img=scene_with_boxes,
222                pt1=(x1, y1),
223                pt2=(x2, y2),
224                color=color.as_rgb(),
225                thickness=-1,
226            )
227
228        vf_cv2.addWeighted(
229            scene_with_boxes, self.opacity, scene, 1 - self.opacity, gamma=0, dst=scene
230        )
231        return scene
232
233
234class CircleAnnotator:
235    def __init__(
236        self,
237        color=ColorPalette.DEFAULT,
238        thickness: int = 2,
239        color_lookup: ColorLookup = ColorLookup.CLASS,
240    ):
241        self.color = color
242        self.thickness: int = thickness
243        self.color_lookup: ColorLookup = color_lookup
244
245    def annotate(
246        self,
247        scene: vf_cv2.Frame,
248        detections: Detections,
249        custom_color_lookup=None,
250    ):
251        for detection_idx in range(len(detections)):
252            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
253            center = ((x1 + x2) // 2, (y1 + y2) // 2)
254            distance = sqrt((x1 - center[0]) ** 2 + (y1 - center[1]) ** 2)
255            color = resolve_color(
256                color=self.color,
257                detections=detections,
258                detection_idx=detection_idx,
259                color_lookup=(
260                    self.color_lookup
261                    if custom_color_lookup is None
262                    else custom_color_lookup
263                ),
264            )
265            vf_cv2.circle(
266                img=scene,
267                center=center,
268                radius=int(distance),
269                color=color.as_rgb(),
270                thickness=self.thickness,
271            )
272
273        return scene
274
275
276class DotAnnotator:
277    def __init__(
278        self,
279        color=ColorPalette.DEFAULT,
280        radius: int = 4,
281        position: Position = Position.CENTER,
282        color_lookup: ColorLookup = ColorLookup.CLASS,
283        outline_thickness: int = 0,
284        outline_color=Color.BLACK,
285    ):
286        self.color = color
287        self.radius: int = radius
288        self.position: Position = position
289        self.color_lookup: ColorLookup = color_lookup
290        self.outline_thickness = outline_thickness
291        self.outline_color = outline_color
292
293    def annotate(
294        self,
295        scene: vf_cv2.Frame,
296        detections: Detections,
297        custom_color_lookup=None,
298    ):
299        xy = detections.get_anchors_coordinates(anchor=self.position)
300        for detection_idx in range(len(detections)):
301            color = resolve_color(
302                color=self.color,
303                detections=detections,
304                detection_idx=detection_idx,
305                color_lookup=(
306                    self.color_lookup
307                    if custom_color_lookup is None
308                    else custom_color_lookup
309                ),
310            )
311            center = (int(xy[detection_idx, 0]), int(xy[detection_idx, 1]))
312
313            vf_cv2.circle(scene, center, self.radius, color.as_rgb(), -1)
314            if self.outline_thickness:
315                outline_color = resolve_color(
316                    color=self.outline_color,
317                    detections=detections,
318                    detection_idx=detection_idx,
319                    color_lookup=(
320                        self.color_lookup
321                        if custom_color_lookup is None
322                        else custom_color_lookup
323                    ),
324                )
325                vf_cv2.circle(
326                    scene,
327                    center,
328                    self.radius,
329                    outline_color.as_rgb(),
330                    self.outline_thickness,
331                )
332        return scene
333
334
335class LabelAnnotator:
336    def __init__(
337        self,
338        color=ColorPalette.DEFAULT,
339        text_color=Color.WHITE,
340        text_scale: float = 0.5,
341        text_thickness: int = 1,
342        text_padding: int = 10,
343        text_position: Position = Position.TOP_LEFT,
344        color_lookup: ColorLookup = ColorLookup.CLASS,
345        border_radius: int = 0,
346        smart_position: bool = False,
347    ):
348        self.border_radius: int = border_radius
349        self.color = color
350        self.text_color = text_color
351        self.text_scale: float = text_scale
352        self.text_thickness: int = text_thickness
353        self.text_padding: int = text_padding
354        self.text_anchor: Position = text_position
355        self.color_lookup: ColorLookup = color_lookup
356        self.smart_position = smart_position
357
358    def annotate(
359        self,
360        scene,
361        detections: Detections,
362        labels,
363        custom_color_lookup=None,
364    ):
365        self._validate_labels(labels, detections)
366
367        labels = self._get_labels_text(detections, labels)
368        label_properties = self._get_label_properties(detections, labels)
369
370        if self.smart_position:
371            xyxy = label_properties[:, :4]
372            xyxy = spread_out_boxes(xyxy)
373            label_properties[:, :4] = xyxy
374
375        self._draw_labels(
376            scene=scene,
377            labels=labels,
378            label_properties=label_properties,
379            detections=detections,
380            custom_color_lookup=custom_color_lookup,
381        )
382
383        return scene
384
385    def _validate_labels(self, labels, detections: Detections):
386        if labels is not None and len(labels) != len(detections):
387            raise ValueError(
388                f"The number of labels ({len(labels)}) does not match the "
389                f"number of detections ({len(detections)}). Each detection "
390                f"should have exactly 1 label."
391            )
392
393    def _get_label_properties(
394        self,
395        detections: Detections,
396        labels,
397    ):
398        label_properties = []
399        anchors_coordinates = detections.get_anchors_coordinates(
400            anchor=self.text_anchor
401        ).astype(int)
402
403        for label, center_coords in zip(labels, anchors_coordinates):
404            (text_w, text_h) = vf_cv2.getTextSize(
405                text=label,
406                fontFace=CV2_FONT,
407                fontScale=self.text_scale,
408                thickness=self.text_thickness,
409            )[0]
410
411            width_padded = text_w + 2 * self.text_padding
412            height_padded = text_h + 2 * self.text_padding
413
414            text_background_xyxy = resolve_text_background_xyxy(
415                center_coordinates=tuple(center_coords),
416                text_wh=(width_padded, height_padded),
417                position=self.text_anchor,
418            )
419
420            label_properties.append(
421                [
422                    *text_background_xyxy,
423                    text_h,
424                ]
425            )
426
427        return np.array(label_properties).reshape(-1, 5)
428
429    @staticmethod
430    def _get_labels_text(detections: Detections, custom_labels):
431        if custom_labels is not None:
432            return custom_labels
433
434        labels = []
435        for idx in range(len(detections)):
436            if CLASS_NAME_DATA_FIELD in detections.data:
437                labels.append(detections.data[CLASS_NAME_DATA_FIELD][idx])
438            elif detections.class_id is not None:
439                labels.append(str(detections.class_id[idx]))
440            else:
441                labels.append(str(idx))
442        return labels
443
444    def _draw_labels(
445        self,
446        scene,
447        labels,
448        label_properties,
449        detections,
450        custom_color_lookup,
451    ) -> None:
452        assert len(labels) == len(label_properties) == len(detections), (
453            f"Number of label properties ({len(label_properties)}), "
454            f"labels ({len(labels)}) and detections ({len(detections)}) "
455            "do not match."
456        )
457
458        color_lookup = (
459            custom_color_lookup
460            if custom_color_lookup is not None
461            else self.color_lookup
462        )
463
464        for idx, label_property in enumerate(label_properties):
465            background_color = resolve_color(
466                color=self.color,
467                detections=detections,
468                detection_idx=idx,
469                color_lookup=color_lookup,
470            )
471            text_color = resolve_color(
472                color=self.text_color,
473                detections=detections,
474                detection_idx=idx,
475                color_lookup=color_lookup,
476            )
477
478            box_xyxy = label_property[:4]
479            text_height_padded = label_property[4]
480            self.draw_rounded_rectangle(
481                scene=scene,
482                xyxy=box_xyxy,
483                color=background_color.as_rgb(),
484                border_radius=self.border_radius,
485            )
486
487            text_x = box_xyxy[0] + self.text_padding
488            text_y = box_xyxy[1] + self.text_padding + text_height_padded
489            vf_cv2.putText(
490                img=scene,
491                text=labels[idx],
492                org=(text_x, text_y),
493                fontFace=CV2_FONT,
494                fontScale=self.text_scale,
495                color=text_color.as_rgb(),
496                thickness=self.text_thickness,
497                lineType=vf_cv2.LINE_AA,
498            )
499
500    @staticmethod
501    def draw_rounded_rectangle(
502        scene: np.ndarray,
503        xyxy,
504        color,
505        border_radius: int,
506    ) -> np.ndarray:
507        x1, y1, x2, y2 = xyxy
508        width = x2 - x1
509        height = y2 - y1
510
511        border_radius = min(border_radius, min(width, height) // 2)
512
513        if border_radius <= 0:
514            vf_cv2.rectangle(
515                img=scene,
516                pt1=(x1, y1),
517                pt2=(x2, y2),
518                color=color,
519                thickness=-1,
520            )
521        else:
522            rectangle_coordinates = [
523                ((x1 + border_radius, y1), (x2 - border_radius, y2)),
524                ((x1, y1 + border_radius), (x2, y2 - border_radius)),
525            ]
526            circle_centers = [
527                (x1 + border_radius, y1 + border_radius),
528                (x2 - border_radius, y1 + border_radius),
529                (x1 + border_radius, y2 - border_radius),
530                (x2 - border_radius, y2 - border_radius),
531            ]
532
533            for coordinates in rectangle_coordinates:
534                vf_cv2.rectangle(
535                    img=scene,
536                    pt1=coordinates[0],
537                    pt2=coordinates[1],
538                    color=color,
539                    thickness=-1,
540                )
541            for center in circle_centers:
542                vf_cv2.circle(
543                    img=scene,
544                    center=center,
545                    radius=border_radius,
546                    color=color,
547                    thickness=-1,
548                )
549        return scene
550
551
552class MaskAnnotator:
553    def __init__(
554        self,
555        color=ColorPalette.DEFAULT,
556        opacity: float = 0.5,
557        color_lookup: ColorLookup = ColorLookup.CLASS,
558    ):
559        self.color = color
560        self.opacity = opacity
561        self.color_lookup: ColorLookup = color_lookup
562
563    def annotate(
564        self,
565        scene,
566        detections: Detections,
567        custom_color_lookup=None,
568    ):
569        if detections.mask is None:
570            return scene
571
572        colored_mask = scene.copy()
573
574        for detection_idx in np.flip(np.argsort(detections.box_area)):
575            color = resolve_color(
576                color=self.color,
577                detections=detections,
578                detection_idx=detection_idx,
579                color_lookup=(
580                    self.color_lookup
581                    if custom_color_lookup is None
582                    else custom_color_lookup
583                ),
584            )
585            mask = detections.mask[detection_idx]
586            colored_mask[mask] = color.as_bgr()
587
588        vf_cv2.addWeighted(
589            colored_mask, self.opacity, scene, 1 - self.opacity, 0, dst=scene
590        )
591        return scene
592
593
594class MaskStreamWriter:
595    def __init__(self, path: str, shape: tuple):
596        # Shape should be (width, height)
597        assert ocv_cv2 is not None, "OpenCV cv2 is required for ExternDetectionsBuilder"
598        assert type(shape) is tuple, "shape must be a tuple"
599        assert len(shape) == 2, "shape must be a tuple of length 2"
600        self._shape = (shape[1], shape[0])
601        self._writer = ocv_cv2.VideoWriter(
602            path, ocv_cv2.VideoWriter_fourcc(*"FFV1"), 1, shape, isColor=False
603        )
604        assert self._writer.isOpened(), f"Failed to open video writer at {path}"
605        self._i = 0
606
607    def write_detections(self, detections: Detections):
608        if len(detections) == 0:
609            return self._i
610
611        mask = detections.mask
612        assert (
613            mask.shape[1:] == self._shape
614        ), f"mask shape ({mask.shape[:1]}) must match the shape of the video ({self._shape})"
615        for i in range(mask.shape[0]):
616            frame_uint8 = detections.mask[i].astype(np.uint8)
617            self._writer.write(frame_uint8)
618            self._i += 1
619        return self._i
620
621    def release(self):
622        self._writer.release()
623
624
625def populate_mask(
626    detections: Detections, mask_stream: vf_cv2.VideoCapture, frame_idx: int
627):
628    assert type(detections) is Detections
629    assert detections.mask is None
630    detections.mask = []
631    assert len(detections) + frame_idx <= len(mask_stream)
632    for i in range(len(detections)):
633        mask = mask_stream[frame_idx + i]
634        assert mask.shape[2] == 1, "mask must be a single channel image"
635        detections.mask.append(mask)
CV2_FONT = 0
class BoxAnnotator:
32class BoxAnnotator:
33    def __init__(
34        self,
35        color=ColorPalette.DEFAULT,
36        thickness=2,
37        color_lookup=ColorLookup.CLASS,
38    ):
39        self.color = color
40        self.thickness = thickness
41        self.color_lookup = color_lookup
42
43    def annotate(
44        self,
45        scene: vf_cv2.Frame,
46        detections: Detections,
47        custom_color_lookup=None,
48    ):
49        for detection_idx in range(len(detections)):
50            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
51            color = resolve_color(
52                color=self.color,
53                detections=detections,
54                detection_idx=detection_idx,
55                color_lookup=(
56                    self.color_lookup
57                    if custom_color_lookup is None
58                    else custom_color_lookup
59                ),
60            )
61            vf_cv2.rectangle(
62                img=scene,
63                pt1=(x1, y1),
64                pt2=(x2, y2),
65                color=color.as_rgb(),
66                thickness=self.thickness,
67            )
68        return scene
BoxAnnotator( color=ColorPalette(colors=[Color(r=163, g=81, b=251), Color(r=255, g=64, b=64), Color(r=255, g=161, b=160), Color(r=255, g=118, b=51), Color(r=255, g=182, b=51), Color(r=209, g=212, b=53), Color(r=76, g=251, b=18), Color(r=148, g=207, b=26), Color(r=64, g=222, b=138), Color(r=27, g=150, b=64), Color(r=0, g=214, b=193), Color(r=46, g=156, b=170), Color(r=0, g=196, b=255), Color(r=54, g=71, b=151), Color(r=102, g=117, b=255), Color(r=0, g=25, b=239), Color(r=134, g=58, b=255), Color(r=83, g=0, b=135), Color(r=205, g=58, b=255), Color(r=255, g=151, b=202), Color(r=255, g=57, b=201)]), thickness=2, color_lookup=<ColorLookup.CLASS: 'class'>)
33    def __init__(
34        self,
35        color=ColorPalette.DEFAULT,
36        thickness=2,
37        color_lookup=ColorLookup.CLASS,
38    ):
39        self.color = color
40        self.thickness = thickness
41        self.color_lookup = color_lookup
color
thickness
color_lookup
def annotate( self, scene: vidformer.cv2.Frame, detections: supervision.detection.core.Detections, custom_color_lookup=None):
43    def annotate(
44        self,
45        scene: vf_cv2.Frame,
46        detections: Detections,
47        custom_color_lookup=None,
48    ):
49        for detection_idx in range(len(detections)):
50            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
51            color = resolve_color(
52                color=self.color,
53                detections=detections,
54                detection_idx=detection_idx,
55                color_lookup=(
56                    self.color_lookup
57                    if custom_color_lookup is None
58                    else custom_color_lookup
59                ),
60            )
61            vf_cv2.rectangle(
62                img=scene,
63                pt1=(x1, y1),
64                pt2=(x2, y2),
65                color=color.as_rgb(),
66                thickness=self.thickness,
67            )
68        return scene
class RoundBoxAnnotator:
 71class RoundBoxAnnotator:
 72    def __init__(
 73        self,
 74        color=ColorPalette.DEFAULT,
 75        thickness: int = 2,
 76        color_lookup: ColorLookup = ColorLookup.CLASS,
 77        roundness: float = 0.6,
 78    ):
 79        self.color = color
 80        self.thickness = thickness
 81        self.color_lookup = color_lookup
 82        if not 0 < roundness <= 1.0:
 83            raise ValueError("roundness attribute must be float between (0, 1.0]")
 84        self.roundness = roundness
 85
 86    def annotate(
 87        self,
 88        scene: vf_cv2.Frame,
 89        detections: _sv.Detections,
 90        custom_color_lookup=None,
 91    ):
 92        for detection_idx in range(len(detections)):
 93            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
 94            color = resolve_color(
 95                color=self.color,
 96                detections=detections,
 97                detection_idx=detection_idx,
 98                color_lookup=(
 99                    self.color_lookup
100                    if custom_color_lookup is None
101                    else custom_color_lookup
102                ),
103            )
104            radius = (
105                int((x2 - x1) // 2 * self.roundness)
106                if abs(x1 - x2) < abs(y1 - y2)
107                else int((y2 - y1) // 2 * self.roundness)
108            )
109            circle_coordinates = [
110                ((x1 + radius), (y1 + radius)),
111                ((x2 - radius), (y1 + radius)),
112                ((x2 - radius), (y2 - radius)),
113                ((x1 + radius), (y2 - radius)),
114            ]
115            line_coordinates = [
116                ((x1 + radius, y1), (x2 - radius, y1)),
117                ((x2, y1 + radius), (x2, y2 - radius)),
118                ((x1 + radius, y2), (x2 - radius, y2)),
119                ((x1, y1 + radius), (x1, y2 - radius)),
120            ]
121            start_angles = (180, 270, 0, 90)
122            end_angles = (270, 360, 90, 180)
123            for center_coordinates, line, start_angle, end_angle in zip(
124                circle_coordinates, line_coordinates, start_angles, end_angles
125            ):
126                vf_cv2.ellipse(
127                    img=scene,
128                    center=center_coordinates,
129                    axes=(radius, radius),
130                    angle=0,
131                    startAngle=start_angle,
132                    endAngle=end_angle,
133                    color=color.as_rgb(),
134                    thickness=self.thickness,
135                )
136                vf_cv2.line(
137                    img=scene,
138                    pt1=line[0],
139                    pt2=line[1],
140                    color=color.as_rgb(),
141                    thickness=self.thickness,
142                )
143        return scene
RoundBoxAnnotator( color=ColorPalette(colors=[Color(r=163, g=81, b=251), Color(r=255, g=64, b=64), Color(r=255, g=161, b=160), Color(r=255, g=118, b=51), Color(r=255, g=182, b=51), Color(r=209, g=212, b=53), Color(r=76, g=251, b=18), Color(r=148, g=207, b=26), Color(r=64, g=222, b=138), Color(r=27, g=150, b=64), Color(r=0, g=214, b=193), Color(r=46, g=156, b=170), Color(r=0, g=196, b=255), Color(r=54, g=71, b=151), Color(r=102, g=117, b=255), Color(r=0, g=25, b=239), Color(r=134, g=58, b=255), Color(r=83, g=0, b=135), Color(r=205, g=58, b=255), Color(r=255, g=151, b=202), Color(r=255, g=57, b=201)]), thickness: int = 2, color_lookup: supervision.annotators.utils.ColorLookup = <ColorLookup.CLASS: 'class'>, roundness: float = 0.6)
72    def __init__(
73        self,
74        color=ColorPalette.DEFAULT,
75        thickness: int = 2,
76        color_lookup: ColorLookup = ColorLookup.CLASS,
77        roundness: float = 0.6,
78    ):
79        self.color = color
80        self.thickness = thickness
81        self.color_lookup = color_lookup
82        if not 0 < roundness <= 1.0:
83            raise ValueError("roundness attribute must be float between (0, 1.0]")
84        self.roundness = roundness
color
thickness
color_lookup
roundness
def annotate( self, scene: vidformer.cv2.Frame, detections: supervision.detection.core.Detections, custom_color_lookup=None):
 86    def annotate(
 87        self,
 88        scene: vf_cv2.Frame,
 89        detections: _sv.Detections,
 90        custom_color_lookup=None,
 91    ):
 92        for detection_idx in range(len(detections)):
 93            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
 94            color = resolve_color(
 95                color=self.color,
 96                detections=detections,
 97                detection_idx=detection_idx,
 98                color_lookup=(
 99                    self.color_lookup
100                    if custom_color_lookup is None
101                    else custom_color_lookup
102                ),
103            )
104            radius = (
105                int((x2 - x1) // 2 * self.roundness)
106                if abs(x1 - x2) < abs(y1 - y2)
107                else int((y2 - y1) // 2 * self.roundness)
108            )
109            circle_coordinates = [
110                ((x1 + radius), (y1 + radius)),
111                ((x2 - radius), (y1 + radius)),
112                ((x2 - radius), (y2 - radius)),
113                ((x1 + radius), (y2 - radius)),
114            ]
115            line_coordinates = [
116                ((x1 + radius, y1), (x2 - radius, y1)),
117                ((x2, y1 + radius), (x2, y2 - radius)),
118                ((x1 + radius, y2), (x2 - radius, y2)),
119                ((x1, y1 + radius), (x1, y2 - radius)),
120            ]
121            start_angles = (180, 270, 0, 90)
122            end_angles = (270, 360, 90, 180)
123            for center_coordinates, line, start_angle, end_angle in zip(
124                circle_coordinates, line_coordinates, start_angles, end_angles
125            ):
126                vf_cv2.ellipse(
127                    img=scene,
128                    center=center_coordinates,
129                    axes=(radius, radius),
130                    angle=0,
131                    startAngle=start_angle,
132                    endAngle=end_angle,
133                    color=color.as_rgb(),
134                    thickness=self.thickness,
135                )
136                vf_cv2.line(
137                    img=scene,
138                    pt1=line[0],
139                    pt2=line[1],
140                    color=color.as_rgb(),
141                    thickness=self.thickness,
142                )
143        return scene
class BoxCornerAnnotator:
146class BoxCornerAnnotator:
147    def __init__(
148        self,
149        color=ColorPalette.DEFAULT,
150        thickness=4,
151        corner_length=15,
152        color_lookup=ColorLookup.CLASS,
153    ):
154        self.color = color
155        self.thickness: int = thickness
156        self.corner_length: int = corner_length
157        self.color_lookup: ColorLookup = color_lookup
158
159    def annotate(
160        self,
161        scene: vf_cv2.Frame,
162        detections: Detections,
163        custom_color_lookup=None,
164    ):
165        for detection_idx in range(len(detections)):
166            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
167            color = resolve_color(
168                color=self.color,
169                detections=detections,
170                detection_idx=detection_idx,
171                color_lookup=(
172                    self.color_lookup
173                    if custom_color_lookup is None
174                    else custom_color_lookup
175                ),
176            )
177            corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
178            for x, y in corners:
179                x_end = x + self.corner_length if x == x1 else x - self.corner_length
180                vf_cv2.line(
181                    scene, (x, y), (x_end, y), color.as_rgb(), thickness=self.thickness
182                )
183
184                y_end = y + self.corner_length if y == y1 else y - self.corner_length
185                vf_cv2.line(
186                    scene, (x, y), (x, y_end), color.as_rgb(), thickness=self.thickness
187                )
188        return scene
BoxCornerAnnotator( color=ColorPalette(colors=[Color(r=163, g=81, b=251), Color(r=255, g=64, b=64), Color(r=255, g=161, b=160), Color(r=255, g=118, b=51), Color(r=255, g=182, b=51), Color(r=209, g=212, b=53), Color(r=76, g=251, b=18), Color(r=148, g=207, b=26), Color(r=64, g=222, b=138), Color(r=27, g=150, b=64), Color(r=0, g=214, b=193), Color(r=46, g=156, b=170), Color(r=0, g=196, b=255), Color(r=54, g=71, b=151), Color(r=102, g=117, b=255), Color(r=0, g=25, b=239), Color(r=134, g=58, b=255), Color(r=83, g=0, b=135), Color(r=205, g=58, b=255), Color(r=255, g=151, b=202), Color(r=255, g=57, b=201)]), thickness=4, corner_length=15, color_lookup=<ColorLookup.CLASS: 'class'>)
147    def __init__(
148        self,
149        color=ColorPalette.DEFAULT,
150        thickness=4,
151        corner_length=15,
152        color_lookup=ColorLookup.CLASS,
153    ):
154        self.color = color
155        self.thickness: int = thickness
156        self.corner_length: int = corner_length
157        self.color_lookup: ColorLookup = color_lookup
color
thickness: int
corner_length: int
color_lookup: supervision.annotators.utils.ColorLookup
def annotate( self, scene: vidformer.cv2.Frame, detections: supervision.detection.core.Detections, custom_color_lookup=None):
159    def annotate(
160        self,
161        scene: vf_cv2.Frame,
162        detections: Detections,
163        custom_color_lookup=None,
164    ):
165        for detection_idx in range(len(detections)):
166            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
167            color = resolve_color(
168                color=self.color,
169                detections=detections,
170                detection_idx=detection_idx,
171                color_lookup=(
172                    self.color_lookup
173                    if custom_color_lookup is None
174                    else custom_color_lookup
175                ),
176            )
177            corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
178            for x, y in corners:
179                x_end = x + self.corner_length if x == x1 else x - self.corner_length
180                vf_cv2.line(
181                    scene, (x, y), (x_end, y), color.as_rgb(), thickness=self.thickness
182                )
183
184                y_end = y + self.corner_length if y == y1 else y - self.corner_length
185                vf_cv2.line(
186                    scene, (x, y), (x, y_end), color.as_rgb(), thickness=self.thickness
187                )
188        return scene
class ColorAnnotator:
191class ColorAnnotator:
192    def __init__(
193        self,
194        color=ColorPalette.DEFAULT,
195        opacity: float = 0.5,
196        color_lookup: ColorLookup = ColorLookup.CLASS,
197    ):
198        self.color = color
199        self.color_lookup: ColorLookup = color_lookup
200        self.opacity = opacity
201
202    def annotate(
203        self,
204        scene: vf_cv2.Frame,
205        detections: Detections,
206        custom_color_lookup=None,
207    ):
208        scene_with_boxes = scene.copy()
209        for detection_idx in range(len(detections)):
210            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
211            color = resolve_color(
212                color=self.color,
213                detections=detections,
214                detection_idx=detection_idx,
215                color_lookup=(
216                    self.color_lookup
217                    if custom_color_lookup is None
218                    else custom_color_lookup
219                ),
220            )
221            vf_cv2.rectangle(
222                img=scene_with_boxes,
223                pt1=(x1, y1),
224                pt2=(x2, y2),
225                color=color.as_rgb(),
226                thickness=-1,
227            )
228
229        vf_cv2.addWeighted(
230            scene_with_boxes, self.opacity, scene, 1 - self.opacity, gamma=0, dst=scene
231        )
232        return scene
ColorAnnotator( color=ColorPalette(colors=[Color(r=163, g=81, b=251), Color(r=255, g=64, b=64), Color(r=255, g=161, b=160), Color(r=255, g=118, b=51), Color(r=255, g=182, b=51), Color(r=209, g=212, b=53), Color(r=76, g=251, b=18), Color(r=148, g=207, b=26), Color(r=64, g=222, b=138), Color(r=27, g=150, b=64), Color(r=0, g=214, b=193), Color(r=46, g=156, b=170), Color(r=0, g=196, b=255), Color(r=54, g=71, b=151), Color(r=102, g=117, b=255), Color(r=0, g=25, b=239), Color(r=134, g=58, b=255), Color(r=83, g=0, b=135), Color(r=205, g=58, b=255), Color(r=255, g=151, b=202), Color(r=255, g=57, b=201)]), opacity: float = 0.5, color_lookup: supervision.annotators.utils.ColorLookup = <ColorLookup.CLASS: 'class'>)
192    def __init__(
193        self,
194        color=ColorPalette.DEFAULT,
195        opacity: float = 0.5,
196        color_lookup: ColorLookup = ColorLookup.CLASS,
197    ):
198        self.color = color
199        self.color_lookup: ColorLookup = color_lookup
200        self.opacity = opacity
color
color_lookup: supervision.annotators.utils.ColorLookup
opacity
def annotate( self, scene: vidformer.cv2.Frame, detections: supervision.detection.core.Detections, custom_color_lookup=None):
202    def annotate(
203        self,
204        scene: vf_cv2.Frame,
205        detections: Detections,
206        custom_color_lookup=None,
207    ):
208        scene_with_boxes = scene.copy()
209        for detection_idx in range(len(detections)):
210            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
211            color = resolve_color(
212                color=self.color,
213                detections=detections,
214                detection_idx=detection_idx,
215                color_lookup=(
216                    self.color_lookup
217                    if custom_color_lookup is None
218                    else custom_color_lookup
219                ),
220            )
221            vf_cv2.rectangle(
222                img=scene_with_boxes,
223                pt1=(x1, y1),
224                pt2=(x2, y2),
225                color=color.as_rgb(),
226                thickness=-1,
227            )
228
229        vf_cv2.addWeighted(
230            scene_with_boxes, self.opacity, scene, 1 - self.opacity, gamma=0, dst=scene
231        )
232        return scene
class CircleAnnotator:
235class CircleAnnotator:
236    def __init__(
237        self,
238        color=ColorPalette.DEFAULT,
239        thickness: int = 2,
240        color_lookup: ColorLookup = ColorLookup.CLASS,
241    ):
242        self.color = color
243        self.thickness: int = thickness
244        self.color_lookup: ColorLookup = color_lookup
245
246    def annotate(
247        self,
248        scene: vf_cv2.Frame,
249        detections: Detections,
250        custom_color_lookup=None,
251    ):
252        for detection_idx in range(len(detections)):
253            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
254            center = ((x1 + x2) // 2, (y1 + y2) // 2)
255            distance = sqrt((x1 - center[0]) ** 2 + (y1 - center[1]) ** 2)
256            color = resolve_color(
257                color=self.color,
258                detections=detections,
259                detection_idx=detection_idx,
260                color_lookup=(
261                    self.color_lookup
262                    if custom_color_lookup is None
263                    else custom_color_lookup
264                ),
265            )
266            vf_cv2.circle(
267                img=scene,
268                center=center,
269                radius=int(distance),
270                color=color.as_rgb(),
271                thickness=self.thickness,
272            )
273
274        return scene
CircleAnnotator( color=ColorPalette(colors=[Color(r=163, g=81, b=251), Color(r=255, g=64, b=64), Color(r=255, g=161, b=160), Color(r=255, g=118, b=51), Color(r=255, g=182, b=51), Color(r=209, g=212, b=53), Color(r=76, g=251, b=18), Color(r=148, g=207, b=26), Color(r=64, g=222, b=138), Color(r=27, g=150, b=64), Color(r=0, g=214, b=193), Color(r=46, g=156, b=170), Color(r=0, g=196, b=255), Color(r=54, g=71, b=151), Color(r=102, g=117, b=255), Color(r=0, g=25, b=239), Color(r=134, g=58, b=255), Color(r=83, g=0, b=135), Color(r=205, g=58, b=255), Color(r=255, g=151, b=202), Color(r=255, g=57, b=201)]), thickness: int = 2, color_lookup: supervision.annotators.utils.ColorLookup = <ColorLookup.CLASS: 'class'>)
236    def __init__(
237        self,
238        color=ColorPalette.DEFAULT,
239        thickness: int = 2,
240        color_lookup: ColorLookup = ColorLookup.CLASS,
241    ):
242        self.color = color
243        self.thickness: int = thickness
244        self.color_lookup: ColorLookup = color_lookup
color
thickness: int
color_lookup: supervision.annotators.utils.ColorLookup
def annotate( self, scene: vidformer.cv2.Frame, detections: supervision.detection.core.Detections, custom_color_lookup=None):
246    def annotate(
247        self,
248        scene: vf_cv2.Frame,
249        detections: Detections,
250        custom_color_lookup=None,
251    ):
252        for detection_idx in range(len(detections)):
253            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
254            center = ((x1 + x2) // 2, (y1 + y2) // 2)
255            distance = sqrt((x1 - center[0]) ** 2 + (y1 - center[1]) ** 2)
256            color = resolve_color(
257                color=self.color,
258                detections=detections,
259                detection_idx=detection_idx,
260                color_lookup=(
261                    self.color_lookup
262                    if custom_color_lookup is None
263                    else custom_color_lookup
264                ),
265            )
266            vf_cv2.circle(
267                img=scene,
268                center=center,
269                radius=int(distance),
270                color=color.as_rgb(),
271                thickness=self.thickness,
272            )
273
274        return scene
class DotAnnotator:
277class DotAnnotator:
278    def __init__(
279        self,
280        color=ColorPalette.DEFAULT,
281        radius: int = 4,
282        position: Position = Position.CENTER,
283        color_lookup: ColorLookup = ColorLookup.CLASS,
284        outline_thickness: int = 0,
285        outline_color=Color.BLACK,
286    ):
287        self.color = color
288        self.radius: int = radius
289        self.position: Position = position
290        self.color_lookup: ColorLookup = color_lookup
291        self.outline_thickness = outline_thickness
292        self.outline_color = outline_color
293
294    def annotate(
295        self,
296        scene: vf_cv2.Frame,
297        detections: Detections,
298        custom_color_lookup=None,
299    ):
300        xy = detections.get_anchors_coordinates(anchor=self.position)
301        for detection_idx in range(len(detections)):
302            color = resolve_color(
303                color=self.color,
304                detections=detections,
305                detection_idx=detection_idx,
306                color_lookup=(
307                    self.color_lookup
308                    if custom_color_lookup is None
309                    else custom_color_lookup
310                ),
311            )
312            center = (int(xy[detection_idx, 0]), int(xy[detection_idx, 1]))
313
314            vf_cv2.circle(scene, center, self.radius, color.as_rgb(), -1)
315            if self.outline_thickness:
316                outline_color = resolve_color(
317                    color=self.outline_color,
318                    detections=detections,
319                    detection_idx=detection_idx,
320                    color_lookup=(
321                        self.color_lookup
322                        if custom_color_lookup is None
323                        else custom_color_lookup
324                    ),
325                )
326                vf_cv2.circle(
327                    scene,
328                    center,
329                    self.radius,
330                    outline_color.as_rgb(),
331                    self.outline_thickness,
332                )
333        return scene
DotAnnotator( color=ColorPalette(colors=[Color(r=163, g=81, b=251), Color(r=255, g=64, b=64), Color(r=255, g=161, b=160), Color(r=255, g=118, b=51), Color(r=255, g=182, b=51), Color(r=209, g=212, b=53), Color(r=76, g=251, b=18), Color(r=148, g=207, b=26), Color(r=64, g=222, b=138), Color(r=27, g=150, b=64), Color(r=0, g=214, b=193), Color(r=46, g=156, b=170), Color(r=0, g=196, b=255), Color(r=54, g=71, b=151), Color(r=102, g=117, b=255), Color(r=0, g=25, b=239), Color(r=134, g=58, b=255), Color(r=83, g=0, b=135), Color(r=205, g=58, b=255), Color(r=255, g=151, b=202), Color(r=255, g=57, b=201)]), radius: int = 4, position: supervision.geometry.core.Position = <Position.CENTER: 'CENTER'>, color_lookup: supervision.annotators.utils.ColorLookup = <ColorLookup.CLASS: 'class'>, outline_thickness: int = 0, outline_color=Color(r=0, g=0, b=0))
278    def __init__(
279        self,
280        color=ColorPalette.DEFAULT,
281        radius: int = 4,
282        position: Position = Position.CENTER,
283        color_lookup: ColorLookup = ColorLookup.CLASS,
284        outline_thickness: int = 0,
285        outline_color=Color.BLACK,
286    ):
287        self.color = color
288        self.radius: int = radius
289        self.position: Position = position
290        self.color_lookup: ColorLookup = color_lookup
291        self.outline_thickness = outline_thickness
292        self.outline_color = outline_color
color
radius: int
position: supervision.geometry.core.Position
color_lookup: supervision.annotators.utils.ColorLookup
outline_thickness
outline_color
def annotate( self, scene: vidformer.cv2.Frame, detections: supervision.detection.core.Detections, custom_color_lookup=None):
294    def annotate(
295        self,
296        scene: vf_cv2.Frame,
297        detections: Detections,
298        custom_color_lookup=None,
299    ):
300        xy = detections.get_anchors_coordinates(anchor=self.position)
301        for detection_idx in range(len(detections)):
302            color = resolve_color(
303                color=self.color,
304                detections=detections,
305                detection_idx=detection_idx,
306                color_lookup=(
307                    self.color_lookup
308                    if custom_color_lookup is None
309                    else custom_color_lookup
310                ),
311            )
312            center = (int(xy[detection_idx, 0]), int(xy[detection_idx, 1]))
313
314            vf_cv2.circle(scene, center, self.radius, color.as_rgb(), -1)
315            if self.outline_thickness:
316                outline_color = resolve_color(
317                    color=self.outline_color,
318                    detections=detections,
319                    detection_idx=detection_idx,
320                    color_lookup=(
321                        self.color_lookup
322                        if custom_color_lookup is None
323                        else custom_color_lookup
324                    ),
325                )
326                vf_cv2.circle(
327                    scene,
328                    center,
329                    self.radius,
330                    outline_color.as_rgb(),
331                    self.outline_thickness,
332                )
333        return scene
class LabelAnnotator:
336class LabelAnnotator:
337    def __init__(
338        self,
339        color=ColorPalette.DEFAULT,
340        text_color=Color.WHITE,
341        text_scale: float = 0.5,
342        text_thickness: int = 1,
343        text_padding: int = 10,
344        text_position: Position = Position.TOP_LEFT,
345        color_lookup: ColorLookup = ColorLookup.CLASS,
346        border_radius: int = 0,
347        smart_position: bool = False,
348    ):
349        self.border_radius: int = border_radius
350        self.color = color
351        self.text_color = text_color
352        self.text_scale: float = text_scale
353        self.text_thickness: int = text_thickness
354        self.text_padding: int = text_padding
355        self.text_anchor: Position = text_position
356        self.color_lookup: ColorLookup = color_lookup
357        self.smart_position = smart_position
358
359    def annotate(
360        self,
361        scene,
362        detections: Detections,
363        labels,
364        custom_color_lookup=None,
365    ):
366        self._validate_labels(labels, detections)
367
368        labels = self._get_labels_text(detections, labels)
369        label_properties = self._get_label_properties(detections, labels)
370
371        if self.smart_position:
372            xyxy = label_properties[:, :4]
373            xyxy = spread_out_boxes(xyxy)
374            label_properties[:, :4] = xyxy
375
376        self._draw_labels(
377            scene=scene,
378            labels=labels,
379            label_properties=label_properties,
380            detections=detections,
381            custom_color_lookup=custom_color_lookup,
382        )
383
384        return scene
385
386    def _validate_labels(self, labels, detections: Detections):
387        if labels is not None and len(labels) != len(detections):
388            raise ValueError(
389                f"The number of labels ({len(labels)}) does not match the "
390                f"number of detections ({len(detections)}). Each detection "
391                f"should have exactly 1 label."
392            )
393
394    def _get_label_properties(
395        self,
396        detections: Detections,
397        labels,
398    ):
399        label_properties = []
400        anchors_coordinates = detections.get_anchors_coordinates(
401            anchor=self.text_anchor
402        ).astype(int)
403
404        for label, center_coords in zip(labels, anchors_coordinates):
405            (text_w, text_h) = vf_cv2.getTextSize(
406                text=label,
407                fontFace=CV2_FONT,
408                fontScale=self.text_scale,
409                thickness=self.text_thickness,
410            )[0]
411
412            width_padded = text_w + 2 * self.text_padding
413            height_padded = text_h + 2 * self.text_padding
414
415            text_background_xyxy = resolve_text_background_xyxy(
416                center_coordinates=tuple(center_coords),
417                text_wh=(width_padded, height_padded),
418                position=self.text_anchor,
419            )
420
421            label_properties.append(
422                [
423                    *text_background_xyxy,
424                    text_h,
425                ]
426            )
427
428        return np.array(label_properties).reshape(-1, 5)
429
430    @staticmethod
431    def _get_labels_text(detections: Detections, custom_labels):
432        if custom_labels is not None:
433            return custom_labels
434
435        labels = []
436        for idx in range(len(detections)):
437            if CLASS_NAME_DATA_FIELD in detections.data:
438                labels.append(detections.data[CLASS_NAME_DATA_FIELD][idx])
439            elif detections.class_id is not None:
440                labels.append(str(detections.class_id[idx]))
441            else:
442                labels.append(str(idx))
443        return labels
444
445    def _draw_labels(
446        self,
447        scene,
448        labels,
449        label_properties,
450        detections,
451        custom_color_lookup,
452    ) -> None:
453        assert len(labels) == len(label_properties) == len(detections), (
454            f"Number of label properties ({len(label_properties)}), "
455            f"labels ({len(labels)}) and detections ({len(detections)}) "
456            "do not match."
457        )
458
459        color_lookup = (
460            custom_color_lookup
461            if custom_color_lookup is not None
462            else self.color_lookup
463        )
464
465        for idx, label_property in enumerate(label_properties):
466            background_color = resolve_color(
467                color=self.color,
468                detections=detections,
469                detection_idx=idx,
470                color_lookup=color_lookup,
471            )
472            text_color = resolve_color(
473                color=self.text_color,
474                detections=detections,
475                detection_idx=idx,
476                color_lookup=color_lookup,
477            )
478
479            box_xyxy = label_property[:4]
480            text_height_padded = label_property[4]
481            self.draw_rounded_rectangle(
482                scene=scene,
483                xyxy=box_xyxy,
484                color=background_color.as_rgb(),
485                border_radius=self.border_radius,
486            )
487
488            text_x = box_xyxy[0] + self.text_padding
489            text_y = box_xyxy[1] + self.text_padding + text_height_padded
490            vf_cv2.putText(
491                img=scene,
492                text=labels[idx],
493                org=(text_x, text_y),
494                fontFace=CV2_FONT,
495                fontScale=self.text_scale,
496                color=text_color.as_rgb(),
497                thickness=self.text_thickness,
498                lineType=vf_cv2.LINE_AA,
499            )
500
501    @staticmethod
502    def draw_rounded_rectangle(
503        scene: np.ndarray,
504        xyxy,
505        color,
506        border_radius: int,
507    ) -> np.ndarray:
508        x1, y1, x2, y2 = xyxy
509        width = x2 - x1
510        height = y2 - y1
511
512        border_radius = min(border_radius, min(width, height) // 2)
513
514        if border_radius <= 0:
515            vf_cv2.rectangle(
516                img=scene,
517                pt1=(x1, y1),
518                pt2=(x2, y2),
519                color=color,
520                thickness=-1,
521            )
522        else:
523            rectangle_coordinates = [
524                ((x1 + border_radius, y1), (x2 - border_radius, y2)),
525                ((x1, y1 + border_radius), (x2, y2 - border_radius)),
526            ]
527            circle_centers = [
528                (x1 + border_radius, y1 + border_radius),
529                (x2 - border_radius, y1 + border_radius),
530                (x1 + border_radius, y2 - border_radius),
531                (x2 - border_radius, y2 - border_radius),
532            ]
533
534            for coordinates in rectangle_coordinates:
535                vf_cv2.rectangle(
536                    img=scene,
537                    pt1=coordinates[0],
538                    pt2=coordinates[1],
539                    color=color,
540                    thickness=-1,
541                )
542            for center in circle_centers:
543                vf_cv2.circle(
544                    img=scene,
545                    center=center,
546                    radius=border_radius,
547                    color=color,
548                    thickness=-1,
549                )
550        return scene
LabelAnnotator( color=ColorPalette(colors=[Color(r=163, g=81, b=251), Color(r=255, g=64, b=64), Color(r=255, g=161, b=160), Color(r=255, g=118, b=51), Color(r=255, g=182, b=51), Color(r=209, g=212, b=53), Color(r=76, g=251, b=18), Color(r=148, g=207, b=26), Color(r=64, g=222, b=138), Color(r=27, g=150, b=64), Color(r=0, g=214, b=193), Color(r=46, g=156, b=170), Color(r=0, g=196, b=255), Color(r=54, g=71, b=151), Color(r=102, g=117, b=255), Color(r=0, g=25, b=239), Color(r=134, g=58, b=255), Color(r=83, g=0, b=135), Color(r=205, g=58, b=255), Color(r=255, g=151, b=202), Color(r=255, g=57, b=201)]), text_color=Color(r=255, g=255, b=255), text_scale: float = 0.5, text_thickness: int = 1, text_padding: int = 10, text_position: supervision.geometry.core.Position = <Position.TOP_LEFT: 'TOP_LEFT'>, color_lookup: supervision.annotators.utils.ColorLookup = <ColorLookup.CLASS: 'class'>, border_radius: int = 0, smart_position: bool = False)
337    def __init__(
338        self,
339        color=ColorPalette.DEFAULT,
340        text_color=Color.WHITE,
341        text_scale: float = 0.5,
342        text_thickness: int = 1,
343        text_padding: int = 10,
344        text_position: Position = Position.TOP_LEFT,
345        color_lookup: ColorLookup = ColorLookup.CLASS,
346        border_radius: int = 0,
347        smart_position: bool = False,
348    ):
349        self.border_radius: int = border_radius
350        self.color = color
351        self.text_color = text_color
352        self.text_scale: float = text_scale
353        self.text_thickness: int = text_thickness
354        self.text_padding: int = text_padding
355        self.text_anchor: Position = text_position
356        self.color_lookup: ColorLookup = color_lookup
357        self.smart_position = smart_position
border_radius: int
color
text_color
text_scale: float
text_thickness: int
text_padding: int
text_anchor: supervision.geometry.core.Position
color_lookup: supervision.annotators.utils.ColorLookup
smart_position
def annotate( self, scene, detections: supervision.detection.core.Detections, labels, custom_color_lookup=None):
359    def annotate(
360        self,
361        scene,
362        detections: Detections,
363        labels,
364        custom_color_lookup=None,
365    ):
366        self._validate_labels(labels, detections)
367
368        labels = self._get_labels_text(detections, labels)
369        label_properties = self._get_label_properties(detections, labels)
370
371        if self.smart_position:
372            xyxy = label_properties[:, :4]
373            xyxy = spread_out_boxes(xyxy)
374            label_properties[:, :4] = xyxy
375
376        self._draw_labels(
377            scene=scene,
378            labels=labels,
379            label_properties=label_properties,
380            detections=detections,
381            custom_color_lookup=custom_color_lookup,
382        )
383
384        return scene
@staticmethod
def draw_rounded_rectangle(scene: numpy.ndarray, xyxy, color, border_radius: int) -> numpy.ndarray:
501    @staticmethod
502    def draw_rounded_rectangle(
503        scene: np.ndarray,
504        xyxy,
505        color,
506        border_radius: int,
507    ) -> np.ndarray:
508        x1, y1, x2, y2 = xyxy
509        width = x2 - x1
510        height = y2 - y1
511
512        border_radius = min(border_radius, min(width, height) // 2)
513
514        if border_radius <= 0:
515            vf_cv2.rectangle(
516                img=scene,
517                pt1=(x1, y1),
518                pt2=(x2, y2),
519                color=color,
520                thickness=-1,
521            )
522        else:
523            rectangle_coordinates = [
524                ((x1 + border_radius, y1), (x2 - border_radius, y2)),
525                ((x1, y1 + border_radius), (x2, y2 - border_radius)),
526            ]
527            circle_centers = [
528                (x1 + border_radius, y1 + border_radius),
529                (x2 - border_radius, y1 + border_radius),
530                (x1 + border_radius, y2 - border_radius),
531                (x2 - border_radius, y2 - border_radius),
532            ]
533
534            for coordinates in rectangle_coordinates:
535                vf_cv2.rectangle(
536                    img=scene,
537                    pt1=coordinates[0],
538                    pt2=coordinates[1],
539                    color=color,
540                    thickness=-1,
541                )
542            for center in circle_centers:
543                vf_cv2.circle(
544                    img=scene,
545                    center=center,
546                    radius=border_radius,
547                    color=color,
548                    thickness=-1,
549                )
550        return scene
class MaskAnnotator:
553class MaskAnnotator:
554    def __init__(
555        self,
556        color=ColorPalette.DEFAULT,
557        opacity: float = 0.5,
558        color_lookup: ColorLookup = ColorLookup.CLASS,
559    ):
560        self.color = color
561        self.opacity = opacity
562        self.color_lookup: ColorLookup = color_lookup
563
564    def annotate(
565        self,
566        scene,
567        detections: Detections,
568        custom_color_lookup=None,
569    ):
570        if detections.mask is None:
571            return scene
572
573        colored_mask = scene.copy()
574
575        for detection_idx in np.flip(np.argsort(detections.box_area)):
576            color = resolve_color(
577                color=self.color,
578                detections=detections,
579                detection_idx=detection_idx,
580                color_lookup=(
581                    self.color_lookup
582                    if custom_color_lookup is None
583                    else custom_color_lookup
584                ),
585            )
586            mask = detections.mask[detection_idx]
587            colored_mask[mask] = color.as_bgr()
588
589        vf_cv2.addWeighted(
590            colored_mask, self.opacity, scene, 1 - self.opacity, 0, dst=scene
591        )
592        return scene
MaskAnnotator( color=ColorPalette(colors=[Color(r=163, g=81, b=251), Color(r=255, g=64, b=64), Color(r=255, g=161, b=160), Color(r=255, g=118, b=51), Color(r=255, g=182, b=51), Color(r=209, g=212, b=53), Color(r=76, g=251, b=18), Color(r=148, g=207, b=26), Color(r=64, g=222, b=138), Color(r=27, g=150, b=64), Color(r=0, g=214, b=193), Color(r=46, g=156, b=170), Color(r=0, g=196, b=255), Color(r=54, g=71, b=151), Color(r=102, g=117, b=255), Color(r=0, g=25, b=239), Color(r=134, g=58, b=255), Color(r=83, g=0, b=135), Color(r=205, g=58, b=255), Color(r=255, g=151, b=202), Color(r=255, g=57, b=201)]), opacity: float = 0.5, color_lookup: supervision.annotators.utils.ColorLookup = <ColorLookup.CLASS: 'class'>)
554    def __init__(
555        self,
556        color=ColorPalette.DEFAULT,
557        opacity: float = 0.5,
558        color_lookup: ColorLookup = ColorLookup.CLASS,
559    ):
560        self.color = color
561        self.opacity = opacity
562        self.color_lookup: ColorLookup = color_lookup
color
opacity
color_lookup: supervision.annotators.utils.ColorLookup
def annotate( self, scene, detections: supervision.detection.core.Detections, custom_color_lookup=None):
564    def annotate(
565        self,
566        scene,
567        detections: Detections,
568        custom_color_lookup=None,
569    ):
570        if detections.mask is None:
571            return scene
572
573        colored_mask = scene.copy()
574
575        for detection_idx in np.flip(np.argsort(detections.box_area)):
576            color = resolve_color(
577                color=self.color,
578                detections=detections,
579                detection_idx=detection_idx,
580                color_lookup=(
581                    self.color_lookup
582                    if custom_color_lookup is None
583                    else custom_color_lookup
584                ),
585            )
586            mask = detections.mask[detection_idx]
587            colored_mask[mask] = color.as_bgr()
588
589        vf_cv2.addWeighted(
590            colored_mask, self.opacity, scene, 1 - self.opacity, 0, dst=scene
591        )
592        return scene
class MaskStreamWriter:
595class MaskStreamWriter:
596    def __init__(self, path: str, shape: tuple):
597        # Shape should be (width, height)
598        assert ocv_cv2 is not None, "OpenCV cv2 is required for ExternDetectionsBuilder"
599        assert type(shape) is tuple, "shape must be a tuple"
600        assert len(shape) == 2, "shape must be a tuple of length 2"
601        self._shape = (shape[1], shape[0])
602        self._writer = ocv_cv2.VideoWriter(
603            path, ocv_cv2.VideoWriter_fourcc(*"FFV1"), 1, shape, isColor=False
604        )
605        assert self._writer.isOpened(), f"Failed to open video writer at {path}"
606        self._i = 0
607
608    def write_detections(self, detections: Detections):
609        if len(detections) == 0:
610            return self._i
611
612        mask = detections.mask
613        assert (
614            mask.shape[1:] == self._shape
615        ), f"mask shape ({mask.shape[:1]}) must match the shape of the video ({self._shape})"
616        for i in range(mask.shape[0]):
617            frame_uint8 = detections.mask[i].astype(np.uint8)
618            self._writer.write(frame_uint8)
619            self._i += 1
620        return self._i
621
622    def release(self):
623        self._writer.release()
MaskStreamWriter(path: str, shape: tuple)
596    def __init__(self, path: str, shape: tuple):
597        # Shape should be (width, height)
598        assert ocv_cv2 is not None, "OpenCV cv2 is required for ExternDetectionsBuilder"
599        assert type(shape) is tuple, "shape must be a tuple"
600        assert len(shape) == 2, "shape must be a tuple of length 2"
601        self._shape = (shape[1], shape[0])
602        self._writer = ocv_cv2.VideoWriter(
603            path, ocv_cv2.VideoWriter_fourcc(*"FFV1"), 1, shape, isColor=False
604        )
605        assert self._writer.isOpened(), f"Failed to open video writer at {path}"
606        self._i = 0
def write_detections(self, detections: supervision.detection.core.Detections):
608    def write_detections(self, detections: Detections):
609        if len(detections) == 0:
610            return self._i
611
612        mask = detections.mask
613        assert (
614            mask.shape[1:] == self._shape
615        ), f"mask shape ({mask.shape[:1]}) must match the shape of the video ({self._shape})"
616        for i in range(mask.shape[0]):
617            frame_uint8 = detections.mask[i].astype(np.uint8)
618            self._writer.write(frame_uint8)
619            self._i += 1
620        return self._i
def release(self):
622    def release(self):
623        self._writer.release()
def populate_mask( detections: supervision.detection.core.Detections, mask_stream: vidformer.cv2.VideoCapture, frame_idx: int):
626def populate_mask(
627    detections: Detections, mask_stream: vf_cv2.VideoCapture, frame_idx: int
628):
629    assert type(detections) is Detections
630    assert detections.mask is None
631    detections.mask = []
632    assert len(detections) + frame_idx <= len(mask_stream)
633    for i in range(len(detections)):
634        mask = mask_stream[frame_idx + i]
635        assert mask.shape[2] == 1, "mask must be a single channel image"
636        detections.mask.append(mask)