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.
96 lines
3.2 KiB
96 lines
3.2 KiB
import os
|
|
import cv2
|
|
import json
|
|
import requests
|
|
from celery import Celery
|
|
from ultralytics import YOLO
|
|
from shapely.geometry import Point, Polygon
|
|
|
|
# 初始化 Celery
|
|
redis_url = os.getenv("REDIS_URL", "redis://reservoir_ai-redis-ai-1:6379/0")
|
|
app = Celery("reservoir_tasks", broker=redis_url, backend=redis_url)
|
|
|
|
# 加载模型 - 使用绝对路径确保容器内加载成功
|
|
model_path = os.path.join(os.path.dirname(__file__), 'yolov8n.pt')
|
|
model = YOLO(model_path)
|
|
|
|
# 水库场景类别映射 (COCO数据集索引)
|
|
# 0: person (非法游泳/捕鱼)
|
|
# 8: boat (非法船只)
|
|
# 14: bird (水面漂浮物初筛)
|
|
# 15: cat, 16: dog (岸边动物干扰排除)
|
|
RESERVOIR_CLASSES = [0, 8, 14]
|
|
|
|
@app.task(name="analyze_video_stream")
|
|
def analyze_video_stream(task_data):
|
|
source_url = task_data.get("source_url")
|
|
webhook_url = task_data.get("webhook_url")
|
|
scene_type = task_data.get("scene_type")
|
|
roi_points = task_data.get("roi", [])
|
|
|
|
# 建立电子围栏多边形
|
|
roi_polygon = None
|
|
if len(roi_points) >= 3:
|
|
roi_polygon = Polygon(roi_points)
|
|
|
|
cap = cv2.VideoCapture(source_url)
|
|
if not cap.isOpened():
|
|
return {"status": "error", "message": f"无法打开视频流: {source_url}"}
|
|
|
|
frame_count = 0
|
|
while cap.isOpened():
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# 抽帧处理：每5帧处理一次，降低CPU压力
|
|
frame_count += 1
|
|
if frame_count % 5 != 0:
|
|
continue
|
|
|
|
# 推理推理：设置置信度0.4，只看特定类别
|
|
results = model.track(
|
|
frame,
|
|
persist=True,
|
|
conf=0.4,
|
|
classes=RESERVOIR_CLASSES,
|
|
verbose=False
|
|
)
|
|
|
|
for r in results:
|
|
if r.boxes:
|
|
boxes = r.boxes
|
|
for box in boxes:
|
|
# 获取坐标
|
|
x1, y1, x2, y2 = box.xyxy[0].tolist()
|
|
conf = float(box.conf[0])
|
|
cls = int(box.cls[0])
|
|
label = model.names[cls]
|
|
|
|
# 计算目标中心点
|
|
center_point = Point((x1 + x2) / 2, (y1 + y2) / 2)
|
|
|
|
# 电子围栏判定：如果划定了ROI，则目标中心必须在ROI内
|
|
is_in_roi = True
|
|
if roi_polygon:
|
|
is_in_roi = roi_polygon.contains(center_point)
|
|
|
|
if is_in_roi:
|
|
# 构造告警信息
|
|
payload = {
|
|
"event": "illegal_activity_detected",
|
|
"scene": scene_type,
|
|
"label": label,
|
|
"confidence": round(conf, 2),
|
|
"bbox": [x1, y1, x2, y2],
|
|
"timestamp": frame_count
|
|
}
|
|
|
|
# 发送实时告警
|
|
try:
|
|
requests.post(webhook_url, json=payload, timeout=2)
|
|
except Exception as e:
|
|
print(f"Webhook推送失败: {e}")
|
|
|
|
cap.release()
|
|
return {"status": "completed", "processed_frames": frame_count}
|