""" Visualize LED denoising: initializer → refinement → final, on basketball court. LED's denoising is a 2-stage process: 1. Initializer: produces 20 diverse trajectory modes (noisy, multimodal) 2. Leapfrog refinement: 5 DDPM steps that clean up each mode (subtle changes) We visualize 5 stages: Step 0: Raw initializer output (mean_estimation only, no variance) Step 1: Initializer + variance scaling (diverse modes) Step 2: After τ=4,3 refinement (2 DDPM steps) Step 3: After τ=2,1 refinement (4 DDPM steps) Step 4: Final prediction τ=0 (5 DDPM steps, fully denoised) For sigma version: trajectory color = uncertainty (green=certain, red=uncertain) For nosigma version: trajectory color = team (blue=home, orange=away, green=ball) """ import os, sys, random import numpy as np import torch import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import matplotlib.cm as cm from matplotlib.colors import Normalize sys.path.insert(0, os.path.dirname(__file__)) from utils.config import Config from data.dataloader_nba import NBADataset, seq_collate from torch.utils.data import DataLoader from models.model_led_initializer import LEDInitializer as InitializationModel from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper NUM_Tau = 5 COURT_IMG = '/mnt/jaewoo4tb/srtp/srtp/raw_data/nba/court.png' COURT_W, COURT_H = 28.0, 15.0 A = 11 HOME_PAST, HOME_FUT = '#a8d5ff', '#187bff' AWAY_PAST, AWAY_FUT = '#ffc9a8', '#ff5b2e' BALL_PAST, BALL_FUT = '#b9f2b9', '#1f9d1f' GT_COLOR = '#333333' def agent_colors(idx): if idx < 5: return HOME_PAST, HOME_FUT elif idx < 10: return AWAY_PAST, AWAY_FUT else: return BALL_PAST, BALL_FUT _court_cache = None def court_img(): global _court_cache if _court_cache is None: _court_cache = plt.imread(COURT_IMG) return _court_cache def load_models(ckpt_path, use_sigma, edge_mode, top_n, device): cfg = Config('led_augment', 'viz') model = CoreDenoisingModel().to(device) cp = torch.load(cfg.pretrained_core_denoising_model, map_location='cpu', weights_only=False) model.load_state_dict(cp['model_dict']); model.eval() model_init = InitializationModel(t_h=10, d_h=6, t_f=20, d_f=2, k_pred=20).to(device) graph = FutureInteractionGraphV6Wrapper( num_agents=11, future_steps=20, past_steps=10, past_channels=6, node_dim=128, top_n=top_n, num_denoise_steps=NUM_Tau, edge_mode=edge_mode).to(device) ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False) model_init.load_state_dict(ckpt['model_initializer_dict']) graph.load_state_dict(ckpt['interaction_graph_dict']) model_init.eval(); graph.eval() return cfg, model, model_init, graph def get_all_stages(model, graph, model_init, past_traj, traj_mask, betas, alphas, abs_sqrt, oma_sqrt, use_sigma, traj_scale, init_pos_np): """Return trajectory predictions at each meaningful stage.""" sample_pred, mean_est, var_est = model_init(past_traj, traj_mask) stages = [] # Stage 0: Mean estimation only (single mode per agent) mean_abs = mean_est.detach().cpu().numpy() # [B*A, T, 2] mean_abs = mean_abs * traj_scale + init_pos_np.reshape(-1, 1, 2) stages.append({ 'trajs': mean_abs[:A, np.newaxis], # [A, 1, T, 2] 'sigma': None, 'title': 'Stage 1: Mean estimation', }) # Stage 1: Initializer with variance scaling (20 diverse modes) sample_pred_scaled = (torch.exp(var_est / 2)[..., None, None] * sample_pred / sample_pred.std(dim=1).mean(dim=(1, 2))[:, None, None, None]) loc = sample_pred_scaled + mean_est[:, None] loc_abs = loc.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2) sigma_np = var_est.detach().cpu().numpy() if use_sigma else None stages.append({ 'trajs': loc_abs[:A], # [A, K=20, T, 2] 'sigma': sigma_np, 'title': 'Stage 2: Initializer (20 modes)', }) # Run denoising and capture intermediate states sigma_input = var_est if use_sigma else None cur_y = loc[:, :10] checkpoints = {3: 'Stage 3: After 2 DDPM steps', 1: 'Stage 4: After 4 DDPM steps', -1: 'Stage 5: Final (5 DDPM steps)'} for i in reversed(range(NUM_Tau)): ef = (1 - alphas[i]) / oma_sqrt[i] beta = betas[i].repeat(past_traj.shape[0]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) eps = model.generate_accelerate(cur_y, beta.squeeze(-1).squeeze(-1), past_traj, traj_mask) y0h = (cur_y - oma_sqrt[i] * eps) / abs_sqrt[i] delta = graph(y0h, past_traj, i, sigma=sigma_input) eps = eps + delta mean = (1 / alphas[i].sqrt()) * (cur_y - ef * eps) z = torch.randn_like(cur_y) cur_y = mean + betas[i].sqrt() * z * 0.00001 if i in checkpoints: # Get y0 estimate (clean prediction) at this point y0_est = (cur_y - oma_sqrt[max(0, i-1)] * eps) / abs_sqrt[max(0, i-1)] if i > 0 else cur_y y0_abs = y0_est.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2) stages.append({ 'trajs': y0_abs[:A], # [A, K=10, T, 2] 'sigma': sigma_np, 'title': checkpoints[i], }) # Stage 5 (final): use cur_y directly # Also do the second half (modes 10-19) cur_y2 = loc[:, 10:] for i in reversed(range(NUM_Tau)): ef = (1 - alphas[i]) / oma_sqrt[i] beta = betas[i].repeat(past_traj.shape[0]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) eps = model.generate_accelerate(cur_y2, beta.squeeze(-1).squeeze(-1), past_traj, traj_mask) y0h = (cur_y2 - oma_sqrt[i] * eps) / abs_sqrt[i] delta = graph(y0h, past_traj, i, sigma=sigma_input) eps = eps + delta mean = (1 / alphas[i].sqrt()) * (cur_y2 - ef * eps) z = torch.randn_like(cur_y2) cur_y2 = mean + betas[i].sqrt() * z * 0.00001 final = torch.cat((cur_y2, cur_y), dim=1) final_abs = final.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2) stages[-1] = { 'trajs': final_abs[:A], # [A, K=20, T, 2] 'sigma': sigma_np, 'title': 'Stage 5: Final (all 20 modes)', } return stages def draw_stage(ax, past_abs, gt_abs, trajs, sigma_vals, title, show_sigma): """Draw one stage on court.""" ax.imshow(court_img(), extent=[0, COURT_W, COURT_H, 0], zorder=0, alpha=0.5) ax.set_xlim(0, COURT_W); ax.set_ylim(COURT_H, 0) ax.axis('off') ax.set_title(title, fontsize=9, pad=3) # Past for a in range(A): pc, _ = agent_colors(a) ax.plot(past_abs[a, :, 0], past_abs[a, :, 1], color=pc, lw=0.8, marker='o', ms=1.5, alpha=0.6, zorder=2) # GT for a in range(A): gt = np.concatenate([past_abs[a, -1:], gt_abs[a]], axis=0) ax.plot(gt[:, 0], gt[:, 1], color=GT_COLOR, lw=0.7, marker='o', ms=1.0, alpha=0.4, zorder=3, linestyle='--') K = trajs.shape[1] K_show = min(K, 10) if show_sigma and sigma_vals is not None: sigma_std = np.exp(sigma_vals[:A, 0] / 2) norm = Normalize(vmin=sigma_std.min() - 0.01, vmax=sigma_std.max() + 0.01) cmap = cm.RdYlGn_r for a in range(A): color = cmap(norm(sigma_std[a])) for k in range(K_show): pred = np.concatenate([past_abs[a, -1:], trajs[a, k]], axis=0) ax.plot(pred[:, 0], pred[:, 1], color=color, lw=0.5, alpha=0.3, zorder=4) # Best mode dists = np.linalg.norm(trajs[a, :K_show] - gt_abs[a:a+1], axis=-1).mean(axis=-1) best_k = dists.argmin() best = np.concatenate([past_abs[a, -1:], trajs[a, best_k]], axis=0) ax.plot(best[:, 0], best[:, 1], color=color, lw=2.0, alpha=0.9, zorder=5, marker='o', ms=2.0) else: for a in range(A): _, fc = agent_colors(a) for k in range(K_show): pred = np.concatenate([past_abs[a, -1:], trajs[a, k]], axis=0) ax.plot(pred[:, 0], pred[:, 1], color=fc, lw=0.5, alpha=0.25, zorder=4) dists = np.linalg.norm(trajs[a, :K_show] - gt_abs[a:a+1], axis=-1).mean(axis=-1) best_k = dists.argmin() best = np.concatenate([past_abs[a, -1:], trajs[a, best_k]], axis=0) ax.plot(best[:, 0], best[:, 1], color=fc, lw=2.0, alpha=0.9, zorder=5, marker='o', ms=2.0) def main(): device = 'cuda:0' ckpt_sigma = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_edge_relpos/models/model_0052.p' ckpt_nosigma = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_nosigma_n3/models/model_0084.p' out_dir = '/mnt/jaewoo4tb/srtp/LED/visualizations/denoising_steps' os.makedirs(out_dir, exist_ok=True) cfg, model_s, init_s, graph_s = load_models( ckpt_sigma, True, 'relpos_only', 5, device) _, model_n, init_n, graph_n = load_models( ckpt_nosigma, False, 'full', 3, device) betas = torch.linspace(1e-5, 1e-2, 100).to(device) alphas = 1 - betas alphas_prod = torch.cumprod(alphas, 0) abs_sqrt = torch.sqrt(alphas_prod) oma_sqrt = torch.sqrt(1 - alphas_prod) traj_scale = cfg.traj_scale test_dset = NBADataset(obs_len=10, pred_len=20, training=False) test_loader = DataLoader(test_dset, batch_size=1, shuffle=False, collate_fn=seq_collate) np.random.seed(42); random.seed(42); torch.manual_seed(42) sample_indices = sorted(random.sample(range(len(test_dset)), 3)) with torch.no_grad(): for sample_idx, data in enumerate(test_loader): if sample_idx not in sample_indices: continue if sample_idx > max(sample_indices): break traj_mean_t = torch.FloatTensor(cfg.traj_mean).cuda().unsqueeze(0).unsqueeze(0).unsqueeze(0) initial_pos = data['pre_motion_3D'].cuda()[:, :, -1:] past_abs_np = data['pre_motion_3D'].numpy().squeeze(0) / (94.0 / 28.0) fut_abs_np = data['fut_motion_3D'].numpy().squeeze(0) / (94.0 / 28.0) init_pos_np = initial_pos.cpu().numpy().squeeze(0) / (94.0 / 28.0) past_traj_abs = ((data['pre_motion_3D'].cuda() - traj_mean_t) / traj_scale).view(-1, 10, 2) past_traj_rel = ((data['pre_motion_3D'].cuda() - initial_pos) / traj_scale).view(-1, 10, 2) past_traj_vel = torch.cat((past_traj_rel[:, 1:] - past_traj_rel[:, :-1], torch.zeros_like(past_traj_rel[:, :1])), dim=1) past_traj = torch.cat((past_traj_abs, past_traj_rel, past_traj_vel), dim=-1) traj_mask = torch.ones(11, 11).cuda() for version, model, init_model, graph, use_sigma, label in [ ('nosigma', model_n, init_n, graph_n, False, 'Without Uncertainty'), ('sigma', model_s, init_s, graph_s, True, 'With Uncertainty'), ]: stages = get_all_stages( model, graph, init_model, past_traj, traj_mask, betas, alphas, abs_sqrt, oma_sqrt, use_sigma, traj_scale, init_pos_np) for si, stage in enumerate(stages): fig, ax = plt.subplots(1, 1, figsize=(7, 5), dpi=200) title = f'{label} — {stage["title"]}' draw_stage(ax, past_abs_np, fut_abs_np, stage['trajs'], stage['sigma'], title, show_sigma=use_sigma) if use_sigma and stage['sigma'] is not None: sigma_std = np.exp(stage['sigma'][:A, 0] / 2) sm = cm.ScalarMappable(cmap=cm.RdYlGn_r, norm=Normalize(vmin=sigma_std.min() - 0.01, vmax=sigma_std.max() + 0.01)) sm.set_array([]) cbar = fig.colorbar(sm, ax=ax, shrink=0.5, pad=0.02) cbar.set_label('σ (uncertainty)', fontsize=7) plt.tight_layout() fname = f'sample_{sample_idx:04d}_{version}_stage{si}.png' fig.savefig(os.path.join(out_dir, fname), bbox_inches='tight', pad_inches=0.02, dpi=200) plt.close(fig) print(f' Sample {sample_idx} {version}: {len(stages)} stages saved') print(f'\nAll saved to {out_dir}/') if __name__ == '__main__': main()