archery/aruco_detector.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ArUco标记检测模块
提供基于ArUco标记的靶心标定和激光点定位功能
"""
import cv2
import numpy as np
import math
import config
from logger_manager import logger_manager


class ArUcoDetector:
    """ArUco标记检测器"""
    
    def __init__(self):
        self.logger = logger_manager.logger
        # 创建ArUco字典和检测器参数
        self.aruco_dict = cv2.aruco.getPredefinedDictionary(config.ARUCO_DICT_TYPE)
        self.detector_params = cv2.aruco.DetectorParameters()
        
        # 设置检测参数
        self.detector_params.minMarkerPerimeterRate = config.ARUCO_MIN_MARKER_PERIMETER_RATE
        self.detector_params.cornerRefinementMethod = config.ARUCO_CORNER_REFINEMENT_METHOD
        
        # 创建检测器
        self.detector = cv2.aruco.ArucoDetector(self.aruco_dict, self.detector_params)
        
        # 预定义靶纸上的标记位置（物理坐标，毫米）
        self.marker_positions_mm = config.ARUCO_MARKER_POSITIONS_MM
        self.marker_ids = config.ARUCO_MARKER_IDS
        self.marker_size_mm = config.ARUCO_MARKER_SIZE_MM
        self.target_paper_size_mm = config.TARGET_PAPER_SIZE_MM
        
        # 靶心偏移（相对于靶纸中心）
        self.target_center_offset_mm = config.TARGET_CENTER_OFFSET_MM
        
        if self.logger:
            self.logger.info(f"[ARUCO] ArUco检测器初始化完成，字典类型: {config.ARUCO_DICT_TYPE}")
    
    def detect_markers(self, frame):
        """
        检测图像中的ArUco标记
        
        Args:
            frame: MaixPy图像帧对象
            
        Returns:
            (corners, ids, rejected) - 检测到的标记角点、ID列表、被拒绝的候选
            如果检测失败返回 (None, None, None)
        """
        try:
            # 转换为OpenCV格式
            from maix import image
            img_cv = image.image2cv(frame, False, False)
            
            # 转换为灰度图（ArUco检测需要）
            if len(img_cv.shape) == 3:
                gray = cv2.cvtColor(img_cv, cv2.COLOR_RGB2GRAY)
            else:
                gray = img_cv
            
            # 检测标记
            corners, ids, rejected = self.detector.detectMarkers(gray)
            
            if self.logger and ids is not None:
                self.logger.debug(f"[ARUCO] 检测到 {len(ids)} 个标记: {ids.flatten().tolist()}")
            
            return corners, ids, rejected
            
        except Exception as e:
            if self.logger:
                self.logger.error(f"[ARUCO] 标记检测失败: {e}")
            return None, None, None
    
    def get_target_center_from_markers(self, corners, ids):
        """
        从检测到的ArUco标记计算靶心位置
        
        Args:
            corners: 标记角点列表
            ids: 标记ID列表
            
        Returns:
            (center_x, center_y, radius, ellipse_params) 或 (None, None, None, None)
            center_x, center_y: 靶心像素坐标
            radius: 估计的靶心半径（像素）
            ellipse_params: 椭圆参数用于透视校正
        """
        if ids is None or len(ids) < 3:
            if self.logger:
                self.logger.debug(f"[ARUCO] 检测到的标记数量不足: {len(ids) if ids is not None else 0} < 3")
            return None, None, None, None
        
        try:
            # 将ID转换为列表便于查找
            detected_ids = ids.flatten().tolist()
            
            # 收集检测到的标记中心点和对应的物理坐标
            image_points = []  # 图像坐标 (像素)
            object_points = []  # 物理坐标 (毫米)
            marker_centers = {}  # 存储每个标记的中心
            
            for i, marker_id in enumerate(detected_ids):
                if marker_id not in self.marker_ids:
                    continue
                
                # 计算标记中心（四个角的平均值）
                corner = corners[i][0]  # shape: (4, 2)
                center_x = np.mean(corner[:, 0])
                center_y = np.mean(corner[:, 1])
                marker_centers[marker_id] = (center_x, center_y)
                
                # 添加到点列表
                image_points.append([center_x, center_y])
                object_points.append(self.marker_positions_mm[marker_id])
            
            if len(image_points) < 3:
                if self.logger:
                    self.logger.debug(f"[ARUCO] 有效标记数量不足: {len(image_points)} < 3")
                return None, None, None, None
            
            # 转换为numpy数组
            image_points = np.array(image_points, dtype=np.float32)
            object_points = np.array(object_points, dtype=np.float32)
            
            # 计算单应性矩阵（Homography）
            # 这建立了物理坐标到图像坐标的映射
            H, status = cv2.findHomography(object_points, image_points, cv2.RANSAC, 5.0)
            
            if H is None:
                if self.logger:
                    self.logger.warning("[ARUCO] 无法计算单应性矩阵")
                return None, None, None, None
            
            # 计算靶心在图像中的位置
            # 靶心物理坐标 = 靶纸中心 + 偏移
            target_center_mm = np.array([[self.target_center_offset_mm[0], 
                                          self.target_center_offset_mm[1]]], dtype=np.float32)
            target_center_mm = target_center_mm.reshape(-1, 1, 2)
            
            # 使用单应性矩阵投影到图像坐标
            target_center_img = cv2.perspectiveTransform(target_center_mm, H)
            center_x = target_center_img[0][0][0]
            center_y = target_center_img[0][0][1]
            
            # 计算靶心半径（像素）
            # 使用已知物理距离和像素距离的比例
            # 选择两个标记计算比例尺
            if len(marker_centers) >= 2:
                # 使用对角线上的标记计算比例尺
                if 0 in marker_centers and 2 in marker_centers:
                    p1_img = np.array(marker_centers[0])
                    p2_img = np.array(marker_centers[2])
                    p1_mm = np.array(self.marker_positions_mm[0])
                    p2_mm = np.array(self.marker_positions_mm[2])
                elif 1 in marker_centers and 3 in marker_centers:
                    p1_img = np.array(marker_centers[1])
                    p2_img = np.array(marker_centers[3])
                    p1_mm = np.array(self.marker_positions_mm[1])
                    p2_mm = np.array(self.marker_positions_mm[3])
                else:
                    # 使用任意两个标记
                    keys = list(marker_centers.keys())
                    p1_img = np.array(marker_centers[keys[0]])
                    p2_img = np.array(marker_centers[keys[1]])
                    p1_mm = np.array(self.marker_positions_mm[keys[0]])
                    p2_mm = np.array(self.marker_positions_mm[keys[1]])
                
                pixel_distance = np.linalg.norm(p1_img - p2_img)
                mm_distance = np.linalg.norm(p1_mm - p2_mm)
                
                if mm_distance > 0:
                    pixels_per_mm = pixel_distance / mm_distance
                    # 标准靶心半径：10环半径约1.22cm = 12.2mm
                    # 但这里我们返回一个估计值，实际环数计算在laser_manager中
                    radius_mm = 122.0  # 整个靶纸的半径约200mm，但靶心区域较小
                    radius = int(radius_mm * pixels_per_mm)
                else:
                    radius = 100  # 默认值
            else:
                radius = 100  # 默认值
            
            # 计算椭圆参数（用于透视校正）
            # 从单应性矩阵可以推导出透视变形
            ellipse_params = self._compute_ellipse_params(H, center_x, center_y)
            
            if self.logger:
                self.logger.info(f"[ARUCO] 靶心计算成功: 中心=({center_x:.1f}, {center_y:.1f}), "
                               f"半径={radius}px, 检测到{len(marker_centers)}个标记")
            
            return (int(center_x), int(center_y)), radius, "aruco", ellipse_params
            
        except Exception as e:
            if self.logger:
                self.logger.error(f"[ARUCO] 计算靶心失败: {e}")
                import traceback
                self.logger.error(traceback.format_exc())
            return None, None, None, None
    
    def _compute_ellipse_params(self, H, center_x, center_y):
        """
        从单应性矩阵计算椭圆参数，用于透视校正
        
        Args:
            H: 单应性矩阵 (3x3)
            center_x, center_y: 靶心图像坐标
            
        Returns:
            ellipse_params: ((center_x, center_y), (width, height), angle)
        """
        try:
            # 在物理坐标系中画一个圆，投影到图像中看变成什么形状
            # 物理圆：半径10mm
            r_mm = 10.0
            angles = np.linspace(0, 2*np.pi, 16)
            circle_mm = np.array([[self.target_center_offset_mm[0] + r_mm * np.cos(a),
                                   self.target_center_offset_mm[1] + r_mm * np.sin(a)] 
                                  for a in angles], dtype=np.float32)
            circle_mm = circle_mm.reshape(-1, 1, 2)
            
            # 投影到图像
            circle_img = cv2.perspectiveTransform(circle_mm, H)
            circle_img = circle_img.reshape(-1, 2)
            
            # 拟合椭圆
            if len(circle_img) >= 5:
                ellipse = cv2.fitEllipse(circle_img.astype(np.float32))
                return ellipse
            else:
                # 从单应性矩阵近似估计
                # 提取缩放和旋转
                # H = K * [R|t] 的近似
                # 这里简化处理：假设没有严重变形
                scale_x = np.linalg.norm(H[0, :2])
                scale_y = np.linalg.norm(H[1, :2])
                avg_scale = (scale_x + scale_y) / 2
                
                width = r_mm * 2 * scale_x
                height = r_mm * 2 * scale_y
                angle = np.degrees(np.arctan2(H[1, 0], H[0, 0]))
                
                return ((center_x, center_y), (width, height), angle)
                
        except Exception as e:
            if self.logger:
                self.logger.debug(f"[ARUCO] 计算椭圆参数失败: {e}")
            return None
    
    def transform_laser_point(self, laser_point, corners, ids):
        """
        将激光点从图像坐标转换到物理坐标（毫米），再计算相对于靶心的偏移
        
        Args:
            laser_point: (x, y) 激光点在图像中的坐标
            corners: 检测到的标记角点
            ids: 检测到的标记ID
            
        Returns:
            (dx_mm, dy_mm) 激光点相对于靶心的偏移（毫米），或 (None, None)
        """
        if laser_point is None or ids is None or len(ids) < 3:
            return None, None
        
        try:
            # 重新计算单应性矩阵（可以优化为缓存）
            detected_ids = ids.flatten().tolist()
            image_points = []
            object_points = []
            
            for i, marker_id in enumerate(detected_ids):
                if marker_id not in self.marker_ids:
                    continue
                corner = corners[i][0]
                center_x = np.mean(corner[:, 0])
                center_y = np.mean(corner[:, 1])
                image_points.append([center_x, center_y])
                object_points.append(self.marker_positions_mm[marker_id])
            
            if len(image_points) < 3:
                return None, None
            
            image_points = np.array(image_points, dtype=np.float32)
            object_points = np.array(object_points, dtype=np.float32)
            
            H, _ = cv2.findHomography(object_points, image_points, cv2.RANSAC, 5.0)
            if H is None:
                return None, None
            
            # 求逆矩阵，将图像坐标转换到物理坐标
            H_inv = np.linalg.inv(H)
            
            laser_img = np.array([[laser_point[0], laser_point[1]]], dtype=np.float32)
            laser_img = laser_img.reshape(-1, 1, 2)
            
            laser_mm = cv2.perspectiveTransform(laser_img, H_inv)
            laser_x_mm = laser_mm[0][0][0]
            laser_y_mm = laser_mm[0][0][1]
            
            # 计算相对于靶心的偏移
            # 注意：Y轴方向可能需要翻转（图像Y向下，物理Y通常向上）
            dx_mm = laser_x_mm - self.target_center_offset_mm[0]
            dy_mm = -(laser_y_mm - self.target_center_offset_mm[1])  # 翻转Y轴
            
            if self.logger:
                self.logger.debug(f"[ARUCO] 激光点转换: 图像({laser_point[0]:.1f}, {laser_point[1]:.1f}) -> "
                                f"物理({laser_x_mm:.1f}, {laser_y_mm:.1f}) -> "
                                f"偏移({dx_mm:.1f}, {dy_mm:.1f})mm")
            
            return dx_mm, dy_mm
            
        except Exception as e:
            if self.logger:
                self.logger.error(f"[ARUCO] 激光点转换失败: {e}")
            return None, None
    
    def draw_debug_info(self, frame, corners, ids, target_center=None, laser_point=None):
        """
        在图像上绘制调试信息
        
        Args:
            frame: MaixPy图像帧
            corners: 标记角点
            ids: 标记ID
            target_center: 计算的靶心位置
            laser_point: 激光点位置
            
        Returns:
            绘制后的图像
        """
        try:
            from maix import image
            img_cv = image.image2cv(frame, False, False).copy()
            
            # 绘制检测到的标记
            if ids is not None:
                cv2.aruco.drawDetectedMarkers(img_cv, corners, ids)
                
                # 绘制标记ID和中心
                for i, marker_id in enumerate(ids.flatten()):
                    corner = corners[i][0]
                    center_x = int(np.mean(corner[:, 0]))
                    center_y = int(np.mean(corner[:, 1]))
                    
                    # 绘制中心点
                    cv2.circle(img_cv, (center_x, center_y), 5, (0, 255, 0), -1)
                    
                    # 绘制ID
                    cv2.putText(img_cv, f"ID:{marker_id}", 
                               (center_x + 10, center_y - 10),
                               cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            
            # 绘制靶心
            if target_center:
                cv2.circle(img_cv, target_center, 8, (255, 0, 0), -1)
                cv2.circle(img_cv, target_center, 50, (255, 0, 0), 2)
                cv2.putText(img_cv, "TARGET", (target_center[0] + 15, target_center[1] - 15),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
            
            # 绘制激光点
            if laser_point:
                cv2.circle(img_cv, (int(laser_point[0]), int(laser_point[1])), 6, (0, 0, 255), -1)
                cv2.putText(img_cv, "LASER", (int(laser_point[0]) + 10, int(laser_point[1]) - 10),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
            
            # 转换回MaixPy图像
            return image.cv2image(img_cv, False, False)
            
        except Exception as e:
            if self.logger:
                self.logger.error(f"[ARUCO] 绘制调试信息失败: {e}")
            return frame


# 创建全局单例实例
aruco_detector = ArUcoDetector()


def detect_target_with_aruco(frame, laser_point=None):
    """
    使用ArUco标记检测靶心的便捷函数
    
    Args:
        frame: MaixPy图像帧
        laser_point: 激光点坐标（可选）
        
    Returns:
        (result_img, center, radius, method, best_radius1, ellipse_params)
        与detect_circle_v3保持相同的返回格式
    """
    detector = aruco_detector
    
    # 检测ArUco标记
    corners, ids, rejected = detector.detect_markers(frame)
    
    # 计算靶心
    center, radius, method, ellipse_params = detector.get_target_center_from_markers(corners, ids)
    
    # 绘制调试信息
    result_img = detector.draw_debug_info(frame, corners, ids, center, laser_point)
    
    # 返回与detect_circle_v3相同的格式
    # best_radius1用于距离估算，这里用radius代替
    return result_img, center, radius, method, radius, ellipse_params


def compute_laser_offset_aruco(laser_point, corners, ids):
    """
    使用ArUco计算激光点相对于靶心的偏移（毫米）
    
    Args:
        laser_point: (x, y) 激光点图像坐标
        corners: ArUco标记角点
        ids: ArUco标记ID
        
    Returns:
        (dx_mm, dy_mm) 偏移量（毫米），或 (None, None)
    """
    return aruco_detector.transform_laser_point(laser_point, corners, ids)