yolo 视频流播放并进行目标识别

news2024/12/22 13:29:13

根据视频流，实时的进行目标识别

一、下载 [lal](https://github.com/q191201771/lal/releases/tag/v0.37.4)
二、安装 [FFmpeg](https://ffmpeg.org/)
三、完整代码演示

需要前置了解YOLO的完整操作使用labelImg标注，YOLO进行目标训练

一、下载 lal

下载解压，并运行lalserver.exe

报错是因为cert.pem和key.pem的原因（无法使用https，无关紧要）

再运行一个输入流

ffmpeg -stream_loop -1 -re -i wjj.mp4 -c:v libx264 -preset veryfast -b:v 3000k -maxrate 3000k -bufsize 6000k -c:a aac -b:a 128k -f flv rtmp://127.0.0.1:1935/live/1

二、安装 FFmpeg

随后将下载下来的FFmpeg进行解压，再进行环境变量配置

~~前面两个软件下载安放位置及视频位置~~

三、完整代码演示

注意，想看见效果需要下载 VLC media player
在这里插入图片描述
他会根据挖机视频的运动，而实时框选目标

完整代码如下：

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


# 加载 YOLO 模型
model = YOLO('runs/detect/dog_detection_finetune4/weights/best.pt')  # 替换为实际模型路径

# 输入 RTMP 流地址
input_rtmp = "rtmp://127.0.0.1:1935/live/1"  # 替换为实际输入流地址
output_rtmp = "rtmp://127.0.0.1:1935/live/2"  # 替换为实际输出流地址

# 打开视频流
cap = cv2.VideoCapture(input_rtmp)
if not cap.isOpened():
    print("无法打开视频流！请检查 RTMP 地址。")
    exit()

# 获取输入流的参数
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))

# 使用 FFmpeg 推流器
ffmpeg_process = (
    ffmpeg
    .input('pipe:', format='rawvideo', pix_fmt='bgr24', s=f"{width}x{height}", r=fps)
    .output(output_rtmp, format='flv', vcodec='libx264', pix_fmt='yuv420p', r=fps)
    .global_args('-loglevel', 'info')  # 打印日志信息
    .overwrite_output()
    .run_async(pipe_stdin=True, pipe_stdout=True, pipe_stderr=True)  # 允许读取日志
)

# 打印 FFmpeg 输出日志
def print_ffmpeg_logs(process):
    for line in process.stderr:
        print(line.decode('utf-8'), end='')

# 处理视频流
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        print("读取视频帧失败，可能是视频结束。")
        break

    # 使用 YOLO 对当前帧进行目标识别
    results = model.predict(frame, imgsz=640, conf=0.5, iou=0.5)
    for result in results:
        boxes = result.boxes
        for box in boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = box.conf[0]
            cls = int(box.cls[0])
            label = f"{model.names[cls]} {conf:.2f}"

            # 绘制检测框和标签
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    # 将处理后的帧传递给 FFmpeg 推流器
    ffmpeg_process.stdin.write(frame.tobytes())

# 打印推流结束日志
print_ffmpeg_logs(ffmpeg_process)

# 释放资源
cap.release()
ffmpeg_process.stdin.close()
ffmpeg_process.wait()
print("推流结束。")