feat：1.识别的时候，记录识别到物体的时间戳；2.视频截取前后2.5s，截取时需要保证物体的识别，返回mp4文件

app/main.py

@@ -8,6 +8,7 @@ import numpy as np
 from fastapi import FastAPI, File, HTTPException, Path, UploadFile, Form
 from pydantic import BaseModel, Field
 from celery.result import AsyncResult
+from starlette.responses import FileResponse
 
 from app.model_registry import MODEL_MAP, SCENE_INFO, iter_models, load_model
 from app.tasks import celery_app, analyze_video_stream, analyze_video_file
@@ -358,6 +359,7 @@ class WebhookPayload(BaseModel):
     stream_url: str = Field(..., description="视频流地址。", example="rtsp://example.com/live")
     details: List[WebhookDetectionItem] = Field(..., description="识别结果列表。")
     msg: str = Field(..., description="告警说明。", example="detected floating event")
+    video_url: Optional[str] = Field(None, description="处理后的视频下载链接")  # 新增
 
 
 def validate_model_type(model_type: str) -> None:
@@ -844,6 +846,19 @@ async def predict_image(
     except Exception as exc:
         raise HTTPException(status_code=500, detail=str(exc)) from exc
 
+@app.get("/v1/download/{filename}")
+async def download_video(filename: str):
+    """下载处理后的视频文件"""
+    video_dir = os.path.join(UPLOAD_DIR, "processed_videos")
+    print("video addr1",video_dir);
+    file_path = os.path.join(video_dir, filename)
+    if not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail="Video not found")
+    # 安全加固：防止路径穿越
+    if not os.path.realpath(file_path).startswith(os.path.realpath(video_dir)):
+        raise HTTPException(status_code=403, detail="Forbidden")
+    return FileResponse(file_path, media_type="video/mp4", filename=filename)
+
 
 if __name__ == "__main__":
     import uvicorn

app/tasks.py

@@ -1,4 +1,6 @@
 import os
+import time
+import uuid
 # Disable CUDA before torch import to avoid encoding errors in container
 #os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
 
@@ -14,6 +16,8 @@ from app.model_registry import MODEL_MAP, iter_models
 
 #REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
 REDIS_URL = os.getenv("CELERY_BROKER_URL", "redis://localhost:6379/0")
+UPLOAD_DIR = os.getenv("UPLOAD_DIR", "/tmp")
+HTTP_URL = os.getenv("HTTP_URL", "http://192.168.1.116:18005")
 
 celery_app = Celery(
     "reservoir_tasks",
@@ -68,6 +72,69 @@ def collect_scene_alerts(
     return scene_alerts
 
 
+def extract_clip_with_annotations(
+    video_path: str,
+    start_time: float,
+    end_time: float,
+    detections_by_frame: Dict[int, List[dict]],
+    output_path: str,
+) -> bool:
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        return False
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    start_frame = max(0, int(start_time * fps))
+    end_frame = min(total_frames - 1, int(end_time * fps))
+    if start_frame >= end_frame:
+        cap.release()
+        return False
+
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+    if not out.isOpened():
+        cap.release()
+        return False
+
+    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
+    for frame_idx in range(start_frame, end_frame + 1):
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # 计算当前帧在原始视频中的时间（秒）
+        current_time = frame_idx / fps
+
+        # 1. 在左上角显示时间戳（黄色文字）
+        timestamp_text = f"{current_time:.2f}s"
+        cv2.putText(frame, timestamp_text, (10, 30),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
+
+        # 2. 如果有检测记录，绘制检测框，并可在框上额外显示时间
+        if frame_idx in detections_by_frame:
+            for det in detections_by_frame[frame_idx]:
+                bbox = det['bbox']
+                label = det['label']
+                conf = det['confidence']
+                x1, y1, x2, y2 = map(int, bbox)
+                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                text = f"{label} {conf:.2f}"
+                cv2.putText(frame, text, (x1, y1 - 5),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+                # 可在框上方额外加时间戳（可选）
+                # cv2.putText(frame, timestamp_text, (x1, y1 - 20), ...)
+
+        out.write(frame)
+
+    cap.release()
+    out.release()
+    return True
+
+
 @celery_app.task(name="analyze_video_stream", bind=True)
 def analyze_video_stream(
     self,
@@ -164,34 +231,39 @@ def analyze_video_file(
     original_filename: str = None,
 ):
     # 分析本地MP4视频文件
-    
+
     try:
         # 更新任务状态
         self.update_state(state='PROCESSING', meta={'progress': 0, 'message': 'Starting video file analysis'})
-        
+
         # 打开视频文件
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
             raise Exception(f"Failed to open video file: {video_path}")
-        
+
         # 获取视频信息
         fps = cap.get(cv2.CAP_PROP_FPS)
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        
+        duration = total_frames / fps if fps > 0 else 0
+
         # 采样间隔（每秒处理2-3帧）
         sample_interval = max(1, int(fps / 2))  # 每秒处理2帧
         frame_count = 0
         processed_frames = 0
         detected_scenes = set()
-        
-        # 检测结果去重（避免短时间内重复告警）
+
+        # 新增：存储所有检测结果 (frame_idx -> list of detections)
+        detections_by_frame: Dict[int, List[dict]] = {}
+        # 存储所有检测的时间戳（用于截取区间）
+        detection_timestamps = []
+
         last_alarm_time = {}
         alarm_cooldown = 30  # 同类型告警冷却时间（秒）
-        
+
         import time
-        
+
         while True:
             ret, frame = cap.read()
             if not ret:
@@ -220,51 +292,87 @@ def analyze_video_file(
             
             # 收集场景告警
             scene_alerts = collect_scene_alerts(frame, roi_polygon=roi_polygon, model_type=model_type)
-            
-            # 发送告警
+
+            # 记录检测结果并发送告警
+            current_time = time.time()
             for current_model_type, alerts in scene_alerts.items():
                 detected_scenes.add(current_model_type)
-                current_time = time.time()
+                # 为每个检测项添加时间戳（秒）
+                timestamp = frame_count / fps if fps > 0 else 0
+                for det in alerts:
+                    det['timestamp'] = timestamp  # 增加时间戳字段
+                # 存储到检测字典
+                detections_by_frame[frame_count] = alerts  # 覆盖或合并？这里假设一帧可能多个场景，但收集函数按场景分，我们统一存
+                # 由于collect_scene_alerts返回的是按场景分组的dict，我们将其扁平化存储
+                # 但为了简单，我们只存储第一个场景？如果同一帧多个场景，我们需要合并。
+                # 修改：将alerts合并到该帧的列表中
+                if frame_count not in detections_by_frame:
+                    detections_by_frame[frame_count] = []
+                detections_by_frame[frame_count].extend(alerts)
+                detection_timestamps.append(timestamp)
+
                 alarm_key = f"{current_model_type}_{webhook_url}"
-                
-                # 检查冷却时间
-                if alarm_key not in last_alarm_time or \
-                   (current_time - last_alarm_time[alarm_key]) > alarm_cooldown:
-                    
+                if alarm_key not in last_alarm_time or (current_time - last_alarm_time[alarm_key]) > alarm_cooldown:
                     payload = {
                         "event": "RESERVOIR_ALARM",
                         "scene": current_model_type,
                         "stream_url": f"local_file://{original_filename if original_filename else video_path}",
-                        "details": alerts,
-                        "msg": f"detected {current_model_type} event at frame {frame_count}",
+                        "details": alerts,  # 已包含timestamp
+                        "msg": f"detected {current_model_type} event at {timestamp:.2f}s",
                     }
-                    
                     try:
                         requests.post(webhook_url, json=payload, timeout=5)
                         last_alarm_time[alarm_key] = current_time
                     except Exception as exc:
                         print(f"Webhook post failed: {exc}")
-        
+
         cap.release()
-        
-        # 发送完成通知
+
+        # 生成截取视频
+        video_url = None
+        if detection_timestamps:
+            min_ts = min(detection_timestamps)
+            max_ts = max(detection_timestamps)
+            start_time = max(0, min_ts - 2.5)
+            end_time = min(duration, max_ts + 2.5)
+            # 确保至少2秒长度
+            if end_time - start_time < 1.0:
+                start_time = max(0, min_ts - 1.5)
+                end_time = min(duration, min_ts + 1.5)
+
+            # 准备输出目录
+            video_dir = os.path.join(UPLOAD_DIR, 'processed_videos')
+            os.makedirs(video_dir, exist_ok=True)
+            output_filename = f"{uuid.uuid4()}.mp4"
+            output_path = os.path.join(video_dir, output_filename)
+
+            success = extract_clip_with_annotations(
+                video_path, start_time, end_time,
+                detections_by_frame, output_path
+            )
+
+            if success:
+                video_url = f"{HTTP_URL}/v1/download/{output_filename}"
+                # 可选：删除原临时文件？稍后统一删除
+
+        # 发送完成通知（增加video_url字段）
         completion_payload = {
             "event": "ANALYSIS_COMPLETE",
             "scene": model_type or "auto_scene",
             "stream_url": f"local_file://{original_filename if original_filename else video_path}",
             "details": [],
-            "msg": f"Video analysis completed. Processed {processed_frames} frames from {total_frames} total frames"
+            "msg": f"Video analysis completed. Processed {processed_frames} frames from {total_frames} total frames",
+            "video_url": video_url  # 新增字段
         }
-        
         try:
             requests.post(webhook_url, json=completion_payload, timeout=5)
         except Exception as exc:
             print(f"Completion webhook failed: {exc}")
-        
-        # 清理临时文件
+
+        # 清理临时文件（原上传的视频）
         if os.path.exists(video_path):
             os.unlink(video_path)
-        
+
         return {
             "status": "completed",
             "processed_frames": processed_frames,
@@ -273,10 +381,10 @@ def analyze_video_file(
             "resolution": {"width": frame_width, "height": frame_height},
             "scene": model_type or "auto_scene",
             "detected_scenes": sorted(detected_scenes),
+            "video_url": video_url,
         }
-        
+
     except Exception as e:
-        # 出错时也要清理临时文件
         if os.path.exists(video_path):
             try:
                 os.unlink(video_path)