参考:
https://supervision.roboflow.com/latest/
https://github.com/roboflow/supervision/tree/develop/examples
版本:
pip install -U supervision
ultralytics-8.1.35 (大于8.1才行,不然可能会有错误AttributeError: ‘Results’ object has no attribute ‘obb’ )
supervision 0.16.0
简单案例:
import cv2
import supervision as sv
from ultralytics import YOLO
image = cv2.imread(...)
model = YOLO('yolov8s.pt')
result = model(image)[0]
detections = sv.Detections.from_ultralytics(result)
len(detections)
跟踪案例
https://github.com/roboflow/supervision/tree/develop/examples/heatmap_and_track
运行结果:
python D:\opencv2\supervision_cv\test.py --source_weights_path "C:\Users\loong\Downloads\yolov8m (1).pt" --source_video_path "C:\Users\loong\Downloads\istockphoto-1047817112-640_adpp_is.mp4" --confidence_threshold 0.3 --iou_threshold 0.5 --target_video_path output_video.mp4
具体代码:
import argparse
import cv2
from ultralytics import YOLO
import supervision as sv
from supervision.assets import VideoAssets, download_assets
def download_video() -> str:
download_assets(VideoAssets.PEOPLE_WALKING)
return VideoAssets.PEOPLE_WALKING.value
def heatmap_and_track(
source_weights_path: str,
source_video_path: str,
target_video_path: str,
confidence_threshold: float = 0.35,
iou_threshold: float = 0.5,
heatmap_alpha: float = 0.5,
radius: int = 25,
track_threshold: float = 0.35,
track_seconds: int = 5,
match_threshold: float = 0.99,
) -> None:
### instantiate model
model = YOLO(source_weights_path)
### heatmap config
heat_map_annotator = sv.HeatMapAnnotator(
position=sv.Position.BOTTOM_CENTER,
opacity=heatmap_alpha,
radius=radius,
kernel_size=25,
top_hue=0,
low_hue=125,
)
### annotation config
label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
### get the video fps
cap = cv2.VideoCapture(source_video_path)
fps = int(cap.get(cv2.CAP_PROP_FPS))
cap.release()
### tracker config
byte_tracker = sv.ByteTrack(
track_thresh=track_threshold,
track_buffer=track_seconds * fps,
match_thresh=match_threshold,
frame_rate=fps,
)
### video config
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
frames_generator = sv.get_video_frames_generator(
source_path=source_video_path, stride=1
)
### Detect, track, annotate, save
with sv.VideoSink(target_path=target_video_path, video_info=video_info) as sink:
for frame in frames_generator:
result = model(
source=frame,
classes=[0], # only person class
conf=confidence_threshold,
iou=iou_threshold,
# show_conf = True,
# save_txt = True,
# save_conf = True,
# save = True,
device=None, # use None = CPU, 0 = single GPU, or [0,1] = dual GPU
# agnostic_nms=True
)[0]
detections = sv.Detections.from_ultralytics(result) # get detections
detections = byte_tracker.update_with_detections(
detections
) # update tracker
### draw heatmap
annotated_frame = heat_map_annotator.annotate(
scene=frame.copy(), detections=detections
)
### draw other attributes from `detections` object
labels = [
f"#{tracker_id}"
for class_id, tracker_id in zip(
detections.class_id, detections.tracker_id
)
]
label_annotator.annotate(
scene=annotated_frame, detections=detections, labels=labels
)
sink.write_frame(frame=annotated_frame)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Heatmap and Tracking with Supervision"
)
parser.add_argument(
"--source_weights_path",
required=True,
help="Path to the source weights file",
type=str,
)
parser.add_argument(
"--source_video_path",
default=download_video(),
help="Path to the source video file",
type=str,
)
parser.add_argument(
"--target_video_path",
default="output.mp4",
help="Path to the target video file (output)",
type=str,
)
parser.add_argument(
"--confidence_threshold",
default=0.35,
help="Confidence threshold for the model",
type=float,
)
parser.add_argument(
"--iou_threshold",
default=0.5,
help="IOU threshold for the model",
type=float,
)
parser.add_argument(
"--heatmap_alpha",
default=0.5,
help="Opacity of the overlay mask, between 0 and 1",
type=float,
)
parser.add_argument(
"--radius",
default=25,
help="Radius of the heat circle",
type=float,
)
parser.add_argument(
"--track_threshold",
default=0.35,
help="Detection confidence threshold for track activation",
type=float,
)
parser.add_argument(
"--track_seconds",
default=5,
help="Number of seconds to buffer when a track is lost",
type=int,
)
parser.add_argument(
"--match_threshold",
default=0.99,
help="Threshold for matching tracks with detections",
type=float,
)
args = parser.parse_args()
heatmap_and_track(
source_weights_path=args.source_weights_path,
source_video_path=args.source_video_path,
target_video_path=args.target_video_path,
confidence_threshold=args.confidence_threshold,
iou_threshold=args.iou_threshold,
heatmap_alpha=args.heatmap_alpha,
radius=args.radius,
track_threshold=args.track_threshold,
track_seconds=args.track_seconds,
match_threshold=args.match_threshold,
)