在opencv中实现多处理以同时测试两个视频以进行对象检测

Question

我正在使用 PyTorch 和 OpenCV 的 YOLO 算法实现对象检测模型。在单个视频上运行我的模型效果很好。但每当我尝试使用多重处理来一次测试更多视频时，它就会冻结。你能解释一下这段代码有什么问题吗？

import torch
import cv2
import time
from multiprocessing import Process

model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/best.pt', force_reload=True)

def detectObject(video,name):
    cap = cv2.VideoCapture(video)
    while cap.isOpened():
        pTime = time.time()
        ret, img = cap.read()
        cTime = time.time()
        fps = str(int(1 / (cTime - pTime)))
        if img is None:
            break
        else:
            results = model(img)
            labels = results.xyxyn[0][:, -1].cpu().numpy()
            cord = results.xyxyn[0][:, :-1].cpu().numpy()
            n = len(labels)
            x_shape, y_shape = img.shape[1], img.shape[0]
            for i in range(n):
                row = cord[i]
                # If score is less than 0.3 we avoid making a prediction.
                if row[4] < 0.3:
                    continue
                x1 = int(row[0] * x_shape)
                y1 = int(row[1] * y_shape)
                x2 = int(row[2] * x_shape)
                y2 = int(row[3] * y_shape)
                bgr = (0, 255, 0)  # color of the box
                classes = model.names  # Get the name of label index
                label_font = cv2.FONT_HERSHEY_COMPLEX  # Font for the label.
                cv2.rectangle(img, (x1, y1), (x2, y2), bgr, 2)  # Plot the boxes
                cv2.putText(img, classes[int(labels[i])], (x1, y1), label_font, 2, bgr, 2)
                cv2.putText(img, f'FPS={fps}', (8, 70), label_font, 3, (100, 255, 0), 3, cv2.LINE_AA)

        img = cv2.resize(img, (700, 700))
        cv2.imshow(name, img)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()

Videos = ['../Dataset/Test1.mp4','../Dataset/Test2.mp4']
for i in Videos:
    process = Process(target=detectObject, args=(i, str(i)))
    process.start()

每次运行该代码时它都会冻结。 这是输出：

Downloading: "https://github.com/ultralytics/yolov5/archive/master.zip" to /home/com/.cache/torch/hub/master.zip
YOLOv5 🚀 2022-6-27 Python-3.9.9 torch-1.11.0+cu102 CPU

Fusing layers... 
YOLOv5s summary: 213 layers, 7023610 parameters, 0 gradients
Adding AutoShape... 

Answer 1

我通过添加

torch multiprocessing

代码让它工作。

from torch.multiprocessing import Pool, Process, set_start_method
try:
     set_start_method('spawn', force=True)
except RuntimeError:
    pass

videos = ['videos/video1.mp4', 'videos/video2.mp4']
for i in videos:
    process = Process(target=detectObject, args=(i, str(i)))
    process.start()

通过这种方式，我可以同时运行多个视频。

Answer 2

对于那些仍然遇到错误的人；尝试下面的代码：祝你好运

import torch
import cv2
import time
# from multiprocessing import Process
from torch.multiprocessing import Pool, Process, set_start_method


model = torch.hub.load("yolov5", 'custom', path='best.pt', source='local', trust_repo=True, force_reload=True)

def detectObject(video,name):
    cap = cv2.VideoCapture(video)
    while cap.isOpened():
        pTime = time.time()
        ret, img = cap.read()
        cTime = time.time()
        fps = str(int(1 / (cTime - pTime)))
        if img is None:
            break
        else:
            results = model(img)
            labels = results.xyxyn[0][:, -1].cpu().numpy()
            cord = results.xyxyn[0][:, :-1].cpu().numpy()
            n = len(labels)
            x_shape, y_shape = img.shape[1], img.shape[0]
            for i in range(n):
                row = cord[i]
                # If score is less than 0.3 we avoid making a prediction.
                if row[4] < 0.3:
                    continue
                x1 = int(row[0] * x_shape)
                y1 = int(row[1] * y_shape)
                x2 = int(row[2] * x_shape)
                y2 = int(row[3] * y_shape)
                bgr = (0, 255, 0)  # color of the box
                classes = model.names  # Get the name of label index
                label_font = cv2.FONT_HERSHEY_COMPLEX  # Font for the label.
                cv2.rectangle(img, (x1, y1), (x2, y2), bgr, 2)  # Plot the boxes
                cv2.putText(img, classes[int(labels[i])], (x1, y1), label_font, 2, bgr, 2)
                cv2.putText(img, f'FPS={fps}', (8, 70), label_font, 3, (100, 255, 0), 3, cv2.LINE_AA)

        img = cv2.resize(img, (700, 700))
        cv2.imshow(name, img)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()

try:
    set_start_method('spawn', force=True)
    videos = ['video/trx_l.mp4', 'video/trx_r.mp4']
    for i in videos:
        process = Process(target=detectObject, args=(i, str(i)))
        process.start()

except RuntimeError:
    pass