import os import time import uuid # Disable CUDA before torch import to avoid encoding errors in container #os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") from typing import Dict, List, Optional import cv2 import numpy as np import requests from celery import Celery 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", broker=REDIS_URL, backend=REDIS_URL, ) def check_roi(point, roi_polygon: List[List[int]]) -> bool: polygon = np.array(roi_polygon, dtype=np.int32) return cv2.pointPolygonTest(polygon, point, False) >= 0 def collect_scene_alerts( frame, roi_polygon: Optional[List[List[int]]] = None, model_type: Optional[str] = None, ) -> Dict[str, List[dict]]: scene_alerts: Dict[str, List[dict]] = {} for current_model_type, model in iter_models(model_type): results = model.predict(frame, conf=0.4, verbose=False) alerts = [] for result in results: boxes = result.boxes if boxes is None: continue for box in boxes: cls = int(box.cls[0]) label = model.names[cls] bbox = box.xyxy[0].cpu().numpy() center_point = ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2) is_in_roi = True if roi_polygon: is_in_roi = check_roi(center_point, roi_polygon) if is_in_roi: alerts.append( { "label": label, "confidence": float(box.conf[0]), "bbox": bbox.tolist(), } ) if alerts: scene_alerts[current_model_type] = 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, stream_url: str, webhook_url: str, model_type: Optional[str] = None, roi_polygon: Optional[List[List[int]]] = None, ): if model_type is not None and model_type not in MODEL_MAP: return {"status": "error", "message": f"unsupported model type: {model_type}"} cap = cv2.VideoCapture(stream_url) if not cap.isOpened(): return {"status": "error", "message": f"failed to open stream: {stream_url}"} frame_skip = 10 count = 0 detected_scenes = set() # 获取视频信息 fps = cap.get(cv2.CAP_PROP_FPS) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if stream_url.endswith('.mp4') else None # 更新任务状态 self.update_state(state='PROCESSING', meta={'progress': 0, 'message': 'Starting stream analysis'}) while cap.isOpened(): success, frame = cap.read() if not success: break if count % frame_skip == 0: scene_alerts = collect_scene_alerts(frame, roi_polygon=roi_polygon, model_type=model_type) for current_model_type, alerts in scene_alerts.items(): detected_scenes.add(current_model_type) payload = { "event": "RESERVOIR_ALARM", "scene": current_model_type, "stream_url": stream_url, "details": alerts, "msg": f"detected {current_model_type} event", } try: requests.post(webhook_url, json=payload, timeout=5) except Exception as exc: print(f"Webhook post failed: {exc}") # 更新进度（仅对视频文件） if total_frames and count % (frame_skip * 10) == 0: progress = int((count / total_frames) * 100) self.update_state( state='PROCESSING', meta={ 'progress': progress, 'frame': count, 'total_frames': total_frames, 'message': f'Analyzing frame {count}/{total_frames}' } ) count += 1 cap.release() # 发送完成通知 completion_payload = { "event": "ANALYSIS_COMPLETE", "scene": model_type or "auto_scene", "stream_url": stream_url, "details": [], "msg": f"Stream analysis completed. Processed {count} frames" } try: requests.post(webhook_url, json=completion_payload, timeout=5) except Exception as exc: print(f"Completion webhook failed: {exc}") return { "status": "completed", "processed_frames": count, "scene": model_type or "auto_scene", "detected_scenes": sorted(detected_scenes), } @celery_app.task(name="analyze_video_file", bind=True) def analyze_video_file( self, video_path: str, webhook_url: str, model_type: Optional[str] = None, roi_polygon: Optional[List[List[int]]] = None, 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: break frame_count += 1 # 按照采样率处理 if frame_count % sample_interval != 0: continue processed_frames += 1 progress = int((frame_count / total_frames) * 100) current_progress = self.request.meta.get('progress', 0) if hasattr(self.request, 'meta') else 0 # 更新进度（每10%更新一次） if progress % 10 == 0 and progress != current_progress: self.update_state( state='PROCESSING', meta={ 'progress': progress, 'frame': frame_count, 'total_frames': total_frames, 'message': f'Analyzing frame {frame_count}/{total_frames}' } ) # 收集场景告警 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) # 为每个检测项添加时间戳（秒） 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: payload = { "event": "RESERVOIR_ALARM", "scene": current_model_type, "stream_url": f"local_file://{original_filename if original_filename else video_path}", "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", "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, "total_frames": total_frames, "fps": fps, "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) except: pass raise e