""" Test script for TAPFormer model on Aedat4 dataset. This script evaluates the model performance on test sequences. Usage: python test_aedat4.py [--config config_aedat4.yaml] """ import os import sys import argparse import yaml import torch import time import numpy as np from LFE_TAP.datasets.TAPFormer_dataset import TAPFormer_dataset from LFE_TAP.evaluator.prediction import TAPFormer_online from LFE_TAP.evaluator.evaluation_pred import EvaluationPredictor from LFE_TAP.utils.visualizer import Visualizer from LFE_TAP.evaluator.evaluator import compute_tapvid_metrics DEFAULT_DEVICE = ('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu') def load_config(config_path): """Load configuration from YAML file.""" with open(config_path, 'r', encoding='utf-8') as f: config = yaml.safe_load(f) return config def parse_args(): """Parse command line arguments.""" parser = argparse.ArgumentParser(description='Test TAPFormer model on Aedat4 dataset') parser.add_argument('--config', type=str, default='config/config_InivTAP_DrivTAP.yaml', help='Path to configuration YAML file') return parser.parse_args() # Parse arguments and load config args = parse_args() config = load_config(args.config) # Extract configuration dataset_dir = config['dataset_dir'] ckpt_root = config['ckpt_root'] EVAL_DATASETS = config['eval_datasets'] model_name = os.path.basename(os.path.dirname(ckpt_root)) representation = config['representation'] stride = config.get('stride') # Auto-detect backbone corr_levels = config.get('corr_levels') backbone = config.get('backbone') # Model configuration window_size = config.get('window_size', 16) corr_radius = config.get('corr_radius', 3) hidden_size = config.get('hidden_size', 384) space_depth = config.get('space_depth', 3) time_depth = config.get('time_depth', 3) # Evaluation settings dt = config.get('dt', 0.0100) grid_size = config.get('grid_size', 0) n_iters = config.get('n_iters', 5) # Visualization and output settings vis_cfg = config.get('visualization', {}) output_cfg = config.get('output', {}) enable_visualization = vis_cfg.get('enable', False) save_results = output_cfg.get('save_results', False) save_trajectory = output_cfg.get('save_trajectory', False) base_output_dir = output_cfg.get('base_dir', 'output/eval_aedat4_subseq') # ========== Model Initialization ========== print("Loading model...") model = TAPFormer_online( window_size=window_size, stride=stride, corr_radius=corr_radius, corr_levels=corr_levels, backbone=backbone, hidden_size=hidden_size, space_depth=space_depth, time_depth=time_depth ) # Load checkpoint state_dict = torch.load(ckpt_root, map_location=DEFAULT_DEVICE) if "model" in state_dict: state_dict = state_dict["model"] model.load_state_dict(state_dict, strict=False) model.eval() print("Model loaded successfully!") # ========== Evaluation ========== print("\n" + "="*50) print("Evaluating on Aedat4 dataset...") print("="*50) datasets = TAPFormer_dataset(os.path.join(dataset_dir), representation=representation, dt=dt) for seq_name in EVAL_DATASETS: sample, gotit = datasets.get_a_seq(seq_name) if not gotit: continue # Setup output directory output_dir = os.path.join(base_output_dir, seq_name, model_name) if enable_visualization or save_results or save_trajectory: os.makedirs(output_dir, exist_ok=True) # Initialize visualizer if needed vis = None if enable_visualization: vis = Visualizer(output_dir, fps=vis_cfg.get('fps', 20)) predictor = EvaluationPredictor( model, grid_size=grid_size, local_grid_size=0, single_point=False, num_uniformly_sampled_pts=0, n_iters=n_iters, ) if torch.cuda.is_available(): predictor.model = predictor.model.cuda() queries = sample.query_points[np.newaxis, ...] queries = queries.to(DEFAULT_DEVICE) sample.video = sample.video[np.newaxis, ...] sample.events = sample.events[np.newaxis, ...] sample.trajectory = sample.trajectory[np.newaxis, ...] sample.visibility = sample.visibility[np.newaxis, ...] start = time.time() pred_tracks = predictor(sample.video, sample.events, queries, img_ifnew=sample.img_ifnew) end = time.time() elapsed_time = (end-start)/sample.events.shape[1] print("time per frame:", elapsed_time) if isinstance(pred_tracks, tuple): pred_trajectory, pred_visibility, _ = pred_tracks else: pred_visibility = None if pred_visibility is None: pred_visibility = torch.zeros_like(sample.visibility) if not pred_visibility.dtype == torch.bool: pred_visibility = pred_visibility > 0.8 pred_occluded = ( torch.logical_not(pred_visibility.clone().permute(0, 2, 1)) .cpu() .numpy() ) pred_tracks = pred_trajectory.permute(0, 2, 1, 3).cpu().numpy() query_points = np.concatenate( ( np.zeros_like(sample.trajectory[:, 0, :, :1]), sample.trajectory[:, 0, :, 1:], ), axis=2, ) gt_tracks = np.transpose(sample.trajectory[:, :, :, 1:].copy(), (0, 2, 1, 3)) gt_occluded = np.transpose(sample.visibility.copy(), (0, 2, 1)) def expand_trajectory(traj, vis): """ traj: numpy array of shape (B, N, T, 2) return: numpy array of shape (B, N, 2*T, 2) """ B, N, T, _ = traj.shape # 初始化结果 expanded = np.zeros((B, N, 2*T, 2), dtype=traj.dtype) expanded_vis = np.zeros((B, N, 2*T), dtype=vis.dtype) expanded_vis[:, :, ::2] = vis expanded_vis[:, :, 1::2] = vis # 原始帧放在偶数位置 expanded[:, :, ::2, :] = traj # 计算速度（v_t = p_t - p_{t-1}） vel = np.zeros_like(traj) vel[:, :, 1:, :] = traj[:, :, 1:, :] - traj[:, :, :-1, :] # 用上一帧速度估算插值帧位置 # p_{t+0.5} = p_t + 0.5 * v_t expanded[:, :, 1::2, :] = traj + 0.5 * vel # 对 t=0 的插值帧（expanded[:, :, 1, :]），因为没有前一帧速度，保持不变 expanded[:, :, 1, :] = traj[:, :, 0, :] return expanded, expanded_vis if pred_trajectory.shape[1] != sample.trajectory.shape[1]: pred_trajectory_new = np.zeros_like(gt_tracks) pred_occluded_new = np.zeros_like(gt_occluded) for i in range(pred_tracks.shape[1]): pred_traj = pred_tracks[:, i, :, :] gt_t = sample.trajectory[:, :, i, 0].squeeze() est_t = sample.segmentation * 1e-6 pred_traj_x, pred_traj_y = pred_traj.squeeze().T pred_traj_x_ = np.interp(gt_t, est_t, pred_traj_x) pred_traj_y_ = np.interp(gt_t, est_t, pred_traj_y) pred_traj_ = np.stack((pred_traj_x_, pred_traj_y_), axis=1) pred_trajectory_new[:, i, :, :] = pred_traj_ pred_occ = pred_occluded[:, i, :].squeeze() indices = np.searchsorted(est_t, gt_t, side='left') # 处理左边界 (新时间点小于所有原始时间) left_mask = indices == 0 pred_occluded_new[:, i, left_mask] = pred_occ[0] # 处理右边界 (新时间点大于所有原始时间) right_mask = indices == len(pred_occ) pred_occluded_new[:, i, right_mask] = pred_occ[-1] # 处理中间区域 (新时间点在原始时间范围内) mid_mask = ~(left_mask | right_mask) mid_indices = indices[mid_mask] # 计算新时间点与左右邻居的距离 left_dist = gt_t[mid_mask] - est_t[mid_indices - 1] right_dist = est_t[mid_indices] - gt_t[mid_mask] # 选择更近的邻居 closer_to_left = left_dist < right_dist pred_occluded_new[:, i,mid_mask] = np.where( closer_to_left, pred_occ[mid_indices - 1], pred_occ[mid_indices] ) pred_tracks = pred_trajectory_new.copy() pred_occluded = pred_occluded_new.copy() out_metrics = compute_tapvid_metrics( query_points, gt_occluded, gt_tracks, pred_occluded, pred_tracks, query_mode="first", ) print("metrics", out_metrics) # Visualization if enable_visualization and vis is not None: vis.visualize( sample.video if isinstance(sample.video, torch.Tensor) else torch.from_numpy(sample.video).float(), sample.events if isinstance(sample.events, torch.Tensor) else torch.from_numpy(sample.events).float(), pred_trajectory, pred_visibility > 0.8, filename=seq_name, video_model="rgb", ) # Save trajectory if requested if save_trajectory: B, T, N, _ = pred_trajectory.shape ind = np.arange(N).reshape(1, 1, N, 1).repeat(T, axis=1) t = sample.segmentation.astype(float) t *= 1e-6 t = t.reshape(1, -1, 1, 1).repeat(N, axis=2) pred_trajectory_with_time = np.concatenate((ind, t, pred_trajectory.cpu().numpy()), axis=3) pred_trajectory_txt = np.transpose(pred_trajectory_with_time, (0, 2, 1, 3)).reshape(-1, 4) traj_path = os.path.join(output_dir, "pred_trajectory.txt") np.savetxt(traj_path, pred_trajectory_txt) print(f"Trajectory saved to {traj_path}") # Save results if requested if save_results: result_path = os.path.join(output_dir, "result.txt") with open(result_path, 'w') as f: f.write("metrics " + str(out_metrics) + "\n") f.write(f"time per frame: {elapsed_time}\n") print(f"Results saved to {result_path}")