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.
394 lines
14 KiB
394 lines
14 KiB
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
|