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