import cv2
import numpy as np
import glob

# 图片储存的路径
img_path = "/home/sunrise/yiliao_ws/imagebag_outputs"

# 棋盘格参数
CHECKERBOARD = (8, 5)  # 内角点数量（列×行）
square_size = 27  # 棋盘格方块实际尺寸（单位：毫米）

# 存储3D和2D点
objpoints = []  # 世界坐标系中的点
imgpoints = []  # 图像坐标系中的点

# 生成世界坐标系中的角点坐标（z=0）
objp = np.zeros((CHECKERBOARD[0]*CHECKERBOARD[1], 3), np.float32)
objp[:, :2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2) * square_size


# 读取所有棋盘格图片
images = glob.glob(f'{img_path}/*.png')

for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 检测角点
    ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD, None)

    if ret:
        objpoints.append(objp)

        # 亚像素角点优化
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        corners2 = cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria)
        imgpoints.append(corners2)

        # 绘制角点（可选）
        # cv2.drawChessboardCorners(img, CHECKERBOARD, corners2, ret)
        # cv2.imshow('Detected Corners', img)
        # cv2.waitKey(500)

# cv2.destroyAllWindows()

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

# 输出标定结果
print("相机内参矩阵:\n", mtx)
print("\n畸变系数:", dist.ravel())

# 计算重投影误差
mean_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
    error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
    mean_error += error
print("\n平均重投影误差:", mean_error / len(objpoints))

# 保存参数
np.savez('cali.npz', mtx=mtx, dist=dist)


# 测试校正效果
img = cv2.imread(fname)  # 替换为实际的图片路径
# ret, img = cap.read()
h, w = img.shape[:2]
newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w,h), 1, (w,h))