You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
285 lines
9.6 KiB
285 lines
9.6 KiB
import os
|
|
# 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")
|
|
|
|
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
|
|
|
|
|
|
@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))
|
|
|
|
# 采样间隔（每秒处理2-3帧）
|
|
sample_interval = max(1, int(fps / 2)) # 每秒处理2帧
|
|
frame_count = 0
|
|
processed_frames = 0
|
|
detected_scenes = set()
|
|
|
|
# 检测结果去重（避免短时间内重复告警）
|
|
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)
|
|
|
|
# 更新进度（每10%更新一次）
|
|
if progress % 10 == 0 and progress != self.request.meta.get('progress', 0):
|
|
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)
|
|
|
|
# 发送告警
|
|
for current_model_type, alerts in scene_alerts.items():
|
|
detected_scenes.add(current_model_type)
|
|
current_time = time.time()
|
|
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,
|
|
"msg": f"detected {current_model_type} event at frame {frame_count}",
|
|
}
|
|
|
|
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()
|
|
|
|
# 发送完成通知
|
|
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"
|
|
}
|
|
|
|
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),
|
|
}
|
|
|
|
except Exception as e:
|
|
# 出错时也要清理临时文件
|
|
if os.path.exists(video_path):
|
|
try:
|
|
os.unlink(video_path)
|
|
except:
|
|
pass
|
|
raise e
|