pmd-python/unit_test.py

import cv2
import os
from tqdm import tqdm
import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt
from scipy.interpolate import griddata
import glob
from scipy import stats
from scipy.optimize import minimize
from collections import defaultdict
from scipy.signal import savgol_filter
from scipy.ndimage import uniform_filter, gaussian_filter, median_filter
import json
import pickle
import shutil
import pandas as pd
from sklearn.linear_model import LinearRegression
from scipy.interpolate import Rbf
from numpy.linalg import lstsq
from scipy.optimize import curve_fit
from aprilgrid import Detector
from src.utils import *
from src.calibration.get_camera_params import generate_3d_points, show_detect_result
from scipy.optimize import least_squares
import random
from scipy.io import savemat, loadmat
from src.pcl_postproc import smooth_pcl, align2ref



def pixel_to_world(points_2d, K, R, T, depth):
    """
    将像素坐标转换为世界坐标系中的 3D 坐标。
    
    Parameters:
    points_2d: np.ndarray
        图像中的2D像素坐标 (N, 2)。
    K: np.ndarray
        相机内参矩阵 (3x3)。
    R: np.ndarray
        相机旋转矩阵 (3x3)。
    T: np.ndarray
        相机平移向量 (3, 1)。
    depth: float or np.ndarray
        点的深度（单位：毫米或其他实际单位）。
        
    Returns:
    points_3d_world: np.ndarray
        世界坐标系中的 3D 点 (N, 3)。
    """
    # 将像素坐标转换为齐次坐标 (u, v, 1)
    image_points_3d = np.hstack((points_2d, np.ones((points_2d.shape[0], 1))))
    norm_points = np.linalg.inv(K).dot(image_points_3d.reshape(-1, 3).T).T
    # 逆投影到3D空间
    object_points = []
    for point in norm_points:
    # 将归一化坐标转换为3D点
        point_3d = np.zeros((3,), dtype=np.float32)
        point_3d[:2] = point[:2] / point[2]  # 归一化坐标转换为3D坐标
        point_3d[2] = 1

        # 使用旋转和平移向量进行逆投影
        # 首先需要将旋转向量转换为旋转矩阵
        #R, _ = cv2.Rodrigues(R.reshape(3, 3))
        object_point = np.dot(R.T, point_3d - T.reshape(3, 1)).flatten()

        object_points.append(object_point)

  
    points_3d_world = np.array(object_points)

    print("Object Points:")
    print(points_3d_world)

    return points_3d_world


def project_to_world(points_2d, K, R, T, dist_coeffs, depth=1.0):
    """
    将相机检测到的2D点转换为世界坐标系下的3D点。
    
    Parameters:
    points_2d: np.ndarray
        图像中的2D点 (N, 2)。
    K: np.ndarray
        相机内参矩阵。
    R: np.ndarray
        相机旋转矩阵 (3x3)。
    T: np.ndarray
        相机平移向量 (3, 1)。
    dist_coeffs: np.ndarray
        相机的畸变系数。
    depth: float or np.ndarray
        点的深度（可以是标量或与点的数量匹配的数组）。
        
    Returns:
    points_3d_world: np.ndarray
        世界坐标系下的3D点 (N, 3)。
    """

    # 去畸变并将2D点转换为相机归一化坐标
    #points_undistorted = cv2.undistortPoints(points_2d, K, dist_coeffs)

    #print('points_undistorted = ', points_undistorted)

    # 将2D点扩展为3D点，相机归一化坐标 (x, y, 1)
    points_3d_camera = np.hstack((points_2d.squeeze(), np.ones((points_2d.shape[0], 1))))

    # 乘以深度，得到实际的相机坐标系中的 3D 点
    points_3d_camera *= depth

    # 将相机坐标系中的点转换为世界坐标系中的点
    points_3d_world = np.dot(R.T, (points_3d_camera.T - T)).T

    return points_3d_world


def plot_3d_points(points_3d, label, color, ax):
    """
    绘制3D点云。
    
    Parameters:
    points_3d: np.ndarray
        3D世界坐标系中的点 (N, 3)。
    label: str
        相机的标签名称。
    color: str
        点的颜色。
    ax: Axes3D
        3D图形对象。
    """
   
    ax.scatter(points_3d[:, 0], points_3d[:, 1], points_3d[:, 2], label=label, color=color)


# 定义投影函数
def project_points(points_3d, rvec, tvec, K, dist_coeffs):
    """
    将 3D 世界坐标点投影到 2D 图像平面。
    
    Parameters:
    points_3d: numpy.ndarray
        3D 点的坐标 (N, 3)。
    rvec: numpy.ndarray
        旋转向量 (3,)。
    tvec: numpy.ndarray
        平移向量 (3,)。
    K: numpy.ndarray
        相机内参矩阵 (3, 3)。
    dist_coeffs: numpy.ndarray
        畸变系数。
    
    Returns:
    image_points: numpy.ndarray
        投影到图像平面的 2D 点 (N, 2)。
    """
    image_points, _ = cv2.projectPoints(points_3d, rvec, tvec, K, dist_coeffs)
    return np.squeeze(image_points)

# def residuals(params, n_cameras, n_points, camera_indices, point_indices, points_2d, K_matrices, dist_coeffs):
#     """
#     计算捆绑调整中的残差，即重投影误差。
    
#     Parameters:
#     params: np.ndarray
#         优化的参数（包括所有相机的外参和所有3D点的坐标）。
#     n_cameras: int
#         相机的数量。
#     n_points: int
#         3D 点的数量。
#     camera_indices: np.ndarray
#         每个观测对应的相机索引。
#     point_indices: np.ndarray
#         每个观测对应的 3D 点索引。
#     points_2d: np.ndarray
#         实际观测到的 2D 图像点。
#     K_matrices: list of np.ndarray
#         每个相机的内参矩阵。
#     dist_coeffs: list of np.ndarray
#         每个相机的畸变系数。
    
#     Returns:
#     residuals: np.ndarray
#         重投影误差。
#     """
#     # 提取旋转和平移向量
#     rvecs = params[:n_cameras * 3].reshape((n_cameras, 3))  # 旋转向量
#     tvecs = params[n_cameras * 3:n_cameras * 6].reshape((n_cameras, 3))  # 平移向量
    
#     # 提取3D点
#     points_3d = params[n_cameras * 6:].reshape((n_points, 3))  # 3D 点
    
#     residuals = []

    
#     # 计算每个观测的重投影误差
#     for i in range(len(camera_indices)):
#         cam_idx = camera_indices[i]  # 当前观测的相机
#         point_idx = point_indices[i]  # 当前观测的 3D 点
#         K = K_matrices[cam_idx]  # 相机的内参矩阵
#         dist = dist_coeffs[cam_idx]  # 相机的畸变系数
        
#         # 投影3D点到相机平面，使用OpenCV的projectPoints函数
#         projected_points, _ = cv2.projectPoints(
#             points_3d[point_idx].reshape(1, -1),  # 当前的3D点
#             rvecs[cam_idx].reshape(-1, 1),        # 当前相机的旋转向量 (Rodrigues形式)
#             tvecs[cam_idx].reshape(-1, 1),        # 当前相机的平移向量
#             K,                                    # 内参矩阵
#             dist                                  # 畸变系数
#         )
        
#         # 将 projected_points 变成 (N, 2) 的形状
#         projected_points = projected_points.squeeze()

#         # 计算残差（重投影误差）
#         residuals.append(projected_points - points_2d[i])

#     return np.concatenate(residuals)



def calculate_normal(points):
    # points 是一个 (4, 3) 的数组，表示 4 个 3D 点
    vec1 = points[1] - points[0]
    vec2 = points[2] - points[0]
    normal = np.cross(vec1, vec2)  # 计算法向量
    return normal / np.linalg.norm(normal)  # 单位化


def point_to_plane_distance(point, plane_normal, plane_point):
    return np.dot(plane_normal, point - plane_point)



def residuals(params, n_cameras, n_points, camera_indices, point_indices, points_2d, K_matrices, dist_coeffs, plane_constraints):
    rvecs = params[:n_cameras * 3].reshape((n_cameras, 3))
    tvecs = params[n_cameras * 3:n_cameras * 6].reshape((n_cameras, 3))
    points_3d = params[n_cameras * 6:].reshape((n_points, 3))

    residuals = []

    # 计算重投影误差
    for i in range(len(camera_indices)):
        cam_idx = camera_indices[i]
        point_idx = point_indices[i]
        K = K_matrices[cam_idx]
        dist = dist_coeffs[cam_idx]

        projected_points, _ = cv2.projectPoints(
            points_3d[point_idx].reshape(1, -1),
            rvecs[cam_idx].reshape(3, 1),
            tvecs[cam_idx].reshape(3, 1),
            K, dist
        )
        projected_points = projected_points.squeeze()
        residual = projected_points - points_2d[i]
        residuals.extend(residual.tolist())  # 将重投影误差的两个分量添加到 residuals

    # 计算平面约束残差
    for constraint in plane_constraints:
        points_idx = constraint
        plane_points = points_3d[points_idx[:3]]  # 取前三个点计算平面
        plane_normal = calculate_normal(plane_points)
        plane_point = plane_points[0]

        for idx in points_idx:
            distance = point_to_plane_distance(points_3d[idx], plane_normal, plane_point)
            residuals.append(distance)  # 添加标量距离值

    return np.array(residuals)

def extract_params_from_result(result, n_cameras, n_points):
    """
    从优化结果中提取旋转向量、平移向量和3D点数据。
    
    Parameters:
    result: OptimizeResult
        scipy.optimize.least_squares 返回的结果。
    n_cameras: int
        相机的数量。
    n_points: int
        3D 点的数量。
        
    Returns:
    rvecs: list of np.ndarray
        优化后的旋转向量。
    tvecs: list of np.ndarray
        优化后的平移向量。
    points_3d: np.ndarray
        优化后的3D点坐标。
    """
    optimized_params = result.x
    
    # 提取旋转和平移向量
    rvecs = optimized_params[:n_cameras * 3].reshape((n_cameras, 3))  # 旋转向量
    tvecs = optimized_params[n_cameras * 3:n_cameras * 6].reshape((n_cameras, 3))  # 平移向量
    
    # 提取3D点
    #points_3d = optimized_params[n_cameras * 6:].reshape((n_points, 3))  # 3D点
    
    return rvecs, tvecs


def compare_params(optimized_rvecs, optimized_tvecs, original_cam_params):
    """
    对比优化后的相机外参和原始标定结果。
    
    Parameters:
    optimized_rvecs: np.ndarray
        优化后的旋转向量。
    optimized_tvecs: np.ndarray
        优化后的平移向量。
    original_cam_params: list of dict
        原始标定结果（包含旋转矩阵和平移向量）。
        
    Returns:
    None
    """
    for cam_idx, original_params in enumerate(original_cam_params):
        # 原始旋转矩阵和转换为旋转向量
        original_rvec = cv2.Rodrigues(original_params['rotation_matrix'])[0].flatten()
        original_tvec = original_params['translation_vector'].flatten()

        # 优化后的旋转向量和平移向量
        optimized_rvec = optimized_rvecs[cam_idx]
        optimized_tvec = optimized_tvecs[cam_idx]

        # 计算旋转向量和平移向量的差异
        rotation_diff = np.linalg.norm(original_rvec - optimized_rvec)
        translation_diff = np.linalg.norm(original_tvec - optimized_tvec)

        print(f"相机 {cam_idx} 差异：")
        print(f"旋转向量差异: {rotation_diff}")
        print(f"平移向量差异: {translation_diff}")



def calculate_reprojection_error(rvecs, tvecs, points_3d, points_2d, K_matrices, dist_coeffs):
    """
    计算重投影误差。
    
    Parameters:
    rvecs: np.ndarray
        相机的旋转向量。
    tvecs: np.ndarray
        相机的平移向量。
    points_3d: np.ndarray
        3D点坐标。
    points_2d: np.ndarray
        2D图像中的对应点。
    K_matrices: list of np.ndarray
        相机内参矩阵。
    dist_coeffs: list of np.ndarray
        畸变系数。
        
    Returns:
    float
        平均重投影误差。
    """
    total_error = 0
    total_points = 0

    for cam_idx in range(len(rvecs)):
        projected_points, _ = cv2.projectPoints(points_3d, rvecs[cam_idx], tvecs[cam_idx], K_matrices[cam_idx], dist_coeffs[cam_idx])
        error = np.linalg.norm(points_2d[cam_idx] - projected_points.squeeze(), axis=1).sum()
        total_error += error
        total_points += len(points_2d[cam_idx])

    return total_error / total_points



def bundle_adjustment(all_object_points, all_corners, cam_params, plane_constraints):
    """
    使用捆绑调整优化相机参数和3D点坐标。
    
    Parameters:
    all_object_points: list of list of np.ndarray
        每个相机中每个图像中的 3D 物体点。
    all_corners: list of list of np.ndarray
        每个相机中的多个图像中的 2D 图像点（每台相机拍摄了多张图片）。
    cam_params: list of dict
        每个相机的内参和外参。
    
    Returns:
    result: scipy.optimize.OptimizeResult
        优化后的参数结果。
    """
    n_cameras = len(cam_params)  # 相机的数量

    # 展平所有 3D 物体点
    all_object_points_flat = []
    for object_points_cam in all_object_points:
        all_object_points_flat.extend(object_points_cam)  # 将每个相机的所有 3D 物体点展平

    all_object_points_flat = np.concatenate(all_object_points_flat)
    n_points = all_object_points_flat.shape[0]

    # 收集相机内参矩阵和畸变系数
    K_matrices = [params['camera_matrix'] for params in cam_params]
    dist_coeffs = [params['distortion_coefficients'] for params in cam_params]

    # 创建相机索引和点索引，用于优化
    camera_indices = []
    point_indices = []
    points_2d = []

    # 记录全局 3D 点的偏移量，用于正确生成 point_indices
    point_offset = 0

    # 遍历每台相机的所有图像
    for cam_idx, (object_points_cam, corners_cam) in enumerate(zip(all_object_points, all_corners)):
        for img_idx, points_2d_list in enumerate(corners_cam):  # 遍历每个相机的每张图像中的2D点
            n_points_in_image = points_2d_list.shape[0]

            # 给每个2D点分配当前相机的相机索引（相机编号 cam_idx）
            camera_indices.extend([cam_idx] * n_points_in_image)
            
            # 给每个2D点分配全局3D点索引
            point_indices.extend(range(point_offset, point_offset + n_points_in_image))
            
            # 增加偏移量
            point_offset += n_points_in_image

            # 收集2D图像点
            points_2d.extend(points_2d_list)

    camera_indices = np.array(camera_indices)
    point_indices = np.array(point_indices)
    points_2d = np.array(points_2d)


    # 初始化外参（旋转和平移向量）和 3D 点
    rvecs = [cv2.Rodrigues(params['rotation_matrix'])[0].flatten() for params in cam_params]
    tvecs = [params['translation_vector'].flatten() for params in cam_params]

    # 将外参和3D点展平为一维数组
    initial_params = np.hstack([
        np.concatenate(rvecs),  # 所有相机的旋转向量
        np.concatenate(tvecs),  # 所有相机的平移向量
        all_object_points_flat.ravel()  # 所有3D点展平
    ])

    # 执行捆绑调整，使用 'trf' 或 'dogbox' 方法
    # result = least_squares(
    #     residuals, initial_params, verbose=2, x_scale='jac', ftol=1e-4, method='trf',  # 使用 'trf' 方法
    #     args=(n_cameras, n_points, camera_indices, point_indices, points_2d, K_matrices, dist_coeffs)
    # )
    result = least_squares(
    residuals, initial_params, verbose=2, x_scale='jac', ftol=1e-4, method='trf',
    args=(n_cameras, n_points, camera_indices, point_indices, points_2d, K_matrices, dist_coeffs, plane_constraints)
)

    # 优化前的重投影误差
    original_rvecs = [cv2.Rodrigues(params['rotation_matrix'])[0] for params in cam_params]
    original_tvecs = [params['translation_vector'] for params in cam_params]

    # original_error = calculate_reprojection_error(
    #     original_rvecs, original_tvecs, original_points_3d, original_points_2d, K_matrices, dist_coeffs
    # )

    # # 优化后的重投影误差
    # optimized_error = calculate_reprojection_error(
    #     optimized_rvecs, optimized_tvecs, points_3d, points_2d, K_matrices, dist_coeffs
    # )

    #print(f"优化前的重投影误差: {original_error}")
    #print(f"优化后的重投影误差: {optimized_error}")

    return result


def visualize_projection(image, points_2d, color='r', title='Projected Points'):
    """
    在图像上绘制投影点。
    
    Parameters:
    image: numpy.ndarray
        图像数据。
    points_2d: numpy.ndarray
        要绘制的 2D 点。
    color: str
        绘制点的颜色。
    title: str
        图像标题。
    """
    plt.figure(figsize=(8, 6))
    plt.imshow(image, cmap='gray')
    plt.scatter(points_2d[:, 0], points_2d[:, 1], c=color, s=50, label='Projected Points')
    plt.legend()
    plt.title(title)
    plt.show()


def compute_reprojection_error(projected_points, observed_points):
    """
    计算投影点与实际观测点之间的误差。
    
    Parameters:
    projected_points: numpy.ndarray
        投影得到的2D点，形状为 (N, 2)。
    observed_points: numpy.ndarray
        实际观测到的2D点，形状为 (N, 2)。
    
    Returns:
    float
        平均重投影误差。
    """
    error = np.linalg.norm(projected_points - observed_points, axis=1)
    return np.mean(error)

# 将 3D 点转换到相机坐标系
def transform_points_to_camera_frame(points, R, T):
    return (R @ points.T + T).T


def detect_april_tag_points(detector, img_file, object_point_single):
    """
    检测单张图像中的 AprilTags,并返回图像中的 2D 角点和相应的 3D 物体点。
    
    Parameters:
    detector: AprilTag 检测器对象
    img_file: 图像文件路径
    object_point_single: 标定板的 3D 物体点
    
    Returns:
    all_corners: np.ndarray
        图像中的 2D 角点 (N, 2)
    all_object_points: np.ndarray
        对应的 3D 物体点 (N, 3)
    """
    #print('Processing image file:', img_file)

    # 读取并转换为灰度图像
    src = cv2.imread(img_file)
    gray = cv2.imread(img_file, 0)
    if gray is None:
        print(f"Failed to load image: {img_file}")
        return np.array([]), np.array([])

    # 检测 AprilTags
    detections = detector.detect(gray)
    all_corners = []
    all_object_points = []
    tag_id_lst = []

    #show_detect_result(src, detections)

    if detections:
        for detection in detections:
            corners = detection.corners  # 检测到的角点 (4, 2)
            tag_id = detection.tag_id    # 标签 ID
            tag_id_lst.append(int(tag_id))


            # 将检测到的角点存储为 np.float32 格式
            all_corners.extend(corners.astype(np.float32))

            # 根据标签 ID 提取相应的 3D 坐标
            if tag_id < len(object_point_single):
                all_object_points.extend(object_point_single[tag_id])
            else:
                print(f"Warning: Tag ID {tag_id} is out of range.")

    if all_object_points and all_corners:
        return np.array(all_corners, dtype=np.float32).reshape(-1, 2), \
               np.array(all_object_points, dtype=np.float32).reshape(-1, 3), \
               tag_id_lst, detections
    else:
        print("No valid detections in this image.")
        return np.array([]), np.array([]), [], []



def single_camera_calibrate_vis():
    config_path = 'D:/code/pmdrecons-python-xtt/config/cfg_3freq_wafer.json'
    para_path = 'D:\\code\\pmdrecons-python-xtt\\config\\world_params.mat'
    world_params = loadmat(para_path)
    camera_matrix = np.array(world_params['camera_matrix'])
    dist_coeffs = np.array(world_params['distortion_coefficients'])
    img_dir = 'D:\\data\\one_cam\\pmd_benchmarks\\set4-20240906\\set4-20240906\\calibration\\'

    img_paths = glob.glob(os.path.join(img_dir, 'world_calibration/*.bmp'))

    debug = 1

    cfg = json.load(open(config_path, 'r'))
    chessboard_size = cfg['world_chessboard_size']
    square_size = cfg['world_square_size']

    # 准备棋盘格的世界坐标
    objp = np.zeros((chessboard_size[0]*chessboard_size[1], 3), np.float32)
    objp[:, :2] = np.mgrid[0:chessboard_size[0], 0:chessboard_size[1]].T.reshape(-1, 2)
    objp *= square_size

    # 存储每个棋盘格在相机坐标系中的位置
    chessboard_positions = []

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

        # 寻找棋盘格角点
        ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

        if ret:
            # 提高角点精度
            criteria = (cv2.TermCriteria_EPS + cv2.TermCriteria_MAX_ITER, 30, 0.001)
            corners_subpix = cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria)

            # 计算棋盘格的姿态（相对于相机）
            success, rvec, tvec = cv2.solvePnP(objp, corners_subpix, camera_matrix, dist_coeffs)
            if success:
                # 计算旋转矩阵
                R, _ = cv2.Rodrigues(rvec)

                # 将棋盘格角点转换到相机坐标系
                # objp_in_camera = R * objp^T + tvec
                objp_in_camera = (R @ objp.T + tvec).T  # 转置以获得 N x 3 的矩阵

                # 存储转换后的棋盘格角点
                chessboard_positions.append(objp_in_camera)
            else:
                print("姿态估计失败：", fname)
        else:
            print("无法找到棋盘格角点：", fname)

    cv2.destroyAllWindows()

    # 绘制多个棋盘格在相机坐标系中的位置
    if len(chessboard_positions) > 0:
        # 创建3D绘图
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        # 绘制每个棋盘格
        for idx, chessboard in enumerate(chessboard_positions):
            x = chessboard[:, 0]
            y = chessboard[:, 1]
            z = chessboard[:, 2]
            ax.scatter(x, y, z, label=f'Chessboard {idx+1}')
            plt.pause(3)
            

        # 绘制相机位置（原点）
        ax.scatter(0, 0, 0, c='r', marker='^', s=100, label='Camera Position')

        # 设置标签和标题
        ax.set_xlabel('X axis')
        ax.set_ylabel('Y axis')
        ax.set_zlabel('Z axis')
        ax.set_title('Multiple Chessboard Positions Relative to Fixed Camera')

        # 设置等比例显示
        all_x = np.hstack([chessboard[:, 0] for chessboard in chessboard_positions])
        all_y = np.hstack([chessboard[:, 1] for chessboard in chessboard_positions])
        all_z = np.hstack([chessboard[:, 2] for chessboard in chessboard_positions])
        ax.set_box_aspect([np.ptp(a) for a in [all_x, all_y, all_z]])

        # 显示图例
        ax.legend()
       

        plt.show()
    else:
        print("未能计算出有效的棋盘格位置。")

    return


def draw(img, corner, imgpts):
    corner = tuple(corner.ravel().astype(int))
    imgpts = imgpts.astype(int)

    # 绘制坐标轴
    img = cv2.line(img, corner, tuple(imgpts[1].ravel()), (0, 0, 255), 5)  # X轴（红色）
    img = cv2.line(img, corner, tuple(imgpts[2].ravel()), (0, 255, 0), 5)  # Y轴（绿色）
    img = cv2.line(img, corner, tuple(imgpts[3].ravel()), (255, 0, 0), 5)  # Z轴（蓝色）

    return img



def verify_coplanarity(points_3d):
    # 使用SVD拟合平面
    centroid = np.mean(points_3d, axis=0)
    centered_points = points_3d - centroid
    U, S, Vt = np.linalg.svd(centered_points)
    normal_vector = Vt[-1, :]
    # 计算每个点到平面的距离
    distances = np.abs(centered_points @ normal_vector)
    mean_distance = np.mean(distances)
    max_distance = np.max(distances)
    print(f"Mean distance to plane: {mean_distance}")
    print(f"Max distance to plane: {max_distance}")
    # 设置一个阈值判断是否共面
    threshold = 1e-3  # 根据需求调整
    if max_distance < threshold:
        print("All points are approximately coplanar.")
    else:
        print("Points are not coplanar.")



def fit_plane(points):
    """拟合一个平面并返回法向量和质心."""
    centroid = np.mean(points, axis=0)  # 计算质心
    centered_points = points - centroid  # 将点移到质心

    # 使用 SVD 分解拟合平面法向量
    _, _, vh = np.linalg.svd(centered_points)
    normal_vector = vh[-1]  # 法向量是 SVD 的最后一列

    return normal_vector, centroid


def rotation_matrix_from_vectors(v1, v2):
    """计算将向量 v1 旋转到 v2 的旋转矩阵."""
    v1 = v1 / np.linalg.norm(v1)
    v2 = v2 / np.linalg.norm(v2)
    axis = np.cross(v1, v2)  # 旋转轴
    angle = np.arccos(np.dot(v1, v2))  # 旋转角度

    # 构建旋转矩阵（Rodrigues 公式）
    K = np.array([
        [0, -axis[2], axis[1]],
        [axis[2], 0, -axis[0]],
        [-axis[1], axis[0], 0]
    ])
    rotation_matrix = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)
    return rotation_matrix


def align_points_to_horizontal(points):
    """将给定的 3D 点集旋转到水平面."""
    normal_vector, centroid = fit_plane(points)  # 拟合平面
    print('normal_vector = ', normal_vector)
    rotation_to_horizontal = rotation_matrix_from_vectors(normal_vector, np.array([0, 0, 1]))
    #print('rotation_to_horizontal = ', rotation_to_horizontal)

    # 将所有点旋转到水平面，并返回结果
    centered_points = points - centroid  # 平移到质心
    rotated_points = (rotation_to_horizontal @ centered_points.T).T + centroid  # 旋转并平移回去

    return rotated_points, rotation_to_horizontal



def camera_calibrate_vis():
    # 加载配置和相机参数
    config_path = 'D:\\code\\pmd-python\\config\\cfg_3freq_wafer.json'
    detector = Detector('t36h11')
    cfg = json.load(open(config_path, 'r'))
    with open(os.path.join('D:\\code\\pmd-python\\config\\', 'cam_params1023.pkl'), 'rb') as pkl_file:
        cam_params = pickle.load(pkl_file)
    
    # 生成 3D 物体点
    world_chessboard_size, world_square_size, world_square_spacing = [10, 6], 35.2, 10.56
    object_point_single = generate_3d_points(world_chessboard_size, world_square_size, world_square_spacing)  
    
    # 保存所有相机的观测数据
    all_points_3d_world = []
    all_image_points = []
    all_camera_centers = []
    K_matrices = []
    dist_coeffs = []
    all_corners = []
    all_object_points = []
    optimized_cam_params = []
    all_tag_lst = []
    all_detections = []

    # 定义四个相机的标定图像路径
    calibration_images_path = [
        'D:\\data\\four_cam\\calibrate\\single_test\\cam0\\',
        'D:\\data\\four_cam\\calibrate\\single_test\\cam1\\',
        'D:\\data\\four_cam\\calibrate\\single_test\\cam2\\',
        'D:\\data\\four_cam\\calibrate\\single_test\\cam3\\'
    ]


    # 遍历每个相机
    for cam_id, params in enumerate(cam_params):
        K_matrices.append(params['camera_matrix'])
        dist_coeffs.append(params['distortion_coefficients'])
        
        camera_corners = []  # 用于当前相机的2D图像点
        camera_object_points = []  # 用于当前相机的3D物体点
       
        # 对每个相机的标定图像提取2D角点
        for img_file in os.listdir(calibration_images_path[cam_id]):
            # 读取图像并检测AprilTags，获取2D角点和对应的3D物体点
            image_points, object_points, tag_id_lst, detections = detect_april_tag_points(
                detector, calibration_images_path[cam_id] + img_file, object_point_single
            )
            # print('image_points  = ', image_points)
            # print('object_points =', object_points)
            print('tag_id_lst = ', tag_id_lst)
            all_tag_lst.append(tag_id_lst)
            all_corners.append(image_points)
            all_object_points.append(object_points)
            all_detections.append(detections)
            undistorted_points = cv2.undistortPoints(image_points,  params['camera_matrix'], params['distortion_coefficients'])
            success, rotation_vector, translation_vector = cv2.solvePnP(object_points, image_points, params['camera_matrix'], 
                                                                       params['distortion_coefficients'],
                                                                       params['rotation_matrix'],
                                                                       params['translation_vector']
                                                                       )
            rotation_matrix, _ = cv2.Rodrigues(rotation_vector)
            # 假设的深度值 (s)
            # 你需要根据场景给出一个合理的深度值
            # 相机的中心
            
            camera_center = -rotation_matrix.T @ translation_vector
            #print('camera_center = ', camera_center)
            depth = camera_center[-1]
            print('undistorted_points shape = ', undistorted_points.shape)
            print('image_points shape = ', image_points.shape)
            print('depth = ', depth)

            # 将 2D 图像点转换为齐次坐标
            #print(image_points.shape)
            image_points_homogeneous = np.hstack([image_points, np.ones((image_points.shape[0], 1), dtype=np.float32)])
            # 批量计算相机坐标系中的 3D 点 (N x 3 矩阵)
            camera_coordinates = (np.linalg.inv(params['camera_matrix']) @ image_points_homogeneous.T).T * depth

            # 将相机坐标转换为世界坐标 (N x 3 矩阵)
            world_coordinates = (np.linalg.inv(rotation_matrix) @ (camera_coordinates - translation_vector.T).T).T
            all_points_3d_world.append(world_coordinates)

    print('len 3d = ', len(all_points_3d_world))

    # adjusted_cam_params = []

    # for i, params in enumerate(cam_params):
    #     world_coordinates = all_points_3d_world[i]  # 获取该相机的 3D 点
    #     #print('world_coordinates = ', world_coordinates)

    #     # 将点集旋转到水平面，并获取旋转矩阵
    #     #aligned_points, rotation_to_horizontal = align_points_to_horizontal(world_coordinates)
    #     aligned_points, rotation_to_horizontal = align2ref(world_coordinates)
        
    #     print('aligned_points = ', rotation_matrix)
    #     print('rotation_to_horizontal = ', rotation_to_horizontal)

    #     # 更新相机的外参（旋转矩阵）
    #     params['rotation_matrix'] = rotation_to_horizontal @ params['rotation_matrix']
    #     adjusted_cam_params.append(params)

    #     # 更新校正后的 3D 点
    #     all_points_3d_world[i] = aligned_points
    
    # fig = plt.figure()
    # ax = fig.add_subplot(111, projection='3d')

    # for i, points in enumerate(all_points_3d_world):
    #     ax.scatter(points[:, 0], points[:, 1], points[:, 2], label=f'Cam {i}')

    # ax.set_xlabel('X')
    # ax.set_ylabel('Y')
    # ax.set_zlabel('Z')
    # plt.legend()
    # plt.show()

    # # 3. 保存调整后的相机参数，以备后续使用
    # with open(os.path.join('D:\\code\\pmd-python\\config\\', 'cam_params1023.pkl'), 'wb') as f:
    #     pickle.dump(adjusted_cam_params, f)
    

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    colors = ['r', 'g', 'b', 'c']  # 颜色列表

    for i in range(len(all_points_3d_world)):
        
        plot_3d_points(all_points_3d_world[i], label=f'Cam {i}', color=[random.random() for _ in range(3)], ax=ax)
        # 显示图例和更新当前图像
        ax.legend(loc='upper right')
        plt.pause(2)  # 暂停 1 秒，逐个显示效果

    
    # for cam_idx, (params, corners) in enumerate(zip(cam_params, all_corners)):
    #     #print('all_points_3d_world[cam_idx] = ', all_points_3d_world[cam_idx])
    #     plot_3d_points(all_points_3d_world[cam_idx], label=f'Cam {cam_idx}', color=colors[cam_idx], ax=ax)

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
   # ax.set_zlim([0, 500])
    plt.legend()
    plt.show()

    #cam_params[3]['rotation_matrix'] = np.eye(3)
    #cam_params[3]['translation_vector'] = np.zeros((3, 1))

    # stereo_calibration_data_0_3 = stereo_calibrate(all_detections, 0, 3, object_point_single, cam_params)
    # stereo_calibration_data_1_3 = stereo_calibrate(all_detections, 1, 3, object_point_single, cam_params)
    # stereo_calibration_data_2_3 = stereo_calibrate(all_detections, 2, 3, object_point_single, cam_params)

    # stereo_calib_data_list = []
    # stereo_calib_data_list.append(stereo_calibration_data_0_3)
    # stereo_calib_data_list.append(stereo_calibration_data_1_3)
    # stereo_calib_data_list.append(stereo_calibration_data_2_3)


    #     # 设置相机 3 的外参不变
    # R_3 = cam_params[3]['rotation_matrix']
    # T_3 = cam_params[3]['translation_vector']

    # for i in [0, 1, 2]:
    #     # 获取相机 i 的初始外参
    #     R_i = cam_params[i]['rotation_matrix']
    #     T_i = cam_params[i]['translation_vector']
        
    #     # 计算初始相对位姿
    #     R_i3_init = R_3 @ R_i.T
    #     T_i3_init = -R_3 @ R_i.T @ T_i + T_3
        
    #     # 获取双目标定结果
    #     stereo_calib_data = stereo_calib_data_list[i]
    #     R_i3_stereo = stereo_calib_data['rotation_matrix']
    #     T_i3_stereo = stereo_calib_data['translation_vector']
        
    #     # 计算位姿差异
    #     delta_R = R_i3_stereo @ R_i3_init.T
    #     delta_T = T_i3_stereo - delta_R @ T_i3_init
        
    #     # 更新相机外参
    #     R_i_new = delta_R @ R_i
    #     T_i_new = delta_R @ T_i + delta_T
        
    #     # 更新 cam_params
    #     cam_params[i]['rotation_matrix'] = R_i_new
    #     cam_params[i]['translation_vector'] = T_i_new
    

    # with open('D:\\code\\pmd-python\\config\\stereo_calib2.pkl', 'wb') as pkl_file:
    #        pickle.dump(cam_params, pkl_file) 

    
def stereo_calibrate(all_detections, cam_id1, cam_id2, object_point_single, cam_params):

    detections1 = all_detections[cam_id1]
    detections2 = all_detections[cam_id2]
   
    detections_dict1 = {d.tag_id: d for d in detections1}
    detections_dict2 = {d.tag_id: d for d in detections2}

    # 用于保存所有图像的检测到的角点和 ID
    objpoints = []  # 3D 点
    imgpoints1 = []  # 相机1的 2D 点
    imgpoints2 = []  # 相机2的 2D 点
    image_size = None

    # 找到两个相机共同检测到的标签 ID
    common_tags = set(detections_dict1.keys()) & set(detections_dict2.keys())
    print('common_tags = ', common_tags)
    if common_tags:
        image_points1 = []
        image_points2 = []
        object_points = []

        for tag_id in common_tags:
            # 获取相机1的角点
            corners1 = detections_dict1[tag_id].corners.astype(np.float32)
            # 获取相机2的角点
            corners2 = detections_dict2[tag_id].corners.astype(np.float32)

            # 获取对应的 3D 点
            if tag_id < len(object_point_single):
                obj_pts = object_point_single[tag_id]
            else:
                print(f"Warning: Tag ID {tag_id} is out of range.")
                continue

            object_points.append(obj_pts)
            image_points1.append(corners1)
            image_points2.append(corners2)

        if object_points and image_points1 and image_points2:
            objpoints.append(np.array(object_points, dtype=np.float32).reshape(-1, 3))
            imgpoints1.append(np.array(image_points1, dtype=np.float32).reshape(-1, 2))
            imgpoints2.append(np.array(image_points2, dtype=np.float32).reshape(-1, 2))
        else:
            print("No valid common detections in this image pair.")
    else:
        print("No common tags detected in this image pair.")
    
    K1 = cam_params[cam_id1]['camera_matrix']
    dist1 = cam_params[cam_id1]['distortion_coefficients']
    K2 = cam_params[cam_id2]['camera_matrix']
    dist2 = cam_params[cam_id2]['distortion_coefficients']

    # 立体标定
    flags = 0
    flags |= cv2.CALIB_FIX_INTRINSIC  # 固定单目内参

    criteria = (cv2.TERM_CRITERIA_MAX_ITER + cv2.TERM_CRITERIA_EPS, 100, 1e-5)

    ret, _, _, _, _, R, T, E, F = cv2.stereoCalibrate(
        objpoints,       # 3D 点
        imgpoints1,      # 相机1的 2D 点
        imgpoints2,      # 相机2的 2D 点
        K1,              # 相机1的内参
        dist1,           # 相机1的畸变系数
        K2,              # 相机2的内参
        dist2,           # 相机2的畸变系数
        image_size,      # 图像尺寸
        criteria=criteria,
        flags=flags
    )

    print(f"Stereo Calibration between Camera {cam_id1} and Camera {cam_id2} completed.")
    print(f"Rotation Matrix R:\n{R}")
    print(f"Translation Vector T:\n{T}")

    # 将结果保存到字典中
    stereo_calibration_data = {
        'camera_matrix1': K1,
        'distortion_coefficients1': dist1,
        'camera_matrix2': K2,
        'distortion_coefficients2': dist2,
        'rotation_matrix': R,
        'translation_vector': T,
        'essential_matrix': E,
        'fundamental_matrix': F,
        'reprojection_error': ret
    }

    # 返回标定结果
    print('stereo_calibration_data = ', stereo_calibration_data)

    return stereo_calibration_data


# ######################## bundle_adjustment
#     # 遍历每个相机
#     for cam_id, params in enumerate(cam_params):
#         K_matrices.append(params['camera_matrix'])
#         dist_coeffs.append(params['distortion_coefficients'])
        
#         camera_corners = []  # 用于当前相机的2D图像点
#         camera_object_points = []  # 用于当前相机的3D物体点

#         # 对每个相机的标定图像提取2D角点
#         for img_file in os.listdir(calibration_images_path[cam_id]):
#             # 读取图像并检测AprilTags，获取2D角点和对应的3D物体点
#             image_points, object_points, tag_lst = detect_april_tag_points(
#                 detector, calibration_images_path[cam_id] + img_file, object_point_single
#             )

#             # 将检测到的2D图像点和3D物体点保存到当前相机的列表
#             if image_points is not None and object_points is not None:
#                 camera_corners.append(image_points)
#                 camera_object_points.append(object_points)
        
#         # 保存当前相机的所有图像点到全局列表
#         all_corners.append(camera_corners)
#         all_object_points.append(camera_object_points)

#     # 调用捆绑调整函数，优化所有相机的内参、外参和3D点
#     plane_constraints = [
#     [0, 1, 2, 3],  # 第一个平面的点索引
#     [4, 5, 6, 7]   # 第二个平面的点索引
#     ]   
#     result = bundle_adjustment(all_object_points, all_corners, cam_params, plane_constraints)
    
#     print(result)
#     optimized_rvecs, optimized_tvecs = extract_params_from_result(result, 4, len(all_object_points))
#     original_rvecs = [cv2.Rodrigues(params['rotation_matrix'])[0] for params in cam_params]
#     original_tvecs = [params['translation_vector'] for params in cam_params]

#     print('optimized_rvecs = ',optimized_rvecs)
    
#     for i in range(4):
#         optimized_rotation_matrices = cv2.Rodrigues(optimized_rvecs[i])[0]
#         #print(original_tvecs[i], optimized_tvecs[i].reshape(-1, 1))
#         calibration_data = {
#         'camera_matrix': cam_params[i]['camera_matrix'],
#         'distortion_coefficients': cam_params[i]['distortion_coefficients'],
#         'rotation_matrix': optimized_rotation_matrices,
#         'translation_vector': optimized_tvecs[i].reshape(-1, 1)
#         }
        
#         optimized_cam_params.append(calibration_data)

#     with open('D:\\code\\pmd-python\\config\\optimized_param1021.pkl', 'wb') as pkl_file:
        
#         pickle.dump(optimized_cam_params, pkl_file)
    
#     # 打印优化后的结果
#     print("Bundle Adjustment 完成!")


    # #绘制相机
    # fig = plt.figure()
    # ax = fig.add_subplot(111, projection='3d')

    # # 绘制每个相机
    # for ii in range(4):
    #     plot_camera(ax, cam_params[ii]['rotation_matrix'], cam_params[ii]['translation_vector'], ii, scale=40, color='r')

    # # 绘制物体 (假设是一个简单的平面物体，位于 z = 0 平面上)
    # plane_x, plane_y = np.meshgrid(np.linspace(-500, 500, 10), np.linspace(-500, 500, 10))
    # plane_z = np.zeros_like(plane_x)  # z = 0 的平面
    # ax.plot_surface(plane_x, plane_y, plane_z, color='blue', alpha=0.3)

    # # 设置轴的范围和标签
    # ax.set_xlabel('X axis')
    # ax.set_ylabel('Y axis')
    # ax.set_zlabel('Z axis')
    # ax.set_xlim([-50, 300])
    # ax.set_ylim([-50, 300])
    # ax.set_zlim([0, 500])

    # plt.show()

def plt_gaussion():

    # 1. Create a Gaussian surface
    X = np.linspace(-3, 3, 100)
    Y = np.linspace(-3, 3, 100)
    X, Y = np.meshgrid(X, Y)
    Z = np.exp(-(X**2 + Y**2))

    # 2. Compute gradients
    Zx, Zy = np.gradient(Z)

    # 3. Plotting the Gaussian surface
    fig = plt.figure(figsize=(14, 6))

    # Plot 1: Gaussian Surface
    ax1 = fig.add_subplot(121, projection='3d')
    ax1.plot_surface(X, Y, Z, cmap='viridis', edgecolor='none')
    ax1.set_title('Gaussian Surface')
    ax1.set_xlabel('X')
    ax1.set_ylabel('Y')
    ax1.set_zlabel('Z')

    # Plot 2: Gradient Field Visualization
    ax2 = fig.add_subplot(122)
    ax2.quiver(X, Y, Zx, Zy, color='r')
    ax2.set_title('Gradient Field of Gaussian Surface')
    ax2.set_xlabel('X')
    ax2.set_ylabel('Y')

    plt.tight_layout()
    plt.show()



if __name__ == '__main__':
    # white_path = 'D:\\data\\four_cam\\1008_storage\\'
    # cam_num = 1
    # #find_notch(white_path, cam_num)
    # for img in os.listdir(white_path):
    #     if img.endswith('.bmp'):
    #         get_white_mask(white_path + img, bin_thresh=12)
    camera_calibrate_vis()
    #plt_gaussion()
    #single_camera_calibrate_vis()