Spaces:
Sleeping
Sleeping
File size: 2,729 Bytes
a31bf96 8d7cbb4 6654dec 8d7cbb4 a31bf96 bdcbf84 | 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 | import supervision as sv
import cv2
import os
def annotate_image(image, detections, class_names, trace_annotator):
"""
Draws bounding boxes, class labels with track IDs, and trace lines on the image.
Args:
image (np.ndarray): The image to annotate.
detections (sv.Detections): Detection results.
class_names (dict): Class ID to name mapping.
trace_annotator (sv.TraceAnnotator): Annotator for trace lines.
Returns:
np.ndarray: Annotated image.
"""
if len(detections) > 0:
# Initialize annotators, trace annotator should be initialized only once
# Draw box
_, width = image.shape[:2]
text_size = 2.0*width/1024
box_thickness = int(6 * width / 1024) # scale box thickness
text_thickness = max(1, int(6 * width / 1024)) # scale text thickness, min 1
box_annotator = sv.BoxAnnotator(thickness=box_thickness)
label_annotator = sv.LabelAnnotator(text_scale=text_size, text_thickness=text_thickness)
# Build labels with track IDs
labels = []
for i in range(len(detections)):
class_id = detections.class_id[i]
conf = detections.confidence[i]
track_id = detections.tracker_id[i] if detections.tracker_id is not None else None
if track_id is not None:
label = f"ID: {track_id} {class_names[class_id]} {conf:.2f}"
else:
label = f"{class_names[class_id]} {conf:.2f}"
labels.append(label)
# Apply all annotators
image = box_annotator.annotate(scene=image, detections=detections)
image = label_annotator.annotate(scene=image, detections=detections, labels=labels)
image = trace_annotator.annotate(scene=image, detections=detections)
return image
def process_video_frame(frame, detections, class_names, trace_annotator, plot_dets=True, show=False, video_writer=None):
"""
Handles each frame of video for annotation, display, and saving.
Args:
frame (np.ndarray): Video frame.
detections (sv.Detections): Detection results.
class_names (dict): Class ID to name mapping.
trace_annotator (sv.TraceAnnotator): Annotator for trace lines.
plot_dets (bool): Save annotated video frame.
show (bool): Display annotated frame in a window.
"""
annotated_frame = annotate_image(frame, detections, class_names, trace_annotator)
if show:
h, w = frame.shape[:2]
scale = min(750 / w, 750 / h)
resized = cv2.resize(annotated_frame, (int(w * scale), int(h * scale)))
cv2.imshow("Detection", resized)
cv2.waitKey(1)
return annotated_frame
|