FootBall-route-drawing-and-modeling/backend/utils/field_transformer.py

import cv2
import numpy as np

class FieldTransformer:
    """用于将视频中的像素坐标转换为场地坐标的工具类"""
    
    def __init__(self):
        # 标准橄榄球场尺寸（单位：米）
        self.field_length = 100  # 大约100码（不包括端区）
        self.field_width = 50    # 大约53.3码
        
        # 默认变换矩阵（假设俯视图）
        self.perspective_matrix = None
        
        # 是否已自动检测场地
        self.field_detected = False
        
        # 场地边界点（图像坐标系）
        self.field_corners = None
    
    def auto_detect_field(self, frame):
        """尝试自动检测球场区域"""
        try:
            # 转换为灰度图
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 应用高斯模糊
            blur = cv2.GaussianBlur(gray, (5, 5), 0)
            
            # 边缘检测
            edges = cv2.Canny(blur, 50, 150)
            
            # 膨胀边缘以连接间隙
            dilated = cv2.dilate(edges, np.ones((3, 3), np.uint8), iterations=2)
            
            # 查找轮廓
            contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            # 按面积排序
            contours = sorted(contours, key=cv2.contourArea, reverse=True)
            
            # 尝试找到矩形区域（假设场地为矩形）
            for contour in contours[:5]:  # 只检查最大的5个轮廓
                # 近似轮廓
                epsilon = 0.02 * cv2.arcLength(contour, True)
                approx = cv2.approxPolyDP(contour, epsilon, True)
                
                # 如果近似为四边形，可能是场地
                if len(approx) == 4:
                    # 进一步处理，确认是否为场地
                    field_area = cv2.contourArea(approx)
                    frame_area = frame.shape[0] * frame.shape[1]
                    
                    # 如果面积占总面积的比例合适，则认为找到了场地
                    if 0.2 < field_area / frame_area < 0.95:
                        # 存储场地角点
                        self.field_corners = approx.reshape(4, 2)
                        
                        # 根据场地角点计算透视变换矩阵
                        self._calculate_perspective_matrix()
                        
                        self.field_detected = True
                        return True
            
            # 未找到合适的场地轮廓，使用默认变换
            self._setup_default_transform(frame)
            return False
            
        except Exception as e:
            print(f"场地自动检测失败: {str(e)}")
            self._setup_default_transform(frame)
            return False
    
    def _setup_default_transform(self, frame):
        """设置默认变换矩阵（假设俯视图）"""
        height, width = frame.shape[:2]
        
        # 默认场地框（图像四角）
        self.field_corners = np.array([
            [0.1 * width, 0.1 * height],     # 左上
            [0.9 * width, 0.1 * height],     # 右上
            [0.9 * width, 0.9 * height],     # 右下
            [0.1 * width, 0.9 * height]      # 左下
        ], dtype=np.float32)
        
        self._calculate_perspective_matrix()
    
    def _calculate_perspective_matrix(self):
        """计算透视变换矩阵"""
        # 标准化角点顺序：左上，右上，右下，左下
        rect = self._order_points(self.field_corners)
        
        # 目标场地坐标（标准场地，单位：米）
        dst = np.array([
            [0, 0],                                  # 左上
            [self.field_length, 0],                  # 右上
            [self.field_length, self.field_width],   # 右下
            [0, self.field_width]                    # 左下
        ], dtype=np.float32)
        
        # 计算透视变换矩阵
        self.perspective_matrix = cv2.getPerspectiveTransform(rect, dst)
    
    def _order_points(self, pts):
        """对四边形角点进行排序: 左上，右上，右下，左下"""
        # 初始化结果数组
        rect = np.zeros((4, 2), dtype=np.float32)
        
        # 计算点的和，左上角的和最小，右下角的和最大
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]  # 左上
        rect[2] = pts[np.argmax(s)]  # 右下
        
        # 计算差，右上角的差最小，左下角的差最大
        diff = np.diff(pts, axis=1)
        rect[1] = pts[np.argmin(diff)]  # 右上
        rect[3] = pts[np.argmax(diff)]  # 左下
        
        return rect
    
    def transform_to_field(self, tracks, frame=None):
        """将球员位置从像素坐标转换到场地坐标
        
        Args:
            tracks: {id: (x, y, w, h)} 形式的字典，x,y是中心点
            frame: 可选的原始视频帧，用于自动检测场地
            
        Returns:
            {id: (field_x, field_y)} 形式的字典
        """
        # 如果没有初始化透视矩阵且提供了frame
        if self.perspective_matrix is None and frame is not None:
            self.auto_detect_field(frame)
        
        # 如果仍未初始化，使用默认值
        if self.perspective_matrix is None:
            # 创建一个假想的frame
            mock_frame = np.zeros((720, 1280, 3), dtype=np.uint8)
            self._setup_default_transform(mock_frame)
        
        # 转换每个轨迹到场地坐标
        field_positions = {}
        
        for track_id, (x, y, w, h) in tracks.items():
            # 转换中心点
            pt = np.array([[x, y]], dtype=np.float32)
            pt = cv2.perspectiveTransform(pt.reshape(-1, 1, 2), self.perspective_matrix)
            
            # 存储场地坐标
            field_positions[track_id] = (float(pt[0][0][0]), float(pt[0][0][1]))
        
        return field_positions
    
    def convert_field_to_image(self, field_positions, frame_shape):
        """将场地坐标转换回图像坐标
        
        Args:
            field_positions: {id: (field_x, field_y)} 形式的字典
            frame_shape: 原始帧的形状(高度, 宽度)
            
        Returns:
            {id: (image_x, image_y)} 形式的字典
        """
        if self.perspective_matrix is None:
            # 创建一个假想的frame
            mock_frame = np.zeros((frame_shape[0], frame_shape[1], 3), dtype=np.uint8)
            self._setup_default_transform(mock_frame)
        
        # 计算逆变换矩阵
        inv_matrix = np.linalg.inv(self.perspective_matrix)
        
        # 转换每个场地坐标到图像坐标
        image_positions = {}
        
        for track_id, (field_x, field_y) in field_positions.items():
            # 转换点
            pt = np.array([[[field_x, field_y]]], dtype=np.float32)
            pt = cv2.perspectiveTransform(pt, inv_matrix)
            
            # 存储图像坐标
            image_positions[track_id] = (float(pt[0][0][0]), float(pt[0][0][1]))
        
        return image_positions
    
    def visualize_field_transform(self, frame):
        """在帧上可视化场地变换"""
        if self.field_corners is None:
            self.auto_detect_field(frame)
            
        vis_frame = frame.copy()
        
        # 绘制检测到的场地边界
        if self.field_corners is not None:
            # 将角点转换为整数坐标
            corners = self.field_corners.astype(np.int32)
            
            # 绘制场地边界
            cv2.polylines(vis_frame, [corners], True, (0, 255, 0), 2)
            
            # 标记角点
            for i, (x, y) in enumerate(corners):
                cv2.circle(vis_frame, (x, y), 5, (0, 0, 255), -1)
                cv2.putText(vis_frame, str(i), (x-10, y-10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
            
            # 添加说明文本
            cv2.putText(vis_frame, "检测到场地", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        else:
            cv2.putText(vis_frame, "未检测到场地，使用默认变换", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        
        return vis_frame