python多线程推理视频

因为是多线程推理所以没有画面预览框，除非把预览框放到主线程

• yolov8
• onnx
• 可以融合目标跟踪

import argparse
import cv2
import numpy as np
import onnxruntime as ort
from multiprocessing import Process,Queue


def pull_rtsp(source,image_queue):
    cap = cv2.VideoCapture(source)
    detect_skip = 3
    detect_skip_index = 1
    while cap.isOpened():
        success, frame = cap.read()
        if success:
            if detect_skip_index % detect_skip == 0:
                detect_skip_index = 1
                image_queue.put(frame)
            else:
                detect_skip_index = detect_skip_index + 1


def detection(onnx_model,confidence_thres,iou_thres,classes,image_queue):
    session = ort.InferenceSession(onnx_model, providers=["CUDAExecutionProvider"])
    model_inputs = session.get_inputs()

    input_shape = model_inputs[0].shape

    input_width = input_shape[2]

    input_height = input_shape[3]
    
    while True:
        if not image_queue.empty() :
            img = image_queue.get()

            img_height, img_width = img.shape[:2]

            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            img = cv2.resize(img, (input_width, input_height))

            image_data = np.array(img) / 255.0

            image_data = np.transpose(image_data, (2, 0, 1))  

            image_data = np.expand_dims(image_data, axis=0).astype(np.float32)

            result = session.run(None, {model_inputs[0].name: image_data})

            outputs = np.transpose(np.squeeze(result[0]))

            rows = outputs.shape[0]

            boxes = []
            scores = []
            class_ids = []

            x_factor = img_width / input_width
            y_factor = img_height / input_height

            for i in range(rows):
                classes_scores = outputs[i][4:]
                max_score = np.amax(classes_scores)

                if max_score >= confidence_thres:

                    class_id = np.argmax(classes_scores)

                    x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]

                    left = int((x - w / 2) * x_factor)
                    top = int((y - h / 2) * y_factor)
                    width = int(w * x_factor)
                    height = int(h * y_factor)

                    class_ids.append(class_id)
                    scores.append(max_score)
                    boxes.append([left, top, width, height])

            indices = cv2.dnn.NMSBoxes(boxes, scores, confidence_thres, iou_thres)

            if len(indices) == 0:
                print('当前画面没有目标')
            else: 
                for i in indices:
                    box = boxes[i]
                    score = scores[i]
                    class_id = class_ids[i]
                    print("i", i)
                    print("class_id", class_id)
            # 画框，并弹出预览，但是因为是多线程，不在主线程，所以无法弹出预览


if __name__ == "__main__":
    
    # pip install numpy==1.23.4
    # pip install onnxruntime==1.18.0 onnxruntime-gpu==1.18.0  opencv-python argparse
    
    try:
        parser = argparse.ArgumentParser()
        parser.add_argument("--model", type=str, default=r"D:\newgit\ultralytics\yolov8n.onnx", help="Input your ONNX model.")
        parser.add_argument("--source", type=str, default=r"D:\newgit\ultralytics\examples\YOLOv8-ONNXRuntime\166959951-1-208.mp4", help="Path to input image.")
        parser.add_argument("--conf-thres", type=float, default=0.5, help="Confidence threshold")
        parser.add_argument("--iou-thres", type=float, default=0.5, help="NMS IoU threshold")
        args = parser.parse_args()
        
        # 定长线程安全的待识别图片队列
        image_queue = Queue(60)
        # 推理后待合成视频队列
        
        # 拉流进程
        p1 = Process(target=pull_rtsp,kwargs={'source':args.source, 'image_queue':image_queue})
        # 推理进程
        p2 = Process(target=detection,kwargs={'onnx_model':args.model,'confidence_thres':args.conf_thres,'iou_thres':args.iou_thres,'classes':args.source,'image_queue':image_queue})
        # 图片合成视频推流到流媒体进程
        
        p1.start()
        p2.start()

    except Exception as e:
        print("error: ", e)