186 lines
8.0 KiB
Python
186 lines
8.0 KiB
Python
import os
|
||
import cv2
|
||
import numpy as np
|
||
|
||
from script.yolo_detector import detect_crop_area
|
||
from script.linknet_segmentor import segment_and_find_endpoints
|
||
from script.reconstruction import visualize_reconstructed_seams, reconstruct_points
|
||
from script.pose_estimation import get_ground_truth_seams, reproject_to_object_coords # 我们只需要真值
|
||
# 导入我们最终的重建流程
|
||
from script.final_reconstruction import final_reconstruction_pipeline, merge_seams
|
||
import itertools # 确保导入itertools
|
||
|
||
|
||
from script.global_optimizer import run_global_optimization, merge_seams
|
||
from script.pose_estimation import get_ground_truth_seams
|
||
from script.reconstruction import visualize_reconstructed_seams
|
||
|
||
def run_full_recognition_pipeline():
|
||
"""
|
||
运行完整的识别流程：YOLO定位 -> LinkNet分割 -> 端点提取。
|
||
"""
|
||
# 1. 定义路径
|
||
base_dir = os.path.dirname(os.path.abspath(__file__))
|
||
data_map = {
|
||
'up': {
|
||
'l_img': os.path.join(base_dir, 'data', 'origin', 'up', 'l1.jpeg'),
|
||
'r_img': os.path.join(base_dir, 'data', 'origin', 'up', 'r1.jpeg'),
|
||
'yolo_model': os.path.join(base_dir, 'module', 'yolov8', 'up.pt'),
|
||
'linknet_models': {
|
||
'line1': os.path.join(base_dir, 'module', 'linknet', 'best_linknet_up_model_line1.pth'),
|
||
'line2': os.path.join(base_dir, 'module', 'linknet', 'best_linknet_up_model_line2.pth')
|
||
}
|
||
},
|
||
'bottom': {
|
||
'l_img': os.path.join(base_dir, 'data', 'origin', 'bottom', 'l1.jpeg'),
|
||
'r_img': os.path.join(base_dir, 'data', 'origin', 'bottom', 'r1.jpeg'),
|
||
'yolo_model': os.path.join(base_dir, 'module', 'yolov8', 'bottom.pt'),
|
||
'linknet_models': {
|
||
'line1': os.path.join(base_dir, 'module', 'linknet', 'best_linknet_bottom_model_line1.pth'),
|
||
'line2': os.path.join(base_dir, 'module', 'linknet', 'best_linknet_bottom_model_line2.pth')
|
||
}
|
||
}
|
||
}
|
||
output_dir = os.path.join(base_dir, 'data', 'processed')
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
all_endpoints = {}
|
||
|
||
for part, paths in data_map.items():
|
||
print(f"\n--- Processing '{part}' part ---")
|
||
for img_path, side in [(paths['l_img'], 'l'), (paths['r_img'], 'r')]:
|
||
print(f"\n-- Analyzing image: {os.path.basename(img_path)} --")
|
||
|
||
crop_box = detect_crop_area(img_path, paths['yolo_model'])
|
||
if not crop_box:
|
||
print(f"Skipping further processing for {os.path.basename(img_path)}.")
|
||
continue
|
||
|
||
original_image_vis = cv2.imread(img_path)
|
||
|
||
for line_name, linknet_path in paths['linknet_models'].items():
|
||
endpoints = segment_and_find_endpoints(original_image_vis, crop_box, linknet_path)
|
||
|
||
if endpoints:
|
||
start_pt, end_pt = endpoints
|
||
result_key = f"{part}_{side}_{line_name}"
|
||
all_endpoints[result_key] = {'start': start_pt, 'end': end_pt}
|
||
|
||
# --- 在可视化图像上绘制结果 (增强版) ---
|
||
# 1. 绘制端点圆圈
|
||
cv2.circle(original_image_vis, start_pt, 15, (0, 255, 0), -1) # 绿色起点
|
||
cv2.circle(original_image_vis, end_pt, 15, (0, 0, 255), -1) # 红色终点
|
||
# 2. 绘制连接线
|
||
cv2.line(original_image_vis, start_pt, end_pt, (255, 0, 0), 4)
|
||
|
||
# 3. 添加文本标签
|
||
# 计算线段中点作为文本放置位置
|
||
mid_point = ((start_pt[0] + end_pt[0]) // 2, (start_pt[1] + end_pt[1]) // 2)
|
||
# 在中点上方放置文本
|
||
text_pos = (mid_point[0], mid_point[1] - 20)
|
||
cv2.putText(original_image_vis,
|
||
result_key,
|
||
text_pos,
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
2, # 字体大小
|
||
(255, 255, 0), # 字体颜色 (青色)
|
||
4, # 字体粗细
|
||
cv2.LINE_AA)
|
||
|
||
# 绘制YOLO框并保存最终的可视化结果
|
||
cv2.rectangle(original_image_vis, (crop_box[0], crop_box[1]), (crop_box[2], crop_box[3]), (0, 255, 255), 4)
|
||
save_path = os.path.join(output_dir, f'{part}_{side}_final_result.jpg')
|
||
cv2.imwrite(save_path, original_image_vis)
|
||
print(f"Saved final visualization to {save_path}")
|
||
# 3. 打印总结
|
||
print("\n--- Final Endpoints Summary (in original image coordinates) ---")
|
||
for name, points in all_endpoints.items():
|
||
print(f"{name}: Start={points['start']}, End={points['end']}")
|
||
|
||
return all_endpoints
|
||
|
||
|
||
def run_3d_reconstruction(all_2d_endpoints):
|
||
"""
|
||
根据识别出的2D端点，重建出三维焊缝。
|
||
"""
|
||
print("\n--- Starting 3D Reconstruction ---")
|
||
|
||
# 这个字典将存储最终的三维坐标
|
||
reconstructed_seams_3d = {}
|
||
|
||
# 需要重建的焊缝对
|
||
# 例如：'up_line1' 对应 up_l_line1 和 up_r_line1
|
||
seam_pairs = ['up_line1', 'up_line2', 'bottom_line1', 'bottom_line2']
|
||
|
||
for seam_name in seam_pairs:
|
||
key_L = f"{seam_name.split('_')[0]}_l_{seam_name.split('_')[1]}" # e.g., 'up_l_line1'
|
||
key_R = f"{seam_name.split('_')[0]}_r_{seam_name.split('_')[1]}" # e.g., 'up_r_line1'
|
||
|
||
# 检查左右相机的点是否都已识别
|
||
if key_L not in all_2d_endpoints or key_R not in all_2d_endpoints:
|
||
print(f"Warning: Missing points for seam '{seam_name}'. Cannot reconstruct.")
|
||
continue
|
||
|
||
# 准备输入点列表：[start_point, end_point]
|
||
points_L = [all_2d_endpoints[key_L]['start'], all_2d_endpoints[key_L]['end']]
|
||
points_R = [all_2d_endpoints[key_R]['start'], all_2d_endpoints[key_R]['end']]
|
||
|
||
# 调用重建函数
|
||
# 假设你的图像尺寸是 4000x3000，如果不是，请修改
|
||
# 这是一个重要的参数，需要与标定时使用的图像尺寸一致！
|
||
points_3d = reconstruct_points(points_L, points_R, image_size=(4000, 3000))
|
||
|
||
reconstructed_seams_3d[seam_name] = {
|
||
'start_3d': points_3d[0],
|
||
'end_3d': points_3d[1]
|
||
}
|
||
|
||
# --- 打印最终的三维坐标结果 ---
|
||
print("\n--- Final 3D Seam Endpoints (in Left Camera Coordinate System, unit: mm) ---")
|
||
for name, points in reconstructed_seams_3d.items():
|
||
start_str = np.array2string(points['start_3d'], formatter={'float_kind': lambda x: "%.3f" % x})
|
||
end_str = np.array2string(points['end_3d'], formatter={'float_kind': lambda x: "%.3f" % x})
|
||
print(f"{name}:")
|
||
print(f" Start 3D: {start_str}")
|
||
print(f" End 3D: {end_str}")
|
||
|
||
return reconstructed_seams_3d
|
||
|
||
|
||
def run_new_reconstruction_pipeline(all_2d_endpoints):
|
||
"""
|
||
使用 solvePnP 的全新重建和拼接流程。
|
||
"""
|
||
print("\n--- Starting NEW Reconstruction Pipeline (with solvePnP) ---")
|
||
|
||
# --- 处理上半部分 ---
|
||
print("\nProcessing 'up' part...")
|
||
reconstructed_up = reproject_to_object_coords(all_2d_endpoints, all_2d_endpoints, part_type='up')
|
||
|
||
# --- 处理下半部分 ---
|
||
print("\nProcessing 'bottom' part...")
|
||
reconstructed_bottom = reproject_to_object_coords(all_2d_endpoints, all_2d_endpoints, part_type='bottom')
|
||
|
||
# --- 合并结果 ---
|
||
final_reconstructed_seams = {}
|
||
if reconstructed_up:
|
||
final_reconstructed_seams.update(reconstructed_up)
|
||
if reconstructed_bottom:
|
||
final_reconstructed_seams.update(reconstructed_bottom)
|
||
|
||
return final_reconstructed_seams
|
||
|
||
|
||
if __name__ == '__main__':
|
||
final_2d_endpoints = run_full_recognition_pipeline()
|
||
ground_truth = get_ground_truth_seams()
|
||
|
||
final_4_seams = {}
|
||
if final_2d_endpoints:
|
||
# 直接调用全局优化
|
||
final_4_seams = run_global_optimization(final_2d_endpoints, ground_truth)
|
||
|
||
final_3_seam_model = {}
|
||
if final_4_seams:
|
||
final_3_seam_model = merge_seams(final_4_seams) |