import cv2
import dlib
import numpy as np
from scipy.optimize import least_squares

# 加载dlib的预训练人脸检测器和关键点检测器模型
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('/data_160TB/2021/yangkun/gaussian-splatting-main-copy/flame/shape_predictor_68_face_landmarks.dat')

# 定义3D模型点（根据68个关键点）
model_points = np.array([
    (0.0, 0.0, 0.0),             # 鼻尖
    (0.0, -330.0, -65.0),        # 下巴
    (-225.0, 170.0, -135.0),     # 左眼角
    (225.0, 170.0, -135.0),      # 右眼角
    (-150.0, -150.0, -125.0),    # 左嘴角
    (150.0, -150.0, -125.0)      # 右嘴角
], dtype=np.float32)//1000

# 定义相机内参矩阵
# 这里 fx, fy 代表焦距，cx, cy 代表光心坐标
fx, fy, cx, cy = 2183.5264, 2183.5264, 256, 256  # 示例值，请根据实际相机参数调整
camera_matrix = np.array([
    [fx, 0, cx],
    [0, fy, cy],
    [0, 0, 1]
], dtype=np.float32)

# 假设无畸变
dist_coeffs = np.zeros((4, 1))

# 准备图像列表
image_files = ['/data_160TB/2021/yangkun/gaussian-splatting-main-copy/eg3d_human/test/person_0/masked/0000.png', 
            '/data_160TB/2021/yangkun/gaussian-splatting-main-copy/eg3d_human/test/person_0/masked/0030.png', 
            '/data_160TB/2021/yangkun/gaussian-splatting-main-copy/eg3d_human/test/person_0/masked/0060.png',
            '/data_160TB/2021/yangkun/gaussian-splatting-main-copy/eg3d_human/test/person_0/masked/0090.png']  # 示例文件名列表
image_points_list = []

# 提取每张图像中的关键点
for file in image_files:
    img = cv2.imread(file)
    if img is None:
        print(f"Error: Could not read image {file}")
        continue
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    faces = detector(gray)
    if len(faces) == 0:
        print(f"Warning: No face detected in image {file}")
        continue
    for face in faces:
        shape = predictor(gray, face)
        image_points = np.array([
            (shape.part(30).x, shape.part(30).y),  # 鼻尖
            (shape.part(8).x, shape.part(8).y),    # 下巴
            (shape.part(36).x, shape.part(36).y),  # 左眼角
            (shape.part(45).x, shape.part(45).y),  # 右眼角
            (shape.part(48).x, shape.part(48).y),  # 左嘴角
            (shape.part(54).x, shape.part(54).y)   # 右嘴角
        ], dtype=np.float32)
        image_points_list.append(image_points)
        break

if len(image_points_list) == 0:
    print("Error: No faces detected in any image")
    exit()

# 假设第一个图像是原点，计算相对于第一张图像的相机姿态
initial_poses = []
for i, image_points in enumerate(image_points_list):
    success, rotation_vector, translation_vector = cv2.solvePnP(
        model_points, image_points, camera_matrix, dist_coeffs)
    if not success:
        print(f"Error: solvePnP failed for image {i+1}")
        continue
    initial_poses.append((rotation_vector, translation_vector))

# 转换为优化问题的参数格式
def pack_params(poses):
    params = []
    for rvec, tvec in poses:
        params.append(rvec.ravel())
        params.append(tvec.ravel())
    return np.concatenate(params)

def unpack_params(params):
    poses = []
    for i in range(0, len(params), 6):
        rvec = params[i:i+3].reshape(3, 1)
        tvec = params[i+3:i+6].reshape(3, 1)
        poses.append((rvec, tvec))
    return poses

# 损失函数
def reprojection_error(params, model_points, image_points_list, camera_matrix, dist_coeffs):
    poses = unpack_params(params)
    errors = []
    for (rvec, tvec), img_points in zip(poses, image_points_list):
        projected_points, _ = cv2.projectPoints(model_points, rvec, tvec, camera_matrix, dist_coeffs)
        error = img_points - projected_points.reshape(-1, 2)
        errors.append(error.ravel())
    return np.concatenate(errors)

# 优化相机姿态
initial_params = pack_params(initial_poses)
result = least_squares(reprojection_error, initial_params, args=(model_points, image_points_list, camera_matrix, dist_coeffs))

# 获取优化后的相机姿态
optimized_poses = unpack_params(result.x)

poses = []
# 输出结果
for i, (rvec, tvec) in enumerate(optimized_poses):
    pose = np.ones((4, 4))
    rotation_matrix, _ = cv2.Rodrigues(rvec)
    print(f"Image {i+1} - Rotation Vector:\n{rotation_matrix}\nTranslation Vector:\n{tvec}")
    pose[0:3, 0:3] = rotation_matrix
    pose[0:3, 3] = tvec[0]
    poses.append(pose)
poses = np.stack(poses, 0)
np.save('/data_160TB/2021/yangkun/gaussian-splatting-main-copy/eg3d_human/test/camera_4.npy', poses)
# 可视化（以第一张图像为基准）
# img = cv2.imread(image_files[0])
# if img is None:
#     print(f"Error: Could not read image {image_files[0]}")
#     exit()
# for i, (rvec, tvec) in enumerate(optimized_poses):
#     axis = np.float32([[200, 0, 0], [0, 200, 0], [0, 0, 200]])
#     imgpts, _ = cv2.projectPoints(axis, rvec, tvec, camera_matrix, dist_coeffs)
#     imgpts = np.int32(imgpts).reshape(-1, 2)
#     img = cv2.line(img, tuple(image_points_list[0][0]), tuple(imgpts[0]), (255, 0, 0), 5)
#     img = cv2.line(img, tuple(image_points_list[0][0]), tuple(imgpts[1]), (0, 255, 0), 5)
#     img = cv2.line(img, tuple(image_points_list[0][0]), tuple(imgpts[2]), (0, 0, 255), 5)

# cv2.imshow('Optimized Result', img)
# cv2.waitKey(0)
# cv2.destroyAllWindows()