import cv2 import numpy as np import time import requests from biz.base_frame_processor import BaseFrameProcessorWorker from algorithm.common.npu_yolo_onnx_person_car_phone import YOLOv8_ONNX from yolox.tracker.byte_tracker import BYTETracker from common.constants import MODEL_ROOT_PATH # ========================= 走廊场景专属配置 ========================= DETECT_MODEL_PATH = 'YOLO_Weight/kanshousuo.onnx' # 犯人检测onnx模型路径 INPUT_SIZE = 640 # 模型输入尺寸 RTSP_FPS = 10 # 视频流目标FPS ALERT_PUSH_INTERVAL = 5 # 相同报警5秒内仅推送1次 ALERT_PUSH_URL = "http://123.57.151.210:10000/picenter/websocket/test/process" # 消失判定:中心点在ROI内消失后,持续无检测的帧数(1.0秒,可微调) ROI_LOST_FRAMES_THRESH = int(0.5 * RTSP_FPS) # ========================= 默认ROI区域配置(当config.yaml未配置时使用) ========================= DEFAULT_DOOR_ROIS = { "left_door_1": { "points": [[0.195, 0.242], [0.265, 0.17], [0.3, 0.63], [0.248, 0.8]], "color": [255, 0, 0] } } # ================================================================================== class PrisonerDoorDetector: def __init__(self, params=None): self.params = params or {} # 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( full_model_path, conf_threshold=0.5, # 置信度阈值,可根据模型精度调整 iou_threshold=0.45, # IOU阈值 input_size=INPUT_SIZE ) # 2. 初始化ByteTracker跟踪器(适配走廊单/多犯人跟踪) class TrackerArgs: track_thresh = 0.25 track_buffer = 20 # 减小缓冲避免跟踪漂移 match_thresh = 0.75 mot20 = False self.tracker = BYTETracker(TrackerArgs(), frame_rate=RTSP_FPS) # 3. 状态变量初始化 self.last_alert_time = 0.0 # 最后报警时间(防重复推送) # 犯人跟踪信息:{track_id: {'is_cx_in_roi': 中心点是否在ROI, 'lost_frames': 消失帧数, 'lost_roi': 消失的ROI名称, 'last_cxcy': 最后中心点坐标}} self.prisoner_track_info = {} self.frame_width = 0 # 帧宽度(动态获取) self.frame_height = 0 # 帧高度(动态获取) self.roi_abs_cache = {} # ROI绝对坐标缓存:{roi_name: np.int32数组} def compute_iou(self, boxA, boxB): """IOU计算:匹配跟踪框与犯人检测框,过滤非犯人目标""" xA = max(boxA[0], boxB[0]) yA = max(boxA[1], boxB[1]) xB = min(boxA[2], boxB[2]) yB = min(boxA[3], boxB[3]) interW = max(0, xB - xA) interH = max(0, yB - yA) interArea = interW * interH boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1]) boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1]) unionArea = boxAArea + boxBArea - interArea return interArea / unionArea if unionArea > 0 else 0.0 def _get_roi_abs(self, roi_name): """相对坐标转绝对像素坐标(适配当前帧分辨率,OpenCV要求int32)""" if roi_name not in self.roi_config: return None roi_rel = np.array(self.roi_config[roi_name], dtype=np.float64) roi_abs = roi_rel * np.array([self.frame_width, self.frame_height]) return roi_abs.astype(np.int32) def is_cxcy_in_roi(self, cx, cy): """判断犯人框**中心点(cx,cy)** 是否在任意ROI内,返回:(是否在ROI, 所在ROI名称)""" for roi_name, roi_abs in self.roi_abs_cache.items(): # OpenCV点在多边形内判定:>=0 表示在内部/边上 if cv2.pointPolygonTest(roi_abs, (cx, cy), False) >= 0: return (True, roi_name) return (False, "outside") # def push_alert(self, camera_id, track_id, lost_roi, last_cxcy, timestamp): # """报警推送:带频率限制,携带消失ROI、最后中心点坐标""" # current_time = time.time() # if current_time - self.last_alert_time < ALERT_PUSH_INTERVAL: # return False # # 构造报警信息(可根据平台要求扩展字段) # alert_info = { # "camera_id": camera_id, # "alert_type": "prisoner_cx_disappear_in_roi", # "prisoner_track_id": track_id, # "disappear_roi": lost_roi, # "last_cx": round(last_cxcy[0], 2), # "last_cy": round(last_cxcy[1], 2), # "timestamp": timestamp, # "details": f"犯人框中心点在{lost_roi}区域内消失,触发报警" # } # # 推送报警请求 # try: # requests.post(ALERT_PUSH_URL, json=alert_info, timeout=3) # print(f"[报警成功] {alert_info}") # self.last_alert_time = current_time # return True # except Exception as e: # print(f"[报警失败] 原因:{str(e)}") # return False def process_frame(self, frame, camera_id: int, timestamp: float) -> dict: """ 核心帧处理: 1. 绘制5个ROI区域 2. 检测+跟踪犯人 3. 判定中心点是否在ROI内 4. 中心点在ROI内消失则累计帧数,达到阈值触发报警 """ self.frame_height, self.frame_width = frame.shape[:2] current_frame_alerts = [] # 本帧报警信息 # ========================= 1. 初始化ROI绝对坐标并绘制ROI ========================= self.roi_abs_cache.clear() for roi_name in self.roi_config: roi_abs = self._get_roi_abs(roi_name) if roi_abs is None: continue self.roi_abs_cache[roi_name] = roi_abs # 绘制ROI多边形(闭合)+ ROI名称标签 roi_draw = roi_abs.reshape((-1, 1, 2)) # OpenCV绘制要求形状 (n,1,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.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) # ========================= 2. 模型推理:仅提取犯人检测框 ========================= detect_results = self.detector(frame) prisoner_dets_xyxy = [] # 仅存犯人检测框 [x1,y1,x2,y2] dets_for_tracker = [] # 跟踪器输入 [x1,y1,x2,y2,conf] if detect_results: for det in detect_results: x1, y1, x2, y2, conf, cls_id = det dets_for_tracker.append([x1, y1, x2, y2, conf]) # 替换为你模型中「犯人」的实际类别ID,此处默认cls_id=1 if cls_id == 1: prisoner_dets_xyxy.append([x1, y1, x2, y2]) # ========================= 3. 目标跟踪:更新犯人跟踪结果 ========================= dets_np = np.array(dets_for_tracker, dtype=np.float32) if dets_for_tracker else np.empty((0, 5)) track_results = self.tracker.update(dets_np, [self.frame_height, self.frame_width], [self.frame_height, self.frame_width]) # ========================= 4. 遍历跟踪结果:判定犯人中心点是否在ROI ========================= current_prisoner_tids = set() # 本帧存在的犯人track_id for track in track_results: track_id = track.track_id track_box = list(map(float, track.tlbr)) # 跟踪框 [x1,y1,x2,y2] # IOU匹配:过滤非犯人目标,仅保留真正的犯人 is_prisoner = False for p_box in prisoner_dets_xyxy: if self.compute_iou(track_box, p_box) > 0.3: is_prisoner = True break if not is_prisoner: continue # 计算犯人框**中心点坐标**(核心判定依据) cx = (track_box[0] + track_box[2]) / 2 cy = (track_box[1] + track_box[3]) / 2 # 判定中心点是否在ROI内,返回(是否在ROI, 所在ROI名称) is_cx_in_roi, current_roi = self.is_cxcy_in_roi(cx, cy) # 更新犯人跟踪信息:记录中心点状态、所在ROI、最后坐标,重置消失帧数 self.prisoner_track_info[track_id] = { "is_cx_in_roi": is_cx_in_roi, "lost_frames": 0, "lost_roi": current_roi, "last_cxcy": (cx, cy) } current_prisoner_tids.add(track_id) # 绘制犯人框+中心点+状态标签(可视化调试) x1, y1, x2, y2 = map(int, track_box) cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2) # 红色犯人框 cv2.circle(frame, (int(cx), int(cy)), 5, (0, 255, 255), -1) # 黄色中心点 cv2.putText(frame, f"Prisoner_{track_id}({current_roi})", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2) # ========================= 5. 核心判定:中心点在ROI内消失则报警 ========================= for track_id in list(self.prisoner_track_info.keys()): if track_id not in current_prisoner_tids: # 犯人本帧消失,获取其最后状态 track_info = self.prisoner_track_info[track_id] # 仅处理「**中心点原本在ROI内**」的消失情况 if track_info["is_cx_in_roi"]: track_info["lost_frames"] += 1 # 累计消失帧数 # 消失帧数达到阈值,触发报警 if track_info["lost_frames"] >= ROI_LOST_FRAMES_THRESH: # self.push_alert( # camera_id=camera_id, # track_id=track_id, # lost_roi=track_info["lost_roi"], # last_cxcy=track_info["last_cxcy"], # timestamp=timestamp # ) # 记录本帧报警信息 current_frame_alerts.append({ "time": timestamp, "camera_id": camera_id, "action": "Indoor Violation", "prisoner_track_id": track_id, "disappear_roi": track_info["lost_roi"], "last_cx": round(track_info["last_cxcy"][0], 2), "last_cy": round(track_info["last_cxcy"][1], 2) }) del self.prisoner_track_info[track_id] # 报警后清除状态,避免重复触发 else: del self.prisoner_track_info[track_id] # 中心点不在ROI的消失,直接清除 # ========================= 6. 绘制辅助信息(摄像头ID、在押犯人数) ========================= cv2.putText(frame, f"Camera: {camera_id}", (20, self.frame_height - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(frame, f"Prisoners: {len(current_prisoner_tids)}", (20, self.frame_height - 50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2) return {"image": frame, "alerts": current_frame_alerts} # ========================= 帧处理线程(对接原有框架,直接复用) ========================= class FrameProcessorWorker(BaseFrameProcessorWorker): """看守所走廊犯人检测 - 5ROI+中心点消失判定""" DETECTOR_FACTORY = lambda params: PrisonerDoorDetector(params) POST_TYPE = 3 # 与原有业务区分,自定义即可 TARGET_FPS = RTSP_FPS