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()