yolo读取视频并检测
from time import sleep
import cv2 as cv, torch
from ultralytics import YOLO
import logging
import time
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
log = logging.getLogger(__name__)
# log.info("-- OpenCV Info --")
# log.info(cv.getBuildInformation())
# log.info("-" * 50)
# device = "cuda:0" if torch.cuda.is_available() else "cpu"
# log.info(f"Using device: {device}")
yolo_model = YOLO("./yolo_model/yolo11s.pt")
image_size = 640
yolo_model.overrides["imgsz"] = image_size
cls_map = {
0: "person",
1: "bicycle",
2: "car",
3: "motorcycle",
5: "bus",
7: "truck",
}
# 帧率
FPS_TARGET = 12
# 帧间隔
T_INTERVAL = 1.0 / FPS_TARGET
# 上一次时间
last_t = time.time()
# 检测间隔
detect_every = 3
# 检测计数
detect_count = 0
cap = cv.VideoCapture(
"https://smart.saas.vppark.cn/oss/1.mp4",
)
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# if time.time() - last_t < T_INTERVAL:
# continue
# last_t = time.time()
detect_count += 1
if detect_count % detect_every == 0 or True:
# 手动 resize
h0, w0 = frame.shape[:2]
scale = min(image_size / h0, image_size / w0)
h1, w1 = int(h0 * scale), int(w0 * scale)
frame_resize = cv.resize(frame, (w1, h1), interpolation=cv.INTER_LINEAR)
# 执行检测
results = yolo_model(
frame_resize, imgsz=image_size, classes=list(cls_map.keys())
)
# 把框映射回原图
for result in results:
boxes_data = result.boxes.data.clone()
boxes_data[..., :4] /= scale
result.boxes.data = boxes_data
react_frame = result.plot(img=frame, line_width=2)
cv.imshow("frame", react_frame)
else:
cv.imshow("frame", frame)
# if cv.waitKey(int(1000 / FPS_TARGET)) & 0xFF == ord("q"):
if cv.waitKey(1) & 0xFF == ord("q"):
break
cap.release()
cv.destroyAllWindows()