174 lines
7.1 KiB
Python
174 lines
7.1 KiB
Python
import cv2
|
||
import numpy as np
|
||
import itertools
|
||
from .reconstruction import get_camera_parameters
|
||
from .pose_estimation import get_ground_truth_seams, estimate_camera_pose
|
||
|
||
|
||
def get_global_transform_from_up_data(all_2d_endpoints):
|
||
"""
|
||
只使用 'up' 部分的数据，计算一个全局的、唯一的“相机->物体”变换矩阵。
|
||
"""
|
||
print("\n--- Calculating Global Transform Matrix using 'up' data ---")
|
||
ground_truth = get_ground_truth_seams()
|
||
|
||
# 1. 准备 'up' 部分的数据
|
||
object_points_3d = []
|
||
image_points_2d_L = []
|
||
|
||
for line_name in ['line1', 'line2']:
|
||
gt_key = f"up_{line_name}"
|
||
key_L = f"up_l_{line_name}"
|
||
object_points_3d.extend([ground_truth[gt_key]['start_3d'], ground_truth[gt_key]['end_3d']])
|
||
image_points_2d_L.extend([all_2d_endpoints[key_L]['start'], all_2d_endpoints[key_L]['end']])
|
||
|
||
# 2. 寻找最佳点对应关系
|
||
best_reprojection_error = float('inf')
|
||
best_pose = None
|
||
|
||
for a, b in itertools.product([0, 1], repeat=2):
|
||
current_image_points_L = list(image_points_2d_L)
|
||
if a: current_image_points_L[0], current_image_points_L[1] = current_image_points_L[1], current_image_points_L[
|
||
0]
|
||
if b: current_image_points_L[2], current_image_points_L[3] = current_image_points_L[3], current_image_points_L[
|
||
2]
|
||
|
||
rvec, tvec = estimate_camera_pose(current_image_points_L, object_points_3d, 'L')
|
||
if rvec is not None:
|
||
projected_points, _ = cv2.projectPoints(np.array(object_points_3d), rvec, tvec,
|
||
get_camera_parameters()[0]['K'], get_camera_parameters()[0]['kc'])
|
||
error = cv2.norm(np.array(current_image_points_L, dtype=np.float32),
|
||
projected_points.reshape(-1, 2).astype(np.float32), cv2.NORM_L2)
|
||
|
||
if error < best_reprojection_error:
|
||
best_reprojection_error = error
|
||
best_pose = (rvec, tvec)
|
||
|
||
if best_pose is None:
|
||
print("Fatal Error: Could not calculate a valid global transform.")
|
||
return None
|
||
|
||
print(f"Global transform calculated with reprojection error: {best_reprojection_error:.2f}")
|
||
|
||
# 3. 构建 4x4 变换矩阵 (从相机坐标系到物体坐标系)
|
||
rvec, tvec = best_pose
|
||
R_cam_from_obj, _ = cv2.Rodrigues(rvec)
|
||
R_obj_from_cam = R_cam_from_obj.T
|
||
t_obj_from_cam = -R_obj_from_cam @ tvec
|
||
|
||
transform_matrix = np.eye(4)
|
||
transform_matrix[:3, :3] = R_obj_from_cam
|
||
transform_matrix[:3, 3] = t_obj_from_cam.flatten()
|
||
|
||
return transform_matrix
|
||
|
||
|
||
def reconstruct_in_camera_coords(points_L, points_R, image_size=(4000, 3000)):
|
||
# ... (这个函数保持不变)
|
||
cam_L, cam_R, extrinsics = get_camera_parameters()
|
||
R1, R2, P1, P2, _, _, _ = cv2.stereoRectify(cam_L['K'], cam_L['kc'], cam_R['K'], cam_R['kc'], image_size,
|
||
extrinsics['R'], extrinsics['T'].flatten())
|
||
points_L_undistorted = cv2.undistortPoints(np.array(points_L, dtype=np.float32), cam_L['K'], cam_L['kc'], P=P1)
|
||
points_R_undistorted = cv2.undistortPoints(np.array(points_R, dtype=np.float32), cam_R['K'], cam_R['kc'], P=P2)
|
||
points_4d_hom = cv2.triangulatePoints(P1, P2, points_L_undistorted.reshape(-1, 2).T,
|
||
points_R_undistorted.reshape(-1, 2).T)
|
||
return (points_4d_hom[:3] / points_4d_hom[3]).T
|
||
|
||
|
||
def final_reconstruction_pipeline(all_2d_endpoints):
|
||
# 1. 计算唯一的、全局的变换矩阵
|
||
global_transform = get_global_transform_from_up_data(all_2d_endpoints)
|
||
if global_transform is None:
|
||
return None
|
||
|
||
reconstructed_4_seams = {}
|
||
|
||
for part_type in ['up', 'bottom']:
|
||
# 2. 对每个部分进行标准双目重建
|
||
points_L, points_R, seam_keys = [], [], []
|
||
for line_name in ['line1', 'line2']:
|
||
key_L = f"{part_type}_l_{line_name}"
|
||
key_R = f"{part_type}_r_{line_name}"
|
||
points_L.extend([all_2d_endpoints[key_L]['start'], all_2d_endpoints[key_L]['end']])
|
||
points_R.extend([all_2d_endpoints[key_R]['start'], all_2d_endpoints[key_R]['end']])
|
||
seam_keys.append(f"{part_type}_{line_name}")
|
||
|
||
points_camL = reconstruct_in_camera_coords(points_L, points_R)
|
||
|
||
# 3. 使用同一个全局矩阵进行变换
|
||
points_camL_hom = np.hstack([points_camL, np.ones((points_camL.shape[0], 1))])
|
||
points_object = (global_transform @ points_camL_hom.T).T[:, :3]
|
||
|
||
# 4. 整理结果
|
||
for i, key in enumerate(seam_keys):
|
||
reconstructed_4_seams[key] = {'start_3d': points_object[i * 2], 'end_3d': points_object[i * 2 + 1]}
|
||
|
||
return reconstructed_4_seams
|
||
|
||
|
||
|
||
|
||
def merge_seams(reconstructed_seams_dict):
|
||
"""
|
||
将重建出的四条焊缝合并为最终的三条焊缝模型。
|
||
|
||
Args:
|
||
reconstructed_seams_dict (dict): 包含 'up_line1', 'up_line2',
|
||
'bottom_line1', 'bottom_line2' 的字典。
|
||
|
||
Returns:
|
||
dict: 包含 'bottom_left', 'middle', 'top_left' 三条最终焊缝的字典。
|
||
"""
|
||
print("\n--- Merging seams into final 3-line model ---")
|
||
|
||
if not all(k in reconstructed_seams_dict for k in ['up_line1', 'up_line2', 'bottom_line1', 'bottom_line2']):
|
||
print("Error: Missing reconstructed seams for merging.")
|
||
return None
|
||
|
||
# 提取所有需要的端点
|
||
bl1_start = reconstructed_seams_dict['bottom_line1']['start_3d']
|
||
bl1_end = reconstructed_seams_dict['bottom_line1']['end_3d']
|
||
|
||
ul2_start = reconstructed_seams_dict['up_line2']['start_3d']
|
||
ul2_end = reconstructed_seams_dict['up_line2']['end_3d']
|
||
|
||
bl2_start = reconstructed_seams_dict['bottom_line2']['start_3d']
|
||
bl2_end = reconstructed_seams_dict['bottom_line2']['end_3d']
|
||
|
||
ul1_start = reconstructed_seams_dict['up_line1']['start_3d']
|
||
ul1_end = reconstructed_seams_dict['up_line1']['end_3d']
|
||
|
||
# --- 定义最终的三条线 ---
|
||
|
||
# 1. 左下焊缝 (bottom_left)
|
||
# 直接使用 bottom_line1。为了保证方向一致，我们让它从X值较小的点指向X值较大的点。
|
||
bottom_left_points = sorted([bl1_start, bl1_end], key=lambda p: p[0])
|
||
|
||
# 2. 中间焊缝 (middle)
|
||
# 这是 up_line2 和 bottom_line2 的合并。理论上它们应该重合。
|
||
# 我们可以取四个点的平均值来得到更稳健的起点和终点。
|
||
# 公共起点应该是 (bl1_end, ul2_start, bl2_start) 的平均值
|
||
middle_start = np.mean([bl1_end, ul2_start, bl2_start], axis=0)
|
||
# 公共终点应该是 (ul2_end, bl2_end, ul1_start) 的平均值
|
||
middle_end = np.mean([ul2_end, bl2_end, ul1_start], axis=0)
|
||
|
||
# 3. 左上焊缝 (top_left)
|
||
# 直接使用 up_line1。
|
||
top_left_points = [ul1_start, ul1_end] # 保持原始方向
|
||
|
||
final_3_seams = {
|
||
'bottom_left': {
|
||
'start_3d': bottom_left_points[0],
|
||
'end_3d': bottom_left_points[1]
|
||
},
|
||
'middle': {
|
||
'start_3d': middle_start,
|
||
'end_3d': middle_end
|
||
},
|
||
'top_left': {
|
||
'start_3d': top_left_points[0],
|
||
'end_3d': top_left_points[1]
|
||
}
|
||
}
|
||
|
||
return final_3_seams |