diff --git a/app/main.py b/app/main.py index b2bb330..938b4fa 100644 --- a/app/main.py +++ b/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 diff --git a/app/tasks.py b/app/tasks.py index 716da57..a919e28 100644 --- a/app/tasks.py +++ b/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)