diff --git a/pipeline/handler/graffiti_visualizer.py b/pipeline/handler/graffiti_visualizer.py deleted file mode 100644 index 43a5025..0000000 --- a/pipeline/handler/graffiti_visualizer.py +++ /dev/null @@ -1,198 +0,0 @@ -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 - ) \ No newline at end of file