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)