HighwayEventDet/pipeline/handler/DrawDirectionProcessor.py

import cv2
import numpy as np
import math
from collections import deque
from ..base_processor import BaseProcessor
from ..pipeline_data import PipelineData


class DrawDirectionProcessor(BaseProcessor):
    """行驶方向绘制处理器：基于YOLO跟踪结果绘制车辆行驶方向"""

    def __init__(self,
                 name: str = "行驶方向绘制处理器",
                 arrow_length: float = 20.0,
                 arrow_thickness: int = 2,
                 arrow_color: tuple = (0, 0, 255),  # 红色箭头
                 show_trajectory: bool = True,
                 trajectory_length: int = 10,  # 轨迹点数量
                 trajectory_color: tuple = (255, 255, 255),
                 trajectory_thickness: int = 1,
                 min_points_for_direction: int = 2,
                 min_move_distance: float = 3.0):  # 最小移动距离，避免抖动
        super().__init__(name)
        self.arrow_length = arrow_length
        self.arrow_thickness = arrow_thickness
        self.arrow_color = arrow_color
        self.show_trajectory = show_trajectory
        self.trajectory_length = trajectory_length
        self.trajectory_color = trajectory_color
        self.trajectory_thickness = trajectory_thickness
        self.min_points_for_direction = min_points_for_direction
        self.min_move_distance = min_move_distance

        # 存储每个track_id的轨迹历史 {track_id: deque([(x1, y1), (x2, y2), ...])}
        self.track_histories = {}

        # 只存储车辆中心点，不存储完整轨迹
        self.center_histories = {}

        # 清理超时的track（超过多少帧没出现就清理）
        self.max_track_lifetime = 30  # 30帧没出现就清理
        self.track_last_seen = {}  # {track_id: last_frame_idx}

    def process(self, data: PipelineData) -> PipelineData:
        """基于YOLO跟踪结果绘制车辆行驶方向"""
        if data.frame is None or data.current_result is None:
            return data

        vis = data.frame.copy()
        frame_idx = data.frame_idx

        # 获取YOLO检测结果
        boxes = getattr(data.current_result, "boxes", None)
        if boxes is None:
            # 没有检测框，清理过期track
            self._cleanup_old_tracks(frame_idx)
            data.frame = vis
            return data

        # 当前帧存在的track_id
        current_track_ids = set()

        # 处理每个检测框
        for box in boxes:
            try:
                # 解析检测框
                xyxy = getattr(box, "xyxy", None)
                if xyxy is None:
                    continue

                # 获取坐标
                x1, y1, x2, y2 = map(float, xyxy[0].tolist())

                # 获取track_id
                track_id = int(getattr(box, "id", -1))
                if track_id == -1:
                    continue  # 没有跟踪ID，跳过

                current_track_ids.add(track_id)

                # 计算中心点
                center_x = (x1 + x2) / 2.0
                center_y = (y1 + y2) / 2.0

                # 更新轨迹历史
                if track_id not in self.center_histories:
                    self.center_histories[track_id] = deque(maxlen=self.trajectory_length)

                # 添加当前中心点到历史
                self.center_histories[track_id].append((center_x, center_y))

                # 更新最后出现帧
                self.track_last_seen[track_id] = frame_idx

                # 绘制该车辆的轨迹和方向
                self._draw_track_info(vis, track_id, (center_x, center_y))

            except Exception as e:
                self.logger.error(f"Error processing box: {e}")
                continue

        # 清理长时间未出现的track
        self._cleanup_old_tracks(frame_idx)

        data.frame = vis
        return data

    def _draw_track_info(self, vis: np.ndarray, track_id: int, current_center: tuple):
        """绘制单个车辆的轨迹和方向"""
        if track_id not in self.center_histories:
            return

        history = self.center_histories[track_id]

        # 需要至少2个点才能绘制方向
        if len(history) < self.min_points_for_direction:
            return

        # 绘制轨迹
        if self.show_trajectory:
            self._draw_trajectory(vis, history)

        # 绘制方向箭头
        self._draw_direction_arrow(vis, history)

        # 绘制当前中心点
        cx, cy = current_center
        cv2.circle(vis, (int(cx), int(cy)), 3, (0, 255, 255), -1)  # 黄色点

        # 绘制track_id文本
        text = f"ID:{track_id}"
        cv2.putText(vis, text,
                    (int(cx) + 10, int(cy) - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                    (255, 255, 255), 1)

    def _draw_trajectory(self, vis: np.ndarray, history: deque):
        """绘制轨迹线"""
        if len(history) < 2:
            return

        # 将deque转换为列表以便索引
        points = list(history)

        # 绘制轨迹线
        for i in range(1, len(points)):
            pt1 = points[i - 1]
            pt2 = points[i]

            # 确保点是有效的
            if (isinstance(pt1, (tuple, list)) and len(pt1) >= 2 and
                    isinstance(pt2, (tuple, list)) and len(pt2) >= 2):
                x1, y1 = int(pt1[0]), int(pt1[1])
                x2, y2 = int(pt2[0]), int(pt2[1])

                # 绘制轨迹线
                cv2.line(vis, (x1, y1), (x2, y2),
                         self.trajectory_color, self.trajectory_thickness)

    def _draw_direction_arrow(self, vis: np.ndarray, history: deque):
        """绘制方向箭头"""
        if len(history) < 2:
            return

        # 获取最近的两个点计算方向
        points = list(history)
        current_point = points[-1]

        # 使用最近的两个点计算方向（避免抖动）
        start_idx = max(0, len(points) - 3)  # 使用最近3个点
        direction_points = points[start_idx:]

        if len(direction_points) < 2:
            return

        # 计算平均方向（首尾点之间的向量）
        start_point = direction_points[0]
        end_point = direction_points[-1]

        dx = end_point[0] - start_point[0]
        dy = end_point[1] - start_point[1]

        # 计算移动距离
        move_distance = math.hypot(dx, dy)

        # 如果移动距离太小，不绘制箭头（避免抖动）
        if move_distance < self.min_move_distance:
            return

        # 归一化方向向量
        if move_distance > 0:
            dx_norm = dx / move_distance
            dy_norm = dy / move_distance
        else:
            return

        # 计算箭头起点（当前点）
        sx, sy = current_point

        # 计算箭头终点
        ex = int(sx + dx_norm * self.arrow_length)
        ey = int(sy + dy_norm * self.arrow_length)

        # 绘制方向箭头
        cv2.arrowedLine(vis,
                        (int(sx), int(sy)),
                        (ex, ey),
                        self.arrow_color,
                        self.arrow_thickness,
                        tipLength=0.3)

    def _cleanup_old_tracks(self, current_frame_idx: int):
        """清理长时间未出现的track"""
        tracks_to_remove = []

        for track_id, last_seen in list(self.track_last_seen.items()):
            # 如果超过最大生命周期没出现，则清理
            if current_frame_idx - last_seen > self.max_track_lifetime:
                tracks_to_remove.append(track_id)

        for track_id in tracks_to_remove:
            if track_id in self.center_histories:
                del self.center_histories[track_id]
            if track_id in self.track_last_seen:
                del self.track_last_seen[track_id]