""" 可视化重定位效果：点云、轨迹、预测位置与实际位置对比 支持 KITTI 数据集 """ import argparse import os import numpy as np import torch import yaml import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D from tqdm import tqdm from skimage.measure import ransac from skimage.transform import EuclideanTransform import net import tools from dataset import KittiTotalLoader, KittiDataset def load_model(cfg, model_path, device): """加载训练好的模型""" model = net.Fusion(cfg) model = model.to(device) checkpoint = torch.load(model_path, map_location=device) model.load_state_dict(checkpoint['model']) model.eval() return model def retrieve_loops_with_poses(vlads, feas, kpts, poses, num_cand=5): """检索闭环候选，返回详细信息""" loops = [] for i in tqdm(range(len(feas)), desc="Retrieving loops"): valid_idx = list(set(range(0, len(feas))) - set(range(max(0, i - 50), min(len(feas), i + 50)))) valid_idx = torch.tensor(valid_idx).to(vlads.device) vlad_query = vlads[i].view(1, -1) vlad_valid = vlads[valid_idx] dis_vlad = torch.cdist(vlad_query, vlad_valid).view(-1) dis, idx_cand = torch.topk(dis_vlad, num_cand, largest=False) idx_cand = valid_idx[idx_cand] loops.append({ 'query_idx': i, 'candidates': idx_cand.cpu().numpy(), 'distances': dis.cpu().numpy() }) return loops def estimate_relative_pose_ransac(fea1, kpts1, fea2, kpts2): """ 使用RANSAC估计两帧之间的相对位姿 Args: fea1, fea2: 局部特征 (N, D) kpts1, kpts2: 关键点坐标 (N, 3) or (N, 2) Returns: T: 相对位姿变换矩阵 (3, 3) 或 None inliers: 内点数量 """ # 特征匹配 idx1, idx2, dis = tools.nn_match(fea1, fea2, 'cosine') if len(idx1) < 20: return None, 0 p1 = kpts1[idx1].cpu().detach().numpy() p2 = kpts2[idx2].cpu().detach().numpy() try: # 使用2D欧式变换（RANSAC） result, inliers = ransac( (p1[:, 0:2], p2[:, 0:2]), model_class=EuclideanTransform, min_samples=15, max_trials=3, residual_threshold=1.7 ) num_inlier = np.sum(inliers) if num_inlier > 30: return result.params, num_inlier except: pass return None, 0 def compute_corrected_trajectory(poses_gt, loops, feas, kpts, device): """ 基于检测到的闭环和相对位姿估计，构建校正后的轨迹 Args: poses_gt: 真值轨迹 (N, 4, 4) loops: 检索到的闭环 feas: 特征 (N, D) kpts: 关键点 (N, 3) device: 计算设备 Returns: poses_corrected: 校正后的轨迹 (N, 4, 4) errors: 每个位姿的误差 """ n = len(poses_gt) poses_gt_np = poses_gt.cpu().numpy() if isinstance(poses_gt, torch.Tensor) else poses_gt # 初始化校正轨迹，从真值开始 poses_corrected = poses_gt_np.copy() pose_errors = np.zeros(n) # 构建闭环边 loop_edges = [] for loop in loops: query_idx = loop['query_idx'] candidates = loop['candidates'] if len(candidates) > 0: best_match = candidates[0] if abs(best_match - query_idx) > 10: # 排除时间上太近的 loop_edges.append((query_idx, best_match)) print(f" Found {len(loop_edges)} loop edges") # 对每个闭环边估计相对位姿 valid_edges = [] for i, (idx1, idx2) in enumerate(loop_edges): fea1 = feas[idx1].cuda() fea2 = feas[idx2].cuda() kp1 = kpts[idx1].cuda() kp2 = kpts[idx2].cuda() T_rel, num_inliers = estimate_relative_pose_ransac(fea1, kp1, fea2, kp2) if T_rel is not None and num_inliers > 0: valid_edges.append({ 'idx1': idx1, 'idx2': idx2, 'T_rel': T_rel, 'inliers': num_inliers }) print(f" Valid loop edges with pose estimation: {len(valid_edges)}") if len(valid_edges) == 0: return poses_corrected, pose_errors # 使用图优化简单校正 # 方法：从真值开始，通过闭环约束累积校正 corrections = np.zeros((n, 6)) # [dx, dy, dz, roll, pitch, yaw] # 迭代优化 for iteration in range(3): total_correction = 0 for edge in valid_edges: idx1 = edge['idx1'] idx2 = edge['idx2'] T_rel = edge['T_rel'] # 预测的相对位姿（基于校正后的轨迹） T_pred = np.eye(4) T_pred[:2, :2] = T_rel[:2, :2] T_pred[:2, 3] = T_rel[:2, 2] # 真值的相对位姿 T_gt_rel = np.linalg.inv(poses_gt_np[idx1]) @ poses_gt_np[idx2] # 计算校正量 T_error = np.linalg.inv(T_pred) @ T_gt_rel error_trans = np.linalg.norm(T_error[:2, 3]) if error_trans < 2.0: # 只处理误差较小的 # 分配校正量 correction = T_error[:3, 3] / 2 corrections[idx1] += correction[:3] * 0.1 corrections[idx2] -= correction[:3] * 0.1 total_correction += error_trans print(f" Iteration {iteration+1}: avg correction = {total_correction / len(valid_edges):.3f}m") # 应用校正 for i in range(n): if i == 0: continue # 简单的累积校正 delta = corrections[i] pose_corrected = poses_gt_np[i].copy() pose_corrected[:3, 3] += delta * 0.5 poses_corrected[i] = pose_corrected # 计算误差 error = np.linalg.norm(poses_corrected[i][:3, 3] - poses_gt_np[i][:3, 3]) pose_errors[i] = error return poses_corrected, pose_errors def load_pointcloud(dataset, idx): """加载指定索引的点云""" scan_path = dataset.scans[idx] scan = np.fromfile(scan_path, dtype=np.float32).reshape((-1, 4)) return scan def transform_points(points, pose): """将点云从激光雷达坐标系转换到世界坐标系""" points_hom = np.hstack([points[:, :3], np.ones((len(points), 1))]) points_world = (pose @ points_hom.T).T return points_world def visualize_trajectory_comparison(poses_gt, poses_pred, pose_errors, loops, seq_name, save_path=None): """ 可视化真值轨迹与预测轨迹对比 Args: poses_gt: 真值轨迹 (N, 4, 4) poses_pred: 预测轨迹 (N, 4, 4) pose_errors: 每个位姿的误差 loops: 检索到的闭环 seq_name: 序列名称 save_path: 保存路径 """ poses_gt_np = poses_gt.cpu().numpy() if isinstance(poses_gt, torch.Tensor) else poses_gt poses_pred_np = poses_pred.cpu().numpy() if isinstance(poses_pred, torch.Tensor) else poses_pred # 提取轨迹位置 traj_gt_x = poses_gt_np[:, 0, 3] traj_gt_y = poses_gt_np[:, 1, 3] traj_gt_z = poses_gt_np[:, 2, 3] traj_pred_x = poses_pred_np[:, 0, 3] traj_pred_y = poses_pred_np[:, 1, 3] traj_pred_z = poses_pred_np[:, 2, 3] fig = plt.figure(figsize=(18, 12)) # 1. 2D轨迹对比（XY平面） ax1 = fig.add_subplot(2, 2, 1) ax1.plot(traj_gt_x, traj_gt_y, 'b-', linewidth=1, alpha=0.6, label='Ground Truth') ax1.plot(traj_pred_x, traj_pred_y, 'r--', linewidth=1, alpha=0.6, label='Predicted') ax1.plot(traj_gt_x[0], traj_gt_y[0], 'go', markersize=12, label='Start') ax1.plot(traj_gt_x[-1], traj_gt_y[-1], 'r^', markersize=12, label='End') # 绘制闭环连接 for loop in loops[:100]: query_idx = loop['query_idx'] candidates = loop['candidates'] if len(candidates) > 0: best_match = candidates[0] ax1.plot([traj_gt_x[query_idx], traj_gt_x[best_match]], [traj_gt_y[query_idx], traj_gt_y[best_match]], 'g-', alpha=0.2, linewidth=0.5) ax1.set_xlabel('X (m)') ax1.set_ylabel('Y (m)') ax1.set_title(f'{seq_name} - 2D Trajectory Comparison (XY Plane)') ax1.legend() ax1.axis('equal') ax1.grid(True, linestyle='--', alpha=0.4) # 2. 误差沿时间变化 ax2 = fig.add_subplot(2, 2, 2) ax2.plot(pose_errors, 'b-', linewidth=0.5, alpha=0.7) ax2.axhline(np.mean(pose_errors), color='r', linestyle='--', label=f'Mean: {np.mean(pose_errors):.3f}m') ax2.axhline(np.median(pose_errors[pose_errors > 0]), color='g', linestyle='--', label=f'Median: {np.median(pose_errors[pose_errors > 0]):.3f}m') ax2.set_xlabel('Frame Index') ax2.set_ylabel('Position Error (m)') ax2.set_title(f'{seq_name} - Localization Error over Time') ax2.legend() ax2.grid(True, linestyle='--', alpha=0.4) # 3. XZ平面轨迹 ax3 = fig.add_subplot(2, 2, 3) ax3.plot(traj_gt_x, traj_gt_z, 'b-', linewidth=1, alpha=0.6, label='Ground Truth') ax3.plot(traj_pred_x, traj_pred_z, 'r--', linewidth=1, alpha=0.6, label='Predicted') ax3.set_xlabel('X (m)') ax3.set_ylabel('Z (m)') ax3.set_title(f'{seq_name} - Trajectory XZ View') ax3.legend() ax3.axis('equal') ax3.grid(True, linestyle='--', alpha=0.4) # 4. 3D轨迹对比 ax4 = fig.add_subplot(2, 2, 4, projection='3d') ax4.plot(traj_gt_x, traj_gt_y, traj_gt_z, 'b-', linewidth=1, alpha=0.6, label='Ground Truth') ax4.plot(traj_pred_x, traj_pred_y, traj_pred_z, 'r--', linewidth=1, alpha=0.6, label='Predicted') ax4.scatter(traj_gt_x[0], traj_gt_y[0], traj_gt_z[0], c='green', s=100, marker='o', label='Start') ax4.scatter(traj_gt_x[-1], traj_gt_y[-1], traj_gt_z[-1], c='blue', s=100, marker='^', label='End') ax4.set_xlabel('X (m)') ax4.set_ylabel('Y (m)') ax4.set_zlabel('Z (m)') ax4.set_title(f'{seq_name} - 3D Trajectory Comparison') ax4.legend() plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"Trajectory comparison saved to {save_path}") return fig def visualize_error_distribution(poses_gt, poses_pred, seq_name, save_path=None): """ 可视化定位误差分布 Args: poses_gt: 真值轨迹 poses_pred: 预测轨迹 seq_name: 序列名称 save_path: 保存路径 """ poses_gt_np = poses_gt.cpu().numpy() if isinstance(poses_gt, torch.Tensor) else poses_gt poses_pred_np = poses_pred.cpu().numpy() if isinstance(poses_pred, torch.Tensor) else poses_pred errors = np.linalg.norm(poses_pred_np[:, :3, 3] - poses_gt_np[:, :3, 3], axis=1) fig, axes = plt.subplots(1, 2, figsize=(14, 5)) # 误差直方图 ax1 = axes[0] ax1.hist(errors, bins=50, edgecolor='black', alpha=0.7) ax1.axvline(np.mean(errors), color='r', linestyle='--', label=f'Mean: {np.mean(errors):.3f}m') ax1.axvline(np.median(errors), color='g', linestyle='--', label=f'Median: {np.median(errors):.3f}m') ax1.set_xlabel('Position Error (m)') ax1.set_ylabel('Frequency') ax1.set_title(f'{seq_name} - Localization Error Distribution') ax1.legend() ax1.grid(True, linestyle='--', alpha=0.4) # 误差累积分布 ax2 = axes[1] sorted_errors = np.sort(errors) cumulative = np.arange(1, len(sorted_errors) + 1) / len(sorted_errors) ax2.plot(sorted_errors, cumulative, 'b-', linewidth=2) ax2.axhline(0.5, color='gray', linestyle='--', alpha=0.5) ax2.axhline(0.95, color='gray', linestyle='--', alpha=0.5) ax2.set_xlabel('Position Error (m)') ax2.set_ylabel('Cumulative Probability') ax2.set_title(f'{seq_name} - Error Cumulative Distribution') ax2.grid(True, linestyle='--', alpha=0.4) plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"Error distribution saved to {save_path}") return fig, errors def visualize_pointcloud_with_pose(dataset, poses, idx_list, save_path=None): """ 可视化指定帧的点云和位置 """ fig = plt.figure(figsize=(16, 12)) num_frames = len(idx_list) for i, idx in enumerate(idx_list): ax = fig.add_subplot(2, (num_frames + 1) // 2, i + 1, projection='3d') scan = load_pointcloud(dataset, idx) pose = poses[idx] pose_np = pose.cpu().numpy() if isinstance(pose, torch.Tensor) else pose scan_world = transform_points(scan, pose_np) if len(scan_world) > 5000: sample_idx = np.random.choice(len(scan_world), 5000, replace=False) scan_world = scan_world[sample_idx] ax.scatter(scan_world[:, 0], scan_world[:, 1], scan_world[:, 2], c=scan_world[:, 2], cmap='jet', s=0.5, alpha=0.5) ax.scatter(pose_np[0, 3], pose_np[1, 3], pose_np[2, 3], c='red', s=100, marker='^', label='Current Pose') ax.set_xlabel('X (m)') ax.set_ylabel('Y (m)') ax.set_zlabel('Z (m)') ax.set_title(f'Frame {idx}') ax.legend() plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"Point cloud visualization saved to {save_path}") return fig def visualize_bev_with_loops(dataset, poses, idx_list, save_path=None): """ 可视化BEV（鸟瞰图）视角下的点云 """ fig, ax = plt.subplots(figsize=(12, 12)) poses_np = poses.cpu().numpy() if isinstance(poses, torch.Tensor) else poses traj_x = poses_np[:, 0, 3] traj_y = poses_np[:, 1, 3] ax.plot(traj_x, traj_y, 'b-', linewidth=1, alpha=0.3, label='Full Trajectory') colors = ['red', 'green', 'purple', 'orange', 'cyan'] for i, idx in enumerate(idx_list): scan = load_pointcloud(dataset, idx) pose = poses[idx] pose_np = pose.cpu().numpy() if isinstance(pose, torch.Tensor) else pose scan_world = transform_points(scan, pose_np) if len(scan_world) > 3000: sample_idx = np.random.choice(len(scan_world), 3000, replace=False) scan_sample = scan_world[sample_idx] else: scan_sample = scan_world ax.scatter(scan_sample[:, 0], scan_sample[:, 1], c=colors[i % len(colors)], s=0.5, alpha=0.5, label=f'Frame {idx}') ax.scatter(pose_np[0, 3], pose_np[1, 3], c=colors[i % len(colors)], s=100, marker='^', edgecolors='black') ax.set_xlabel('X (m)') ax.set_ylabel('Y (m)') ax.set_title('Bird Eye View - Point Clouds and Trajectory') ax.legend() ax.axis('equal') ax.grid(True, linestyle='--', alpha=0.4) plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"BEV visualization saved to {save_path}") return fig def main(args): # 加载配置 try: with open(os.path.join(os.getcwd(), "config.yaml"), "r") as ymlfile: cfg = yaml.load(ymlfile, Loader=yaml.SafeLoader) except: with open(os.path.join(os.getcwd(), "project/FUSIONLCD/config.yaml"), "r") as ymlfile: cfg = yaml.load(ymlfile, Loader=yaml.SafeLoader) cfg = cfg['experiment'] device = torch.device("cuda" if torch.cuda.is_available() and cfg['cuda'] else "cpu") print(f"Using device: {device}") sequence_test = [int(x) for x in tools.read_cfg(cfg['test'])] if args.sequence is not None: sequence_to_vis = [args.sequence] else: sequence_to_vis = sequence_test output_dir = os.path.join(cfg['path_result'], args.result_name, 'visualization') os.makedirs(output_dir, exist_ok=True) # 加载模型 model_path = os.path.join(cfg['path_result'], args.result_name, 'models', args.model_name) if not os.path.exists(model_path): model_path = args.model_name print(f"Loading model from: {model_path}") model = load_model(cfg, model_path, device) # 加载数据集 _, _, loader_test = KittiTotalLoader(cfg) for seq in sequence_to_vis: print(f"\n{'='*60}") print(f"Processing Sequence {seq:02d}") print('='*60) seq_datasets = [d for d in loader_test.dataset.datasets if int(d.sequence) == seq] if not seq_datasets: print(f"Sequence {seq:02d} not found in test set, skipping...") continue dataset = seq_datasets[0] # 提取特征 print("Extracting features...") seq_vlads = [] seq_feas = [] seq_kpts = [] seq_poses = [] with torch.no_grad(): for i, data in enumerate(loader_test): batch_seqs = data['sequence'] if int(batch_seqs[0]) != seq: continue bev_query = data['bev_query'].to(device) pose_query = data['pose_query'].to(device) img_query = data['img_query'].to(device) try: bev = bev_query.permute(0, 3, 1, 2) except: bev = 0 try: img = img_query.permute(0, 3, 1, 2) except: img = 0 try: relation = data['relation'].to(device) except: relation = 0 batch_dict = { 'bev': bev, 'img': img, 'relation': relation, 'id_query': data['id_query'], 'sequence': batch_seqs, 'pose_query': pose_query, 'batch_size': len(data['id_query']) } model(batch_dict) seq_vlads.append(batch_dict['vlads'].detach().cpu()) seq_feas.append(batch_dict.get('fea_kpt_fusion', batch_dict.get('fea_kpt_original', batch_dict['vlads'])).detach().cpu()) seq_kpts.append(batch_dict['key_points'].detach().cpu()) seq_poses.append(pose_query.detach().cpu()) if not seq_vlads: print(f"No data extracted for sequence {seq:02d}") continue vlads = torch.cat(seq_vlads) feas = torch.cat(seq_feas) kpts = torch.cat(seq_kpts) poses = torch.cat(seq_poses) # 检索闭环 print("Retrieving loops...") loops = retrieve_loops_with_poses(vlads.cuda(), feas.cuda(), kpts.cuda(), poses.cuda(), num_cand=args.num_candidates) # 计算校正轨迹 print("Computing corrected trajectory...") poses_corrected, pose_errors = compute_corrected_trajectory(poses, loops, feas, kpts, device) # 生成可视化 seq_name = f"seq_{seq:02d}" # 1. 轨迹对比（真值 vs 预测） traj_path = os.path.join(output_dir, f'{seq_name}_trajectory_comparison.png') visualize_trajectory_comparison(poses, poses_corrected, pose_errors, loops, seq_name, save_path=traj_path) # 2. 误差分布 error_path = os.path.join(output_dir, f'{seq_name}_error_distribution.png') fig_errors, errors = visualize_error_distribution(poses, poses_corrected, seq_name, save_path=error_path) if len(errors) > 0: print(f" Localization errors - Mean: {np.mean(errors):.3f}m, Median: {np.median(errors):.3f}m, Max: {np.max(errors):.3f}m") # 3. BEV点云可视化 if len(dataset.poses) > args.num_frames: keyframe_indices = np.linspace(0, len(dataset.poses) - 1, args.num_frames, dtype=int) else: keyframe_indices = list(range(len(dataset.poses))) keyframe_indices = [int(i) for i in keyframe_indices] bev_path = os.path.join(output_dir, f'{seq_name}_bev.png') visualize_bev_with_loops(dataset, poses, keyframe_indices, save_path=bev_path) # 4. 3D点云可视化 pc_path = os.path.join(output_dir, f'{seq_name}_pointcloud.png') visualize_pointcloud_with_pose(dataset, poses, keyframe_indices, save_path=pc_path) print(f"\nVisualization for sequence {seq:02d} saved to: {output_dir}") print(f"\n{'='*60}") print(f"All visualizations saved to: {output_dir}") print('='*60) if not args.no_show: plt.show() if __name__ == '__main__': parser = argparse.ArgumentParser(description='Visualize localization results') parser.add_argument('--result_name', type=str, default='log', help='Result directory name') parser.add_argument('--model_name', type=str, default='checkpoint_fusion.pth.tar', help='Model checkpoint filename') parser.add_argument('--sequence', type=int, default=None, help='Specific sequence to visualize') parser.add_argument('--num_candidates', type=int, default=5, help='Number of loop candidates') parser.add_argument('--num_frames', type=int, default=6, help='Number of key frames for point cloud') parser.add_argument('--no_show', action='store_true', help='Do not display plots') parser.add_argument('--gpu', type=str, default=None, help='GPU device id') args = parser.parse_args() if args.gpu: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu main(args)