gha / data_tool /pose_estimation.py

Add files using upload-large-folder tool

3217786 verified 3 months ago

5.39 kB

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()

	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()