提前配置好yolo8的虚拟环境,同时补充安装包 supervision
#开 发 者 :刘 海 超
#开 发 时 间:2023-05-18 14:20
import cv2
import supervision as sv
from ultralytics import YOLO
def main():
# 保存视频
writer = cv2.VideoWriter('webcam_yolo.mp4',
cv2.VideoWriter_fourcc(*'mp4v'),
7,
(1280, 720))
# define resolution
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
# specify the model
model = YOLO("yolov8m.pt")
# 自定义边框
box_annotator = sv.BoxAnnotator(
thickness=2,
text_thickness=2,
text_scale=1
)
while True:
ret, frame = cap.read()
result = model(frame, agnostic_nms=True)[0]
detections = sv.Detections.from_yolov8(result)
labels = [
f"{model.model.names[class_id]} {confidence:0.2f}"
for _, _, confidence, class_id,_
in detections
]
frame = box_annotator.annotate(
scene=frame,
detections=detections,
labels=labels
)
writer.write(frame)
cv2.imshow("yolov8", frame)
if (cv2.waitKey(30) == 27): # break with escape key
break
cap.release()
writer.release()
cv2.destroyAllWindows()
if __name__ == "__main__":
main()