import cv2
import numpy as np

# 设置棋盘格的大小
chessboard_size = (6, 9)

# 设置棋盘格的物理尺寸
square_size = 4.0  # 单位为厘米

# 准备用于标定的物体空间坐标和对应的图像坐标
objpoints = []  # 3D points in real world space
imgpoints = []  # 2D points in image plane.

# 准备用于标定的图像列表
images = []

# 加载图像并查找棋盘格角点
for i in range(1, 7):  # 假设有20张图像
    filename = f'data/iphone_0.5/{i:02d}.jpg'
    img = cv2.imread(filename)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 在灰度图上查找棋盘格角点
    ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

    # 如果找到角点，则添加到标定数据中
    if ret:
        objp = np.zeros((np.prod(chessboard_size), 3), np.float32)
        objp[:, :2] = np.indices(chessboard_size).T.reshape(-1, 2) * square_size
        objpoints.append(objp)
        imgpoints.append(corners)

        # 在图像上绘制角点并添加到标定数据中
        cv2.drawChessboardCorners(img, chessboard_size, corners, ret)
        images.append(img)

# 进行相机标定
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)

# 打印内部参数矩阵和畸变系数
print("Camera Matrix:")
print(mtx)
print("\nDistortion Coefficients:")
print(dist)

# 显示标定结果图像
cv2.imshow('Chessboard Corners', images[0])
cv2.waitKey(0)
cv2.destroyAllWindows()