vidformer.supervision
vidformer.supervision is the supervision frontend for vidformer.
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'>)
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
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'>)
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'>)
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'>)
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
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
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'>)
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
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)