模型、区域点、颜色从配置中读取

This commit is contained in:
zqc
2026-04-01 17:00:36 +08:00
parent 54a906211a
commit 59ab409d52
2 changed files with 86 additions and 31 deletions

View File

@@ -5,9 +5,10 @@ import requests
from biz.base_frame_processor import BaseFrameProcessorWorker from biz.base_frame_processor import BaseFrameProcessorWorker
from algorithm.common.npu_yolo_onnx_person_car_phone import YOLOv8_ONNX from algorithm.common.npu_yolo_onnx_person_car_phone import YOLOv8_ONNX
from yolox.tracker.byte_tracker import BYTETracker from yolox.tracker.byte_tracker import BYTETracker
from common.constants import MODEL_ROOT_PATH
# ========================= 走廊场景专属配置 ========================= # ========================= 走廊场景专属配置 =========================
MODEL_PATH = 'YOLO_Weight/kanshousuo.onnx' # 犯人检测onnx模型路径 DETECT_MODEL_PATH = 'YOLO_Weight/kanshousuo.onnx' # 犯人检测onnx模型路径
INPUT_SIZE = 640 # 模型输入尺寸 INPUT_SIZE = 640 # 模型输入尺寸
RTSP_FPS = 10 # 视频流目标FPS RTSP_FPS = 10 # 视频流目标FPS
ALERT_PUSH_INTERVAL = 5 # 相同报警5秒内仅推送1次 ALERT_PUSH_INTERVAL = 5 # 相同报警5秒内仅推送1次
@@ -15,24 +16,35 @@ ALERT_PUSH_URL = "http://123.57.151.210:10000/picenter/websocket/test/process"
# 消失判定中心点在ROI内消失后持续无检测的帧数1.0秒,可微调) # 消失判定中心点在ROI内消失后持续无检测的帧数1.0秒,可微调)
ROI_LOST_FRAMES_THRESH = int(0.5 * RTSP_FPS) ROI_LOST_FRAMES_THRESH = int(0.5 * RTSP_FPS)
# ========================= 5个ROI区域配置相对坐标,适配任意分辨率 ========================= # ========================= 默认ROI区域配置当config.yaml未配置时使用 =========================
# 格式:{ROI名称: [[x1,y1], [x2,y2], ...], ...} (多边形顶点,顺/逆时针均可) DEFAULT_DOOR_ROIS = {
# 相对坐标x/y 0~10=左/上1=右/下),可直接根据场景调整 "left_door_1": {
ROI_CONFIG = { "points": [[0.195, 0.242], [0.265, 0.17], [0.3, 0.63], [0.248, 0.8]],
"left_door_1": [[0.195, 0.242], [0.265, 0.17], [0.3, 0.63] ,[0.248, 0.8]], # 左侧1门ROI "color": [255, 0, 0]
"left_door_2": [[0.3, 0.1], [0.34, 0.08], [0.35, 0.43], [0.322, 0.52]], # 左侧2门 ROI }
"left_door_3": [[0.355, 0.06], [0.42, 0], [0.42, 0.18], [0.362, 0.36]], # 左侧3门ROI
"right_door_1": [[0.735, 0.142], [0.81, 0.22], [0.78, 0.8], [0.715, 0.65]], # 右侧1门 ROI
"right_door_2": [[0.65, 0.06], [0.7, 0.09], [0.69, 0.5], [0.65, 0.4]] # 右侧2门ROI
} }
# ================================================================================== # ==================================================================================
class PrisonerDoorDetector: class PrisonerDoorDetector:
def __init__(self, params=None): def __init__(self, params=None):
self.params = params or {} self.params = params or {}
# 1. 加载YOLO模型仅提取犯人检测结果
# 0. 从params解析ROI配置无则使用默认值
door_rois_config = self.params.get('door_rois', DEFAULT_DOOR_ROIS)
self.roi_config = {}
self.roi_colors = {}
for door_name, door_cfg in door_rois_config.items():
self.roi_config[door_name] = door_cfg['points']
self.roi_colors[door_name] = tuple(door_cfg['color'])
model_path = self.params.get('model_path')
if model_path:
full_model_path = f"{MODEL_ROOT_PATH}/{model_path}"
else:
full_model_path = DETECT_MODEL_PATH
self.detector = YOLOv8_ONNX( self.detector = YOLOv8_ONNX(
MODEL_PATH, full_model_path,
conf_threshold=0.5, # 置信度阈值,可根据模型精度调整 conf_threshold=0.5, # 置信度阈值,可根据模型精度调整
iou_threshold=0.45, # IOU阈值 iou_threshold=0.45, # IOU阈值
input_size=INPUT_SIZE input_size=INPUT_SIZE
@@ -70,9 +82,9 @@ class PrisonerDoorDetector:
def _get_roi_abs(self, roi_name): def _get_roi_abs(self, roi_name):
"""相对坐标转绝对像素坐标适配当前帧分辨率OpenCV要求int32""" """相对坐标转绝对像素坐标适配当前帧分辨率OpenCV要求int32"""
if roi_name not in ROI_CONFIG: if roi_name not in self.roi_config:
return None return None
roi_rel = np.array(ROI_CONFIG[roi_name], dtype=np.float64) roi_rel = np.array(self.roi_config[roi_name], dtype=np.float64)
roi_abs = roi_rel * np.array([self.frame_width, self.frame_height]) roi_abs = roi_rel * np.array([self.frame_width, self.frame_height])
return roi_abs.astype(np.int32) return roi_abs.astype(np.int32)
@@ -119,23 +131,19 @@ class PrisonerDoorDetector:
self.frame_height, self.frame_width = frame.shape[:2] self.frame_height, self.frame_width = frame.shape[:2]
current_frame_alerts = [] # 本帧报警信息 current_frame_alerts = [] # 本帧报警信息
# ========================= 1. 初始化ROI绝对坐标并绘制5个ROI ========================= # ========================= 1. 初始化ROI绝对坐标并绘制ROI =========================
roi_colors = { # 各ROI绘制颜色自定义区分
"left_door_1": (255, 0, 0), "left_door_2": (0, 255, 0),
"left_door_3": (0, 0, 255), "right_door_1": (255, 255, 0),
"right_door_2": (255, 165, 0)
}
self.roi_abs_cache.clear() self.roi_abs_cache.clear()
for roi_name, _ in ROI_CONFIG.items(): for roi_name in self.roi_config:
roi_abs = self._get_roi_abs(roi_name) roi_abs = self._get_roi_abs(roi_name)
if roi_abs is None: if roi_abs is None:
continue continue
self.roi_abs_cache[roi_name] = roi_abs self.roi_abs_cache[roi_name] = roi_abs
# 绘制ROI多边形闭合+ ROI名称标签 # 绘制ROI多边形闭合+ ROI名称标签
roi_draw = roi_abs.reshape((-1, 1, 2)) # OpenCV绘制要求形状 (n,1,2) roi_draw = roi_abs.reshape((-1, 1, 2)) # OpenCV绘制要求形状 (n,1,2)
cv2.polylines(frame, [roi_draw], isClosed=True, color=roi_colors[roi_name], thickness=2) color = self.roi_colors.get(roi_name, (255, 255, 255))
cv2.polylines(frame, [roi_draw], isClosed=True, color=color, thickness=2)
cv2.putText(frame, roi_name, (roi_abs[0][0], roi_abs[0][1] - 5), cv2.putText(frame, roi_name, (roi_abs[0][0], roi_abs[0][1] - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, roi_colors[roi_name], 2) cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
# ========================= 2. 模型推理:仅提取犯人检测框 ========================= # ========================= 2. 模型推理:仅提取犯人检测框 =========================
detect_results = self.detector(frame) detect_results = self.detector(frame)

View File

@@ -48,22 +48,69 @@ video_clip_retention_seconds: 3600 # 视频文件
video_clip_default_segment_duration: 2 # 默认分片时长fallback video_clip_default_segment_duration: 2 # 默认分片时长fallback
service_groups: service_groups:
- name: "kadian_group" # 服务组名称 #- name: "kadian_group" # 服务组名称
# video_source_type: "hls"
# ws_host: "0.0.0.0" # WebSocket 服务地址
# ws_port: 8765 # WebSocket 服务端口
# algorithm: "checkpoint" # 算法类型
# cameras: # 该组下的摄像头列表
# - id: 8
# index: "12345"
# name: Entrance
# params:
# model_path: "Kadian_sanshijiazi.onnx"
# roi_points:
# - [0.15, 0.001]
# - [0.5, 0.001]
# - [1.0, 0.8]
# - [0.35, 1.0]
- name: "indoor_group" # 服务组名称
video_source_type: "hls" video_source_type: "hls"
ws_host: "0.0.0.0" # WebSocket 服务地址 ws_host: "0.0.0.0" # WebSocket 服务地址
ws_port: 8765 # WebSocket 服务端口 ws_port: 8765 # WebSocket 服务端口
algorithm: "checkpoint" # 算法类型 algorithm: "indoor" # 算法类型
cameras: # 该组下的摄像头列表 cameras: # 该组下的摄像头列表
- id: 8 - id: 8
index: "12345" index: "12345"
name: Entrance name: Entrance
params: params:
model_path: "Kadian_sanshijiazi.onnx" model_path: "kanshousuo.onnx"
roi_points: door_rois:
- [0.15, 0.001] left_door_1:
- [0.5, 0.001] points:
- [1.0, 0.8] - [0.195, 0.242]
- [0.35, 1.0] - [0.265, 0.17]
- [0.3, 0.63]
- [0.248, 0.8]
color: [255, 0, 0]
left_door_2:
points:
- [0.3, 0.1]
- [0.34, 0.08]
- [0.35, 0.43]
- [0.322, 0.52]
color: [0, 255, 0]
left_door_3:
points:
- [0.355, 0.06]
- [0.42, 0.0]
- [0.42, 0.18]
- [0.362, 0.36]
color: [0, 0, 255]
right_door_1:
points:
- [0.735, 0.142]
- [0.81, 0.22]
- [0.78, 0.8]
- [0.715, 0.65]
color: [255, 255, 0]
right_door_2:
points:
- [0.65, 0.06]
- [0.7, 0.09]
- [0.69, 0.5]
- [0.65, 0.4]
color: [255, 165, 0]
#- name: "prison_group" # 服务组名称 #- name: "prison_group" # 服务组名称
# video_source_type: "hls" # video_source_type: "hls"
# ws_host: "0.0.0.0" # WebSocket 服务地址 # ws_host: "0.0.0.0" # WebSocket 服务地址