|
|
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 |
|
|
|
|
|
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 |
