直接返回mp4格式文件路径-推流成功版本

app.py

@@ -18,6 +18,9 @@ from pydantic import BaseModel
 from ultralytics import YOLO
 import os
 import glob
+import subprocess
+import signal
+import time
 from typing import Optional
 
 # 设置日志格式和级别
@@ -362,11 +365,15 @@ class StreamParams(BaseModel):
     model: 推理模型路径
     source: 拉流地址（如rtsp/http视频流）
     stream_url: 推流地址（如rtmp推流地址）
+    fps: 输出帧率
+    bitrate: 输出码率
     其他参数同 predict
     """
     model: str = "yolov12m.pt"
     source: str = None
     stream_url: str = None
+    fps: int = 25
+    bitrate: str = "2000k"
     conf: float = 0.25
     iou: Optional[float] = 0.7
     imgsz: int = 640
@@ -377,43 +384,141 @@ class StreamParams(BaseModel):
 def yolov12_stream(params: StreamParams):
     """
     RESTful POST接口：/yolov12/stream
-    接收视频拉流地址和推流地址，调用YOLO模型推理，将推理后的视频推送到推流地址。
+    接收视频拉流地址和推流地址，调用YOLO模型推理，使用ffmpeg将推理后的视频推送到推流地址。
     返回格式：{"code": 0/1, "msg": "success/错误原因", "result": None}
     """
-    import cv2
-    import logging
     logging.info("收到/yolov12/stream请求")
     logging.info(f"请求参数: {params}")
+    
+    # 全局变量用于存储进程引用
+    ffmpeg_process = None
+    
+    def cleanup_process():
+        """清理ffmpeg进程"""
+        nonlocal ffmpeg_process
+        if ffmpeg_process:
+            try:
+                ffmpeg_process.terminate()
+                ffmpeg_process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                ffmpeg_process.kill()
+            except Exception as e:
+                logging.warning(f"清理ffmpeg进程时出错: {e}")
+    
     try:
         model = YOLO(params.model)
         cap = cv2.VideoCapture(params.source)
         if not cap.isOpened():
             return {"code": 1, "msg": f"无法打开视频流: {params.source}", "result": None}
+        
+        # 获取视频流信息
         fps = cap.get(cv2.CAP_PROP_FPS)
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        # 推流地址通常为rtmp/rtsp等
-        fourcc = cv2.VideoWriter_fourcc(*'flv1') if params.stream_url.startswith('rtmp') else cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(params.stream_url, fourcc, fps if fps > 0 else 25, (width, height))
-        if not out.isOpened():
-            cap.release()
-            return {"code": 1, "msg": f"无法打开推流地址: {params.stream_url}", "result": None}
+        
+        # 使用实际帧率，如果获取不到则使用参数中的fps
+        output_fps = fps if fps > 0 else params.fps
+        
+        # 构建ffmpeg命令
+        ffmpeg_cmd = [
+            'ffmpeg',
+            '-f', 'rawvideo',
+            '-pix_fmt', 'bgr24',
+            '-s', f'{width}x{height}',
+            '-r', str(output_fps),
+            '-i', '-',  # 从stdin读取
+            '-c:v', 'libx264',
+            '-preset', 'ultrafast',
+            '-tune', 'zerolatency',
+            '-b:v', params.bitrate,
+            '-maxrate', params.bitrate,
+            '-bufsize', '4000k',
+            '-g', str(output_fps * 2),  # GOP大小
+            '-f', 'flv' if params.stream_url.startswith('rtmp') else 'mpegts',
+            '-y',  # 覆盖输出文件
+            params.stream_url
+        ]
+        
+        logging.info(f"启动ffmpeg命令: {' '.join(ffmpeg_cmd)}")
+        
+        # 启动ffmpeg进程
+        ffmpeg_process = subprocess.Popen(
+            ffmpeg_cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            bufsize=0
+        )
+        
+        # 等待ffmpeg启动
+        time.sleep(1)
+        
+        if ffmpeg_process.poll() is not None:
+            # ffmpeg进程异常退出
+            stderr_output = ffmpeg_process.stderr.read().decode() if ffmpeg_process.stderr else "未知错误"
+            return {"code": 1, "msg": f"ffmpeg启动失败: {stderr_output}", "result": None}
+        
         frame_count = 0
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-            # 推理
-            results = model.predict(source=frame, imgsz=params.imgsz, conf=params.conf, iou=params.iou, device=params.device)
-            annotated_frame = results[0].plot()
-            out.write(annotated_frame)
-            frame_count += 1
-        cap.release()
-        out.release()
-        logging.info(f"推理并推流完成，共处理帧数: {frame_count}")
-        return {"code": 0, "msg": "success", "result": None}
+        start_time = time.time()
+        
+        try:
+            while cap.isOpened():
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                
+                # YOLO推理
+                try:
+                    predict_kwargs = {
+                        'source': frame,
+                        'imgsz': params.imgsz,
+                        'conf': params.conf,
+                        'device': params.device
+                    }
+                    
+                    # 只有当iou不为None时才添加到参数中
+                    if params.iou is not None:
+                        predict_kwargs['iou'] = params.iou
+                    
+                    results = model.predict(**predict_kwargs)
+                    annotated_frame = results[0].plot()
+                except Exception as predict_error:
+                    logging.error(f"YOLO推理出错: {predict_error}")
+                    # 如果推理失败，使用原始帧
+                    annotated_frame = frame
+                
+                # 将处理后的帧写入ffmpeg
+                try:
+                    ffmpeg_process.stdin.write(annotated_frame.tobytes())
+                    ffmpeg_process.stdin.flush()
+                    frame_count += 1
+                    
+                    # 每100帧输出一次进度
+                    if frame_count % 100 == 0:
+                        elapsed_time = time.time() - start_time
+                        current_fps = frame_count / elapsed_time
+                        logging.info(f"已处理 {frame_count} 帧，当前FPS: {current_fps:.2f}")
+                        
+                except IOError as e:
+                    logging.error(f"写入ffmpeg时出错: {e}")
+                    break
+                    
+        except KeyboardInterrupt:
+            logging.info("收到中断信号，停止处理")
+        finally:
+            # 清理资源
+            cap.release()
+            cleanup_process()
+        
+        elapsed_time = time.time() - start_time
+        avg_fps = frame_count / elapsed_time if elapsed_time > 0 else 0
+        
+        logging.info(f"推理并推流完成，共处理帧数: {frame_count}，平均FPS: {avg_fps:.2f}")
+        return {"code": 0, "msg": "success", "result": {"frames_processed": frame_count, "avg_fps": avg_fps}}
+        
     except Exception as e:
         logging.error(f"/yolov12/stream 发生异常: {e}")
+        cleanup_process()
         return {"code": 1, "msg": str(e), "result": None}
 
 # 全局异常处理器：参数校验失败时统一返回格式