38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
import subprocess
|
||
import numpy as np
|
||
from pathlib import Path
|
||
from algo import YoloModel
|
||
|
||
|
||
def run():
|
||
FILE = Path(__file__).resolve()
|
||
ROOT = FILE.parents[0] # root directory
|
||
model = YoloModel(ROOT / "chache/best.pt")
|
||
width, height = 640, 640
|
||
fps = 0.2 # 每 5 秒一帧
|
||
command = [
|
||
'ffmpeg',
|
||
'-i', 'rtsp://admin:xkrs0425@172.16.20.15:554/D6/sub/av_stream',
|
||
'-an', # 不处理音频
|
||
'-f', 'image2pipe', # 输出为图像流
|
||
'-pix_fmt', 'bgr24', # OpenCV 默认是 BGR 格式
|
||
'-r', str(fps), # 每秒 0.2 帧 => 每 5 秒一帧
|
||
'-vcodec', 'rawvideo',
|
||
'-'
|
||
]
|
||
# 启动 FFmpeg 子进程
|
||
pipe = subprocess.Popen(command, stdout=subprocess.PIPE, bufsize=10 ** 8)
|
||
while True:
|
||
# 读取一帧原始数据（大小 = width * height * 3）0
|
||
raw_frame = pipe.stdout.read(width * height * 3)
|
||
if not raw_frame:
|
||
break
|
||
# 转换为 NumPy 数组并 reshape 成图像
|
||
frame = np.frombuffer(raw_frame, dtype=np.uint8).reshape((height, width, 3))
|
||
print(f"Frame shape: {frame.shape}, dtype: {frame.dtype}, min/max: {frame.min(), frame.max()}")
|
||
# # 使用 YOLO 进行推理
|
||
result = model.predict(frame)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
run() |