Spaces:
Running
Running
File size: 3,915 Bytes
3ed4749 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import os
import cv2
import numpy as np
import supervision as sv
from PIL import Image, ImageDraw, ImageFont
from supervision.annotators.utils import resolve_color
# visualization toos based on supervision
BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator(thickness=2)
class LabelAnnotator(sv.LabelAnnotator):
@staticmethod
def resolve_text_background_xyxy(
center_coordinates,
text_wh,
position,
):
center_x, center_y = center_coordinates
text_w, text_h = text_wh
return center_x, center_y, center_x + text_w, center_y + text_h
def _draw_labels(
self,
scene: np.ndarray,
labels: list[str],
label_properties: np.ndarray,
detections,
custom_color_lookup,
) -> None:
assert len(labels) == len(label_properties) == len(detections), (
f"Number of label properties ({len(label_properties)}), "
f"labels ({len(labels)}) and detections ({len(detections)}) "
"do not match."
)
color_lookup = (
custom_color_lookup
if custom_color_lookup is not None
else self.color_lookup
)
font = ImageFont.truetype("simhei.ttf", int(30 * self.text_scale))
for idx, label_property in enumerate(label_properties):
background_color = resolve_color(
color=self.color,
detections=detections,
detection_idx=idx,
color_lookup=color_lookup,
)
text_color = resolve_color(
color=self.text_color,
detections=detections,
detection_idx=idx,
color_lookup=color_lookup,
)
box_xyxy = label_property[:4]
text_height_padded = label_property[4]
self.draw_rounded_rectangle(
scene=scene,
xyxy=box_xyxy,
color=background_color.as_bgr(),
border_radius=self.border_radius,
)
text_x = box_xyxy[0] + self.text_padding
text_y = box_xyxy[1]
scene_pil = Image.fromarray(cv2.cvtColor(scene, cv2.COLOR_BGR2RGB))
draw = ImageDraw.Draw(scene_pil)
draw.text(
(text_x, text_y),
labels[idx],
font=font,
fill=(text_color.r, text_color.g, text_color.b),
)
scene[:] = cv2.cvtColor(np.array(scene_pil), cv2.COLOR_RGB2BGR)
LABEL_ANNOTATOR = LabelAnnotator(text_padding=4,
text_scale=0.5,
text_thickness=1)
POINT_ANNOTATOR = sv.DotAnnotator(radius=6)
def draw_boxes_points_with_labels(
cv2_image,
boxes=None,
points=None,
classes=None,
output_path=None,
):
annotated_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB)
if boxes is not None and boxes.size:
detections = sv.Detections(
xyxy=boxes,
class_id=np.arange(len(boxes)),
confidence=np.ones(len(boxes))
)
annotated_image = BOUNDING_BOX_ANNOTATOR.annotate(
annotated_image, detections)
if points is not None and points.size:
points = np.concatenate([points, points], axis=1)
detections = sv.Detections(
xyxy=points,
class_id=np.arange(len(points)),
confidence=np.ones(len(points))
)
annotated_image = POINT_ANNOTATOR.annotate(
annotated_image, detections,
)
if classes:
annotated_image = LABEL_ANNOTATOR.annotate(
annotated_image, detections, labels=classes
)
if output_path:
cv2.imwrite(
output_path,
cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR)
)
return annotated_image |