import cv2
import yaml
import numpy as np

# 定义函数读取标定数据
def read_calibration_data(calibration_file):
    with open(calibration_file, 'r') as f:
        calib_data = yaml.safe_load(f)
        cameraMatrix_l = np.array(calib_data['camera_matrix_left']['data']).reshape(3, 3)
        distCoeffs_l = np.array(calib_data['dist_coeff_left']['data'])
        cameraMatrix_r = np.array(calib_data['camera_matrix_right']['data']).reshape(3, 3)
        distCoeffs_r = np.array(calib_data['dist_coeff_right']['data'])
        R = np.array(calib_data['R']['data']).reshape(3, 3)
        T = np.array(calib_data['T']['data']).reshape(3, 1)
    return cameraMatrix_l, distCoeffs_l, cameraMatrix_r, distCoeffs_r, R, T

# 定义函数对图像进行矫正
def rectify_images(left_image_path, right_image_path, calibration_file):
    # 读取标定数据
    cameraMatrix_l, distCoeffs_l, cameraMatrix_r, distCoeffs_r, R, T = read_calibration_data(calibration_file)

    # 读取左右图像
    img_left = cv2.imread(left_image_path)
    img_right = cv2.imread(right_image_path)

    # 获取图像尺寸（假设左右图像尺寸相同）
    img_size = img_left.shape[:2][::-1]

    # 立体校正
    R1, R2, P1, P2, Q, roi1, roi2 = cv2.stereoRectify(cameraMatrix_l, distCoeffs_l,
                                                     cameraMatrix_r, distCoeffs_r,
                                                     img_size, R, T)

    # 计算映射参数
    map1_l, map2_l = cv2.initUndistortRectifyMap(cameraMatrix_l, distCoeffs_l, R1, P1, img_size, cv2.CV_32FC1)
    map1_r, map2_r = cv2.initUndistortRectifyMap(cameraMatrix_r, distCoeffs_r, R2, P2, img_size, cv2.CV_32FC1)

    # 应用映射并显示结果
    rectified_img_l = cv2.remap(img_left, map1_l, map2_l, cv2.INTER_LINEAR)
    rectified_img_r = cv2.remap(img_right, map1_r, map2_r, cv2.INTER_LINEAR)

    # 合并图像显示
    combined_img = np.hstack((rectified_img_l, rectified_img_r))
    cv2.imshow('Rectified Images', combined_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# 设置路径和文件名
left_image_path = "left/left_WIN_20241023_14_54_55_Pro.jpg"
right_image_path = "right/right_WIN_20241023_14_54_55_Pro.jpg"
calibration_file = "calibration_parameters.yaml"

# 调用函数进行图像矫正
rectify_images(left_image_path, right_image_path, calibration_file)