位姿可视化

visualize_localization.py

@@ -0,0 +1,611 @@
+"""
+可视化重定位效果：点云、轨迹、预测位置与实际位置对比
+支持 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)