import numpy as np import cv2 import math from ..base_processor import BaseProcessor from ..pipeline_data import PipelineData class GraffitiVisualizer(BaseProcessor): """涂鸦可视化处理器：绘制热力图和方向箭头""" def __init__(self, name: str = "涂鸦可视化处理器"): super().__init__(name) # 可调参数 self.MAX_HEAT = 10.0 self.ARROW_LEN_MIN = 10.0 self.GRID_STEP = 40 self.GRID_HEAT_THRESHOLD = 0.6 self.GRID_DRAW_EVERY = 2 self.FULL_DECAY_EVERY = 30 self.frame_idx = 0 def process(self, data: PipelineData) -> PipelineData: """绘制涂鸦可视化""" if data.frame is None: return data self.frame_idx = data.frame_idx vis = data.frame.copy() # 获取涂鸦数据 heat_map = data.get_data("graffiti", "heat_map") dir_x_map = data.get_data("graffiti", "dir_x_map") dir_y_map = data.get_data("graffiti", "dir_y_map") track_hist = data.get_data("graffiti", "track_hist") current_events = data.get_data("retrograde", "current_events", {}) all_events = data.get_data("retrograde", "all_events", {}) # 绘制热力图 if heat_map is not None: heat_norm = np.clip(heat_map / self.MAX_HEAT, 0.0, 1.0) heat_u8 = (heat_norm * 255).astype(np.uint8) red_overlay = np.zeros_like(vis) red_overlay[:, :, 2] = heat_u8 alpha = 0.45 vis = cv2.addWeighted(vis, 1.0, red_overlay, alpha, 0) # 绘制网格方向箭头 if (dir_x_map is not None and dir_y_map is not None and self.frame_idx % self.GRID_DRAW_EVERY == 0): self._draw_grid_arrows(vis, heat_map, dir_x_map, dir_y_map) # 绘制车辆轨迹和事件 if track_hist is not None: self._draw_tracks_and_events(vis, track_hist, current_events, all_events) # 绘制检测框和方向 if data.current_result is not None: self._draw_detections(vis, data.current_result, track_hist) # 添加调试信息 self._draw_debug_info(vis) data.frame = vis return data def _draw_grid_arrows(self, vis, heat_map, dir_x_map, dir_y_map): """在网格上绘制方向箭头""" h, w = heat_map.shape for gy in range(0, h, self.GRID_STEP): for gx in range(0, w, self.GRID_STEP): x0, y0 = gx, gy x1 = min(gx + self.GRID_STEP, w) y1 = min(gy + self.GRID_STEP, h) win = heat_map[y0:y1, x0:x1] if win.size == 0: continue avg_heat = float(win.mean()) if avg_heat / self.MAX_HEAT >= self.GRID_HEAT_THRESHOLD: avg_dx = float(dir_x_map[y0:y1, x0:x1].sum()) avg_dy = float(dir_y_map[y0:y1, x0:x1].sum()) if avg_dx == 0 and avg_dy == 0: continue cell_center = (x0 + (x1 - x0) // 2, y0 + (y1 - y0) // 2) self._draw_arrow( vis, cell_center, (avg_dx * 5.0, avg_dy * 5.0), color=(0, 128, 255), thickness=2 ) def _draw_tracks_and_events(self, vis, track_hist, current_events, all_events): """绘制车辆轨迹和事件标签""" for tid, hist in track_hist.items(): if len(hist) == 0: continue # 绘制轨迹线 pts = np.array(hist, dtype=np.int32) for i in range(1, len(pts)): cv2.line(vis, tuple(pts[i - 1]), tuple(pts[i]), (255, 255, 255), 1) # 绘制当前位置 cx, cy = hist[-1] cv2.circle(vis, (int(cx), int(cy)), 3, (255, 255, 255), -1) # 绘制当前事件 if tid in current_events: event = current_events[tid] txt = f"id:{tid} {event}" cv2.putText( vis, txt, (int(cx) + 6, int(cy) - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2 ) # 绘制近期事件（过去5帧内） elif tid in all_events: last_frame_idx, event = all_events[tid] if event is not None and self.frame_idx - last_frame_idx <= 5: txt = f"id:{tid} {event}" cv2.putText( vis, txt, (int(cx) + 6, int(cy) - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2 ) def _draw_detections(self, vis, result, track_hist): """绘制检测框和车辆方向""" boxes = getattr(result, "boxes", None) if boxes is None: return for box in boxes: try: xyxy = getattr(box, "xyxy", None) if xyxy is None: continue x1, y1, x2, y2 = map(int, map(float, xyxy[0].tolist())) track_id = int(getattr(box, "id", -1)) cls = int(getattr(box, "cls", -1)) # 绘制检测框 cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 255, 255), 2) # 绘制车辆ID cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 cv2.putText( vis, f"id:{track_id}", (cx + 6, cy - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1 ) # 绘制车辆方向箭头 if track_id in track_hist and len(track_hist[track_id]) >= 2: hist = track_hist[track_id] p0 = hist[-2] p1 = hist[-1] v = (p1[0] - p0[0], p1[1] - p0[1]) self._draw_arrow( vis, (cx, cy), (v[0], v[1]), color=(0, 0, 255), thickness=2 ) except Exception: continue def _draw_arrow(self, img, start, vec, color=(0, 255, 0), thickness=2, tip_length=0.3): """绘制箭头""" sx, sy = int(round(start[0])), int(round(start[1])) vx, vy = float(vec[0]), float(vec[1]) norm = math.hypot(vx, vy) if norm < 1e-3: return scale = 1.0 if norm < self.ARROW_LEN_MIN: scale = self.ARROW_LEN_MIN / (norm + 1e-6) ex = int(round(sx + vx * scale)) ey = int(round(sy + vy * scale)) cv2.arrowedLine( img, (sx, sy), (ex, ey), color, thickness, tipLength=tip_length ) def _draw_debug_info(self, vis): """绘制调试信息""" cv2.putText( vis, f"Frame:{self.frame_idx}", (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200, 200, 200), 2 ) cv2.putText( vis, f"HeatMax:{self.MAX_HEAT} TrustThr:3.0", (10, 45), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1 )