import cv2
import numpy as np

# 准备棋盘格参数
pattern_size = (9, 6)  # 9×6的棋盘格
square_size = 25       # 方格大小25mm

# 存储角点坐标
obj_points = []  # 世界坐标系中的3D点
img_points = []  # 图像坐标系中的2D点

# 读取标定图像
images = ['calib1.jpg', 'calib2.jpg', ...]

for img_file in images:
    img = cv2.imread(img_file)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 查找棋盘格角点
    ret, corners = cv2.findChessboardCorners(gray, pattern_size)
    
    if ret:
        obj_points.append(objp)  # 添加世界坐标点
        img_points.append(corners)  # 添加图像坐标点
        
        # 可视化角点
        cv2.drawChessboardCorners(img, pattern_size, corners, ret)
        cv2.imshow('Corners', img)
        cv2.waitKey(500)

# 执行标定
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(
    obj_points, img_points, gray.shape[::-1], None, None
)

print("相机内参矩阵:")
print(mtx)
print("畸变系数:")
print(dist)

import cv2
import apriltag
import numpy as np
import os

class AprilTagCalibrator:
    def __init__(self, tag_size_mm=40, spacing_mm=20, grid_size=(4, 4)):
        self.tag_size_mm = tag_size_mm
        self.spacing_mm = spacing_mm
        self.grid_size = grid_size
        
        # 计算每个Tag中心的世界坐标
        self.tag_centers = self.calculate_tag_centers()
        
        # 创建检测器
        self.detector = apriltag.Detector(apriltag.DetectorOptions(families="tag36h11"))
    
    def calculate_tag_centers(self):
        """计算每个Tag中心的世界坐标"""
        tag_centers = {}
        
        for i in range(self.grid_size[0]):
            for j in range(self.grid_size[1]):
                tag_id = i * self.grid_size[1] + j
                
                x = j * (self.tag_size_mm + self.spacing_mm)
                y = i * (self.tag_size_mm + self.spacing_mm)
                z = 0
                
                tag_centers[tag_id] = [x, y, z]
        
        # 转换为相对于中心的坐标
        center_x = (self.grid_size[1] - 1) * (self.tag_size_mm + self.spacing_mm) / 2
        center_y = (self.grid_size[0] - 1) * (self.tag_size_mm + self.spacing_mm) / 2
        
        for tag_id in tag_centers:
            tag_centers[tag_id][0] -= center_x
            tag_centers[tag_id][1] -= center_y
        
        return tag_centers
    
    def get_tag_corners_3d(self, tag_center):
        """获取Tag四个角点的世界坐标"""
        cx, cy, cz = tag_center
        
        corner_offsets = [
            [-self.tag_size_mm/2, -self.tag_size_mm/2, 0],
            [ self.tag_size_mm/2, -self.tag_size_mm/2, 0],
            [ self.tag_size_mm/2,  self.tag_size_mm/2, 0],
            [-self.tag_size_mm/2,  self.tag_size_mm/2, 0]
        ]
        
        corners_3d = []
        for offset in corner_offsets:
            corners_3d.append([
                cx + offset[0],
                cy + offset[1],
                cz + offset[2]
            ])
        
        return corners_3d
    
    def process_calibration_images(self, image_paths):
        """处理所有标定图像"""
        all_object_points = []  # 3D世界点列表
        all_image_points = []   # 2D像素点列表
        valid_images = 0
        
        for img_path in image_paths:
            print(f"正在处理: {img_path}")
            
            image = cv2.imread(img_path)
            if image is None:
                print(f"跳过: 无法读取图像 {img_path}")
                continue
            
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            detections = self.detector.detect(gray)
            
            if len(detections) < 4:  # 至少检测到4个Tag
                print(f"跳过: 检测到 {len(detections)} 个Tag，不足4个")
                continue
            
            object_points = []
            image_points = []
            
            for detection in detections:
                tag_id = detection.tag_id
                
                if tag_id not in self.tag_centers:
                    continue
                
                # 获取该Tag的四个角点的世界坐标
                tag_corners_3d = self.get_tag_corners_3d(self.tag_centers[tag_id])
                
                # 获取该Tag的四个角点的图像坐标
                tag_corners_2d = np.array(detection.corners, dtype=np.float32)
                
                # 添加到当前图像的对应关系中
                object_points.extend(tag_corners_3d)
                image_points.extend(tag_corners_2d)
            
            if len(object_points) >= 8:  # 至少8个点才有效
                all_object_points.append(np.array(object_points))
                all_image_points.append(np.array(image_points))
                valid_images += 1
                print(f"✓ 处理成功，检测到 {len(object_points)//4} 个Tag")
            else:
                print("跳过: 对应点不足")
        
        print(f"\n共处理 {len(image_paths)} 张图像，有效 {valid_images} 张")
        
        return all_object_points, all_image_points
    
    def calibrate_camera(self, image_paths):
        """执行相机标定"""
        # 获取3D-2D对应关系
        object_points, image_points = self.process_calibration_images(image_paths)
        
        if len(object_points) == 0:
            raise ValueError("没有有效的标定图像")
        
        # 获取图像尺寸
        first_image = cv2.imread(image_paths[0])
        image_size = (first_image.shape[1], first_image.shape[0])
        
        # 执行标定
        ret, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(
            object_points,     # 3D世界点列表
            image_points,      # 2D像素点列表
            image_size,        # 图像分辨率
            None, None         # 内参和畸变系数（None表示需要计算）
        )
        
        return {
            'reprojection_error': ret,
            'camera_matrix': camera_matrix,
            'distortion_coefficients': dist_coeffs,
            'rotation_vectors': rvecs,
            'translation_vectors': tvecs,
            'image_size': image_size
        }

# 使用示例
if __name__ == "__main__":
    # 初始化标定器
    calibrator = AprilTagCalibrator(tag_size_mm=40, spacing_mm=20, grid_size=(4, 4))
    
    # 设置图像路径
    image_paths = [
        "calib_001.jpg",
        "calib_002.jpg",
        "calib_003.jpg",
        # ... 添加更多图像
    ]
    
    try:
        # 执行标定
        calibration_result = calibrator.calibrate_camera(image_paths)
        
        # 输出结果
        print("\n" + "="*50)
        print("✅ 相机标定完成！")
        print("="*50)
        
        print(f"重投影误差: {calibration_result['reprojection_error']:.4f} 像素")
        print(f"图像尺寸: {calibration_result['image_size']}")
        
        print("\n� 相机内参矩阵:")
        print(calibration_result['camera_matrix'])
        
        print("\n� 畸变系数:")
        print(calibration_result['distortion_coefficients'])
        
        print("\n� 旋转向量 (每张图像):")
        for i, rvec in enumerate(calibration_result['rotation_vectors']):
            print(f"  图像 {i+1}: {rvec.flatten()}")
        
        print("\n� 平移向量 (每张图像):")
        for i, tvec in enumerate(calibration_result['translation_vectors']):
            print(f"  图像 {i+1}: {tvec.flatten()}")
        
        # 保存结果
        np.savez('camera_calibration.npz',
                camera_matrix=calibration_result['camera_matrix'],
                dist_coeffs=calibration_result['distortion_coefficients'],
                reprojection_error=calibration_result['reprojection_error'])
        
        print("\n� 标定结果已保存到 camera_calibration.npz")
        
    except Exception as e:
        print(f"❌ 标定失败: {e}")