| """ |
| Draw each denoising step as a separate image. |
| Past=blue, GT=red, Prediction=green (darkness by uncertainty). |
| """ |
|
|
| import os, sys, random |
| import numpy as np |
| import torch |
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
|
|
| sys.path.insert(0, '/mnt/jaewoo4tb/srtp/LED') |
|
|
| 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 |
| from models.model_diffusion import TransformerDenoisingModel |
| 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 |
| OUT_DIR = '/mnt/jaewoo4tb/srtp/NeurIPS_2026_SRT/Viz_uncertainty' |
|
|
| PAST_COLOR = '#2962FF' |
| GT_COLOR = '#E91E63' |
|
|
| _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): |
| os.chdir('/mnt/jaewoo4tb/srtp/LED') |
| cfg = Config('led_augment', 'viz') |
| model = TransformerDenoisingModel().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() |
| init_m = LEDInitializer(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) |
| init_m.load_state_dict(ckpt['model_initializer_dict']) |
| graph.load_state_dict(ckpt['interaction_graph_dict']) |
| init_m.eval(); graph.eval() |
| return cfg, model, init_m, graph |
|
|
|
|
| def get_all_steps(model, graph, init_m, past_traj, traj_mask, |
| betas, alphas, abs_sqrt, oma_sqrt, |
| use_sigma, traj_scale, init_pos_np): |
| sample_pred, mean_est, var_est = init_m(past_traj, traj_mask) |
| sample_pred = (torch.exp(var_est / 2)[..., None, None] |
| * sample_pred / sample_pred.std(dim=1).mean(dim=(1, 2))[:, None, None, None]) |
| loc = sample_pred + mean_est[:, None] |
| sigma_input = var_est if use_sigma else None |
| sigma_np = var_est.detach().cpu().numpy() if use_sigma else None |
|
|
| steps = [] |
|
|
| |
| loc_abs = loc.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2) |
| steps.append({'trajs': loc_abs[:A], 'sigma': sigma_np, 'label': 'Initializer'}) |
|
|
| |
| cur_y = 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_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 |
|
|
| cur_abs = cur_y.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2) |
| steps.append({'trajs': cur_abs[:A], 'sigma': sigma_np, 'label': f'Step {NUM_Tau - i}/{NUM_Tau}'}) |
|
|
| return steps |
|
|
|
|
| def unc_color(sigma_vals, a): |
| if sigma_vals is None: |
| return (0.2, 0.65, 0.1, 0.8) |
| sigma_std = np.exp(sigma_vals[:A, 0] / 2) |
| s_min, s_max = sigma_std.min(), sigma_std.max() |
| norm = (sigma_std[a] - s_min) / (s_max - s_min + 1e-8) |
| dark = np.array([0.05, 0.35, 0.0]) |
| light = np.array([0.55, 0.82, 0.25]) |
| c = dark + norm * (light - dark) |
| return (*c, 0.85) |
|
|
|
|
| def draw(ax, past, gt, trajs, sigma, show_sigma): |
| ax.imshow(court_img(), extent=[0, COURT_W, COURT_H, 0], zorder=0, alpha=0.45) |
| ax.set_xlim(0, COURT_W); ax.set_ylim(COURT_H, 0) |
| ax.axis('off') |
|
|
| K = min(5, trajs.shape[1]) |
|
|
| for a in range(A): |
| color = unc_color(sigma, a) if show_sigma else (0.2, 0.65, 0.1, 0.8) |
| dists = np.linalg.norm(trajs[a, :K] - gt[a:a+1], axis=-1).mean(axis=-1) |
| best_k = dists.argmin() |
|
|
| best = np.concatenate([past[a, -1:], trajs[a, best_k]], axis=0) |
| ax.plot(best[:, 0], best[:, 1], color=color, lw=1.8, |
| marker='o', ms=1.8, markevery=3, zorder=5) |
|
|
| for a in range(A): |
| ax.plot(past[a, :, 0], past[a, :, 1], color=PAST_COLOR, |
| lw=1.2, marker='o', ms=2.0, markevery=2, alpha=0.85, zorder=6) |
|
|
| for a in range(A): |
| g = np.concatenate([past[a, -1:], gt[a]], axis=0) |
| ax.plot(g[:, 0], g[:, 1], color=GT_COLOR, lw=1.0, |
| marker='o', ms=1.5, markevery=3, alpha=0.7, zorder=4) |
|
|
|
|
| def main(): |
| device = 'cuda:0' |
| os.environ['CUDA_VISIBLE_DEVICES'] = '3' |
| torch.cuda.set_device(0) |
|
|
| ckpt_s = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_edge_relpos/models/model_0052.p' |
| ckpt_n = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_nosigma_n3/models/model_0084.p' |
|
|
| cfg, model_s, init_s, graph_s = load_models(ckpt_s, True, 'relpos_only', 5, device) |
| _, model_n, init_n, graph_n = load_models(ckpt_n, 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 idx, data in enumerate(test_loader): |
| if idx not in sample_indices: |
| continue |
| if 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_np = data['pre_motion_3D'].numpy().squeeze(0) |
| fut_np = data['fut_motion_3D'].numpy().squeeze(0) |
| init_np = initial_pos.cpu().numpy().squeeze(0) |
|
|
| pa = ((data['pre_motion_3D'].cuda() - traj_mean_t) / traj_scale).view(-1, 10, 2) |
| pr = ((data['pre_motion_3D'].cuda() - initial_pos) / traj_scale).view(-1, 10, 2) |
| pv = torch.cat((pr[:, 1:] - pr[:, :-1], torch.zeros_like(pr[:, :1])), dim=1) |
| past_traj = torch.cat((pa, pr, pv), dim=-1) |
| mask = torch.ones(11, 11).cuda() |
|
|
| for version, model, init_m, graph, use_sigma, tag in [ |
| ('nosigma', model_n, init_n, graph_n, False, 'no_uncertainty'), |
| ('sigma', model_s, init_s, graph_s, True, 'with_uncertainty'), |
| ]: |
| steps = get_all_steps(model, graph, init_m, past_traj, mask, |
| betas, alphas, abs_sqrt, oma_sqrt, |
| use_sigma, traj_scale, init_np) |
|
|
| for si, step in enumerate(steps): |
| fig, ax = plt.subplots(1, 1, figsize=(7, 5), dpi=200) |
| draw(ax, past_np, fut_np, step['trajs'], step['sigma'], use_sigma) |
| plt.tight_layout() |
| fname = f'sample_{idx:04d}_{tag}_step{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 {idx} {tag}: {len(steps)} images') |
|
|
| print(f'All saved to {OUT_DIR}/') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
