上一次介绍了如何使用最基本的 Yolo-word来做检测,现在我们在加opencv来做个实时检测的例子
基本思路
1、读取离线视频流
2、将视频帧给yolo识别
3、根据识别结果 对视频进行绘制边框、加文字之类的
完整代码如下:
import datetimefrom ultralytics import YOLO
import cv2
from loguru import logger as log#加载YOLO模型
model = YOLO('model/yolov8s-world.pt')resize_width = 1920
resize_height = 1080def predict(chosen_model, img, classes = [], conf = 0.5):img = cv2.resize(img, (resize_width, resize_height))if classes:results = chosen_model.predict(img, classes = classes, conf = conf, save_txt = False)else:results = chosen_model.predict(img, conf = conf, save_txt = False)return resultsdef predict_and_detect(chosen_model, img, classes = [], conf = 0.5):img = cv2.resize(img, (resize_width, resize_height))cv2.putText(img, f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",(10, 20),cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 1)results = predict(chosen_model, img, classes, conf = conf)person = 0for result in results:for box in result.boxes:# 如果标签是人的,将盒子做成绿色,并在盒子上用大号字体打印置信度if result.names[int(box.cls[0])] == "person":person += 1cv2.rectangle(img, (int(box.xyxy[0][0]), int(box.xyxy[0][1])),(int(box.xyxy[0][2]), int(box.xyxy[0][3])), (0, 255, 0), 2)cv2.putText(img, f"{result.names[int(box.cls[0])]} {box.conf[0]:.2f}",(int(box.xyxy[0][0]), int(box.xyxy[0][1]) - 10),cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)else:cv2.rectangle(img, (int(box.xyxy[0][0]), int(box.xyxy[0][1])),(int(box.xyxy[0][2]), int(box.xyxy[0][3])), (0, 0, 255), 2)cv2.putText(img, f"{result.names[int(box.cls[0])]} {box.conf[0]:.2f}",(int(box.xyxy[0][0]), int(box.xyxy[0][1]) - 10),cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 1)if person > 0:log.error(f"当前发现有{person}个人")return img, resultsdef main():# 在处理下一个帧之前跳过的帧数skip_frames = 2frame_count = 0cap = cv2.VideoCapture(0)while True:ret, frame = cap.read()if not ret:breakframe_count = 1 + frame_countif frame_count % skip_frames != 0:continueresult_frame, _ = predict_and_detect(model, frame)cv2.imshow("video", result_frame)if cv2.waitKey(1) & 0xFF == ord('q'):breakcap.release()cv2.destroyAllWindows()if __name__ == '__main__':main()