""" 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 = [] # Initializer 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'}) # Denoise first 10 modes 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:] # Data is already in court units (0-28 x 0-15), no further scaling needed 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()