在GPU下使用最新的opencv4.4（yolov4和yolov4-tiny）对rtsp视频流进行检测并保存

前置条件（opencv4.4带cuda的编译安装）:https://blog.csdn.net/qq_34717531/article/details/107763872

import cv2
import time

CONFIDENCE_THRESHOLD = 0.25#置信度阀值
NMS_THRESHOLD = 0.4#非极大值抑制阀值
COLORS = [(0, 255, 255), (255, 255, 0), (0, 255, 0), (255, 0, 0)]#颜色

class_names = []#初始化一个列表以存储类名
#每个列表的元素中的末尾都会有一个“\n”，这是由于换行的时候产生的，但我们需要文件中将
#每一个类变成列表的每一个元素，不希望存在“\n”
#python strip()方法用于移除字符串头尾指定的字符
with open("coco.names", "r") as f:
    class_names = [cname.strip() for cname in f.readlines()]

#保存处理后视频的一套
cap = cv2.VideoCapture("1.mp4")
#cap = cv2.VideoCapture(rtsp)
fps = int(cap.get(cv2.CAP_PROP_FPS))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fourcc = cv2.VideoWriter_fourcc(*"MJPG")
out = cv2.VideoWriter('11.avi', fourcc, int(fps), (int(width),int(height)))

#opencv的dnn模块(NVIDIA GPU的推理模块)
net = cv2.dnn.readNet("yolov4-tiny.weights", "yolov4-tiny.cfg")
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

model = cv2.dnn_DetectionModel(net)#检测模型加载
model.setInputParams(size=(416, 416), scale=1/256)#调整尺寸大小

while cv2.waitKey(1) < 1:
    (grabbed, frame) = cap.read()#获取每一帧
    if not grabbed:
        exit()
    
    start = time.time()
    classes, scores, boxes = model.detect(frame, CONFIDENCE_THRESHOLD, NMS_THRESHOLD)
    end = time.time()

    start_drawing = time.time()
    for (classid, score, box) in zip(classes, scores, boxes):#循环检测
        color = COLORS[int(classid) % len(COLORS)]
        label = "%s : %f" % (class_names[classid[0]], score)
        cv2.rectangle(frame, box, color, 2)
        cv2.putText(frame, label, (box[0], box[1] ), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
    end_drawing = time.time()
    
    fps_label = "FPS: %.2f (excluding drawing time of %.2fms)" % (1 / (end - start),(end_drawing - start_drawing) * 1000)
    cv2.putText(frame, fps_label, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 0, 255), 2)
    out.write(frame)
    cv2.imshow("detections", frame)
cap.release()
out.release()
cv2.destoryAllWindows()
    

yolov4 / yolov4-tiny对比与darknet要快（我的显卡是GTX1080）。