【yolov5改编系列之四】多进程调用摄像头（也可以是视频流，如：rtsp://，http://）使两个检测文件同时运行，兵分两路

引言

本节内容是采用多进程，使两个detect.py文件同时检测摄像头（也可以是视频流，如：rtsp://，http://）采集到的图像（视频帧）。

使用方法：在官方yolov5-6.0文件夹中创建三个.py文件，将本节三个部分的内容分别放进去，配置好环境，然后运行process.py。运行结果如下图所示。

功能：你可以将其功能拓展，比如两个检测文件可以使用自己训练的不同的权重文件，分别执行不同的功能，分别输出，互不干扰。

基本原理：detect11.py和detect22.py两个文件是由一般的检测图片的文件改编的。一方面改成了可以调用的形式，另一方面其内部是循环接收外部传入图片的形式。

process.py

import cv2
from multiprocessing import Process, RawArray, Lock
from detect22 import run2
from detect11 import run1

# 定义摄像头参数
WIDTH = 640
HEIGHT = 640

# 创建共享内存空间
Image_raw = RawArray('B', WIDTH * HEIGHT * 3)
lock = Lock()
# 将摄像头的视频帧存储到共享内存空间
def img2memory(lock, img, raws):
    h, w, _ = img.shape
    lock.acquire()  # 获取线程锁，确保在修改共享内存时只有一个线程访问
    memoryview(raws).cast('B')[:img.size] = img.ravel()  # 将图像数据写入共享内存的对应位置
    lock.release()  # 释放线程锁，允许其他线程访问共享内存
    return w, h  # 返回图像的宽度和高度

#读取摄像头
def read_camera(lock,  raws, fr=0):
        # address = 'rtsp://{}@{}:554/h264/ch1/main/av_stream'  # 创建 RTSP 地址
        address = 0
        cap = cv2.VideoCapture(address)  # 打开 RTSP 地址，创建 VideoCapture 对象
        while True:  # 在 flag 为真且时间限制内的循环中读取帧数据
            # 读取帧
            ret, img = cap.read() # 读取一帧图像，将结果存储在 ret, img 变量中
            # 调整帧的大小为共享内存空间的宽度和高度
            img = cv2.resize(img, (WIDTH, HEIGHT))
            if not ret: continue  # 如果没有成功读取到帧，则跳过当前循环
            img2memory(lock, img, raws)  # 调用 img2memory 函数，将帧数据写入共享内存

def main():
    # 创建运行子进程
    process1 = Process(target=run1, args=(Image_raw,))
    process2 = Process(target=run2, args=(Image_raw,))
    process1.start()
    process2.start()

    # 创建读取摄像头的线程
    process_main = Process(target=read_camera, args=(lock, Image_raw))
    process_main.start()

    # 主线程等待所有线程结束
    process_main.join()
    process1.join()
    process2.join()

if __name__ == '__main__':
    main()

detect11.py

import cv2
import numpy as np
import torch

from models.experimental import attempt_load
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.plots import Annotator, colors
from utils.augmentations import letterbox
from utils.torch_utils import select_device

def run1(image_path):
    img_size = 640
    stride = 32
    weights = 'yolov5s.pt'  # 模型权重文件路径
    device = 'cpu'  # 设置设备类型
    # image_path = 'data/images/bus.jpg'  # 输入图像路径（也可以是绝对路径）
    save_path = 'run1/11.jpg'  # 输出图像保存路径（也可以是绝对路径）
    view_img = True  # 是否显示检测结果的图像窗口
    half = False

    device = select_device(device)
    half &= device.type != 'cpu'  # half precision only supported on CUDA

    # 导入模型
    model = attempt_load(weights, map_location=device)  # 加载模型
    img_size = check_img_size(img_size, s=stride)  # 检查图像尺寸是否符合规范
    names = model.names  # 获取模型输出的类别名称（people、bus等）

    while True:
        # Padded resize
        # img0 = cv2.imread(image_path)  # 读取输入图像
        img0 = np.frombuffer(image_path, dtype=np.uint8).reshape((640, 640, 3))

        img = letterbox(img0, img_size, stride=stride, auto=True)[0]  # 对输入图像进行填充和调整大小

        # Convert
        img = img.transpose((2, 0, 1))[::-1]  # 图像通道转换和颜色通道转换
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(device)
        img = img.float() / 255.0   # 归一化图像数据
        img = img[None]     # [h w c] -> [1 h w c]

        # inference
        pred = model(img)[0]  # 进行目标检测
        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45, classes=None, max_det=1000)  # 非最大抑制，保留置信度最高的目标框

        # plot label
        det = pred[0]
        annotator = Annotator(img0, line_width=3, example=str(names))  # 创建绘制标签的对象
        if len(det):
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()  # 将检测结果的坐标转换到原始图像坐标系中
            for *xyxy, conf, cls in reversed(det):  # 遍历每个目标框
                c = int(cls)  # 目标类别
                label = f'{names[c]} {conf:.2f}'  # 标签文字内容
                annotator.box_label(xyxy, label, color=colors(c, True))  # 绘制标签框和文字

        # write image
        im0 = annotator.result()  # 获取带有标签的图像
        cv2.imwrite(save_path, im0)  # 保存图像
        if view_img:
            im0 = cv2.resize(im0, None, fx=1, fy=1, interpolation=cv2.INTER_CUBIC) # 按比例修改需要展示的图像大小
            cv2.imshow(str('image-1'), im0)  # 显示图像
            cv2.waitKey(1)
        print(f'Inference {image_path} finish, save to {save_path}') # 打印图片来源和输出保存路径

detect22.py

import cv2
import numpy as np
import torch

from models.experimental import attempt_load
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.plots import Annotator, colors
from utils.augmentations import letterbox
from utils.torch_utils import select_device

def run2(image_path):
    img_size = 640
    stride = 32
    weights = 'yolov5s.pt'  # 模型权重文件路径
    device = 'cpu'  # 设置设备类型
    # image_path = 'data/images/bus.jpg'  # 输入图像路径（也可以是绝对路径）
    save_path = 'run1/11.jpg'  # 输出图像保存路径（也可以是绝对路径）
    view_img = True  # 是否显示检测结果的图像窗口
    half = False

    device = select_device(device)
    half &= device.type != 'cpu'  # half precision only supported on CUDA

    # 导入模型
    model = attempt_load(weights, map_location=device)  # 加载模型
    img_size = check_img_size(img_size, s=stride)  # 检查图像尺寸是否符合规范
    names = model.names  # 获取模型输出的类别名称（people、bus等）

    while True:
        # Padded resize
        # img0 = cv2.imread(image_path)  # 读取输入图像

        img0 = np.frombuffer(image_path, dtype=np.uint8).reshape((640, 640, 3))
        img = letterbox(img0, img_size, stride=stride, auto=True)[0]  # 对输入图像进行填充和调整大小

        # Convert
        img = img.transpose((2, 0, 1))[::-1]  # 图像通道转换和颜色通道转换
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(device)
        img = img.float() / 255.0   # 归一化图像数据
        img = img[None]     # [h w c] -> [1 h w c]

        # inference
        pred = model(img)[0]  # 进行目标检测
        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45, classes=None, max_det=1000)  # 非最大抑制，保留置信度最高的目标框

        # plot label
        det = pred[0]
        annotator = Annotator(img0, line_width=3, example=str(names))  # 创建绘制标签的对象
        if len(det):
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()  # 将检测结果的坐标转换到原始图像坐标系中
            for *xyxy, conf, cls in reversed(det):  # 遍历每个目标框
                c = int(cls)  # 目标类别
                label = f'{names[c]} {conf:.2f}'  # 标签文字内容
                annotator.box_label(xyxy, label, color=colors(c, True))  # 绘制标签框和文字

        # write image
        im0 = annotator.result()  # 获取带有标签的图像
        cv2.imwrite(save_path, im0)  # 保存图像
        if view_img:
            im0 = cv2.resize(im0, None, fx=1, fy=1, interpolation=cv2.INTER_CUBIC) # 按比例修改需要展示的图像大小
            cv2.imshow(str('image-2'), im0)  # 显示图像
            cv2.waitKey(1)
        print(f'Inference {image_path} finish, save to {save_path}') # 打印图片来源和输出保存路径