接摄像头 rtsp流 模型检测

import cv2
import numpy as np
from ultralytics import YOLO
import time

# 初始化YOLO模型
model = YOLO(r'C:\Users\DELL\Desktop\yolov8\runs\detect\train18\weights\best.pt')  # 请确保模型路径正确

# RTSP视频流URL
rtsp_url = "rtsp://admin:P@ssword@192.168.1.64:554/Streaming/Channels/1"

# 捕获RTSP视频流
cap = cv2.VideoCapture(rtsp_url)

if not cap.isOpened():
    print("Error: Cannot open camera.")
    exit()

# 设置帧处理间隔（例如每隔0.1秒处理一帧）
interval = 0.2
last_time_processed = time.time()

# 创建窗口
window_name = 'YOLO Detection'
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
# cv2.resizeWindow(window_name, 640, 640)

# 获取屏幕分辨率
screen_res = (1920, 1080)  # 根据你的屏幕分辨率进行调整

# 计算窗口位置
center_x = screen_res[0] // 2 - 320  # 屏幕宽度的一半减去窗口宽度的一半
center_y = screen_res[1] // 2 - 320  # 屏幕高度的一半减去窗口高度的一半

# 移动窗口到屏幕中间
cv2.moveWindow(window_name, center_x, center_y)

while True:
    ret, frame = cap.read()
    if not ret:
        break

    current_time = time.time()
    if current_time - last_time_processed >= interval:
        last_time_processed = current_time

        # 使用YOLO模型进行检测
        results = model(frame, stream=True, conf=0.5, classes=[0, 1, 2], imgsz=(1088, 1920))

        for r in results:
            # 获取检测结果
            boxes = r.boxes.xyxy.cpu().numpy()
            confidences = r.boxes.conf.cpu().numpy()
            class_ids = r.boxes.cls.cpu().numpy()

            for box, confidence, class_id in zip(boxes, confidences, class_ids):
                if confidence >= 0.5:  # 只显示置信度大于等于0.5的检测框
                    x1, y1, x2, y2 = map(int, box)
                    label = f'{model.names[int(class_id)]}: {confidence:.2f}'
                    color = (0, 255, 0)  # 绿色框
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

            # 将检测结果叠加到原始帧上
            result_image = frame

            # 显示结果
            cv2.imshow(window_name, result_image)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

cap.release()
cv2.destroyAllWindows()