""" Visualize LED denoising with/without uncertainty — matching qual_uncertainty.png style. For each denoising step: - Past: blue dots+lines - Future GT: red dots+lines - Prediction: green, with darkness proportional to uncertainty (dark green = certain, light green = uncertain) Left column: No uncertainty (all predictions same green) Right column: With uncertainty (per-agent green intensity from σ) Produces one image per sample with rows = denoising stages, columns = [no uncertainty, with uncertainty]. """ import os, sys, random import numpy as np import torch import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import matplotlib.colors as mcolors from matplotlib.lines import Line2D 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 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 OUT_DIR = '/mnt/jaewoo4tb/srtp/NeurIPS_2026_SRT/Viz_uncertainty' # Colors matching the reference figure PAST_COLOR = '#2962FF' # blue GT_COLOR = '#E91E63' # red/pink PRED_BASE = np.array([0.2, 0.6, 0.1]) # base green for predictions _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 = 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_stages(model, graph, model_init, past_traj, traj_mask, betas, alphas, abs_sqrt, oma_sqrt, use_sigma, traj_scale, init_pos_np): """Get prediction at each meaningful stage.""" sample_pred, mean_est, var_est = model_init(past_traj, traj_mask) 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] sigma_input = var_est if use_sigma else None sigma_np = var_est.detach().cpu().numpy() if use_sigma else None stages = [] # Stage 0: Initializer output (before any denoising) loc_abs = loc.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2) stages.append({'trajs': loc_abs[:A], 'sigma': sigma_np, 'label': 'Initializer'}) # Run denoising, capture at key steps cur_y = loc[:, :10] capture_at = {3: 'After 2 steps', 1: 'After 4 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 capture_at: cur_abs = cur_y.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2) stages.append({'trajs': cur_abs[:A], 'sigma': sigma_np, 'label': capture_at[i]}) # Final: run second half too 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.append({'trajs': final_abs[:A], 'sigma': sigma_np, 'label': 'Final'}) return stages def uncertainty_to_green(sigma_vals, a): """Convert per-agent uncertainty to green color intensity. High certainty (low σ) → dark green, Low certainty (high σ) → light green. """ if sigma_vals is None: return (0.2, 0.65, 0.1, 0.8) # default green sigma_std = np.exp(sigma_vals[:A, 0] / 2) # actual std # Normalize to [0, 1] range s_min, s_max = sigma_std.min(), sigma_std.max() if s_max - s_min < 1e-6: norm_val = 0.5 else: norm_val = (sigma_std[a] - s_min) / (s_max - s_min) # Dark green (certain) to light yellow-green (uncertain) # Interpolate: dark (0.1, 0.4, 0.05) ↔ light (0.6, 0.85, 0.3) dark = np.array([0.05, 0.35, 0.0]) light = np.array([0.55, 0.82, 0.25]) color = dark + norm_val * (light - dark) return (*color, 0.85) def draw_stage(ax, past_abs, gt_abs, trajs, sigma_vals, show_sigma, xlim=None, ylim=None): """Draw one stage on court.""" ax.imshow(court_img(), extent=[0, COURT_W, COURT_H, 0], zorder=0, alpha=0.45) if xlim: ax.set_xlim(*xlim) else: ax.set_xlim(0, COURT_W) if ylim: ax.set_ylim(*ylim) else: ax.set_ylim(COURT_H, 0) ax.axis('off') K = trajs.shape[1] K_show = min(5, K) # Find best mode per agent best_modes = [] for a in range(A): if K > 1: dists = np.linalg.norm(trajs[a, :K_show] - gt_abs[a:a+1], axis=-1).mean(axis=-1) best_modes.append(dists.argmin()) else: best_modes.append(0) # Draw predictions (green, intensity by uncertainty) for a in range(A): if show_sigma: color = uncertainty_to_green(sigma_vals, a) else: color = (0.2, 0.65, 0.1, 0.7) # Light modes for k in range(K_show): pred = trajs[a, k] ax.plot(pred[:, 0], pred[:, 1], color=color, lw=0.4, alpha=0.25, zorder=3) # Best mode (thicker) best = trajs[a, best_modes[a]] pred_line = np.concatenate([past_abs[a, -1:], best], axis=0) ax.plot(pred_line[:, 0], pred_line[:, 1], color=color, lw=1.8, marker='o', ms=1.8, markevery=3, zorder=5) # Draw past (blue) for a in range(A): ax.plot(past_abs[a, :, 0], past_abs[a, :, 1], color=PAST_COLOR, lw=1.2, marker='o', ms=2.0, markevery=2, alpha=0.85, zorder=6) # Draw GT future (red) 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=1.0, marker='o', ms=1.5, markevery=3, alpha=0.7, zorder=4) def make_figure(stages_nosigma, stages_sigma, past_abs, gt_abs, sample_idx, zoom_region=None): """Create the full comparison figure.""" n_stages = len(stages_nosigma) # If zoom region provided, add zoomed row at bottom n_rows = n_stages + (1 if zoom_region else 0) fig, axes = plt.subplots(n_rows, 2, figsize=(12, 3.0 * n_rows), dpi=200) if n_rows == 1: axes = axes.reshape(1, 2) for row in range(n_stages): # Left: no uncertainty draw_stage(axes[row, 0], past_abs, gt_abs, stages_nosigma[row]['trajs'], None, show_sigma=False) if row == 0: axes[row, 0].set_title('No uncertainty', fontsize=11, fontweight='bold') # Right: with uncertainty draw_stage(axes[row, 1], past_abs, gt_abs, stages_sigma[row]['trajs'], stages_sigma[row]['sigma'], show_sigma=True) if row == 0: axes[row, 1].set_title('Using uncertainty', fontsize=11, fontweight='bold') # Row label axes[row, 0].text(0.02, 0.95, stages_nosigma[row]['label'], transform=axes[row, 0].transAxes, fontsize=8, verticalalignment='top', fontweight='bold', bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.7)) # Zoomed bottom row if zoom_region and n_rows > n_stages: xl, yl = zoom_region for col in range(2): stages = stages_nosigma if col == 0 else stages_sigma sigma = None if col == 0 else stages[-1]['sigma'] draw_stage(axes[-1, col], past_abs, gt_abs, stages[-1]['trajs'], sigma, show_sigma=(col == 1), xlim=xl, ylim=yl) # Draw zoom box on the row above from matplotlib.patches import Rectangle rect = Rectangle((xl[0], yl[1]), xl[1]-xl[0], yl[0]-yl[1], linewidth=1.5, edgecolor='black', facecolor='none', zorder=10) axes[-2, col].add_patch(rect) # Add noise level indicator on the right side # Gray gradient arrow from τ^K (top) to 0 (bottom) noise_ax = fig.add_axes([0.92, 0.15, 0.03, 0.7]) gradient = np.linspace(0.3, 1.0, 256).reshape(256, 1) noise_ax.imshow(gradient, aspect='auto', cmap='Greys_r', extent=[0, 1, 0, 1]) noise_ax.set_xticks([]) noise_ax.set_yticks([0, 1]) noise_ax.set_yticklabels(['0', 'τ$^K$'], fontsize=8) noise_ax.set_ylabel('Noise level', fontsize=8, rotation=270, labelpad=12) noise_ax.yaxis.set_label_position('right') # Legend legend_elements = [ Line2D([0], [0], color=PAST_COLOR, lw=2, marker='o', ms=4, label='Past'), Line2D([0], [0], color=GT_COLOR, lw=2, marker='o', ms=4, label='Future'), Line2D([0], [0], color=(0.2, 0.65, 0.1), lw=2, marker='o', ms=4, label='Prediction'), ] fig.legend(handles=legend_elements, loc='upper center', ncol=3, fontsize=9, frameon=True, fancybox=True, shadow=True, bbox_to_anchor=(0.45, 0.98)) # Uncertainty colorbar for the right column from matplotlib.cm import ScalarMappable from matplotlib.colors import LinearSegmentedColormap, Normalize dark = (0.05, 0.35, 0.0) light = (0.55, 0.82, 0.25) cmap_unc = LinearSegmentedColormap.from_list('unc', [dark, light]) sm = ScalarMappable(cmap=cmap_unc, norm=Normalize(0, 1)) sm.set_array([]) cbar_ax = fig.add_axes([0.52, 0.96, 0.15, 0.012]) cbar = fig.colorbar(sm, cax=cbar_ax, orientation='horizontal') cbar.set_ticks([0, 1]) cbar.set_ticklabels(['Certain', 'Uncertain'], fontsize=7) cbar_ax.set_title('Uncertainty level', fontsize=7, pad=2) plt.subplots_adjust(hspace=0.05, wspace=0.02, right=0.90, top=0.93) save_path = os.path.join(OUT_DIR, f'denoising_uncertainty_sample_{sample_idx:04d}.png') fig.savefig(save_path, bbox_inches='tight', pad_inches=0.05, dpi=200) plt.close(fig) print(f'Saved: {save_path}') def main(): device = 'cuda:0' os.environ['CUDA_VISIBLE_DEVICES'] = '3' # use GPU 3 torch.cuda.set_device(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' 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)), 5)) 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:] # Data is already in court units (0-28 x 0-15), no further scaling needed past_abs_np = data['pre_motion_3D'].numpy().squeeze(0) fut_abs_np = data['fut_motion_3D'].numpy().squeeze(0) init_pos_np = initial_pos.cpu().numpy().squeeze(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() # Get stages for both versions stages_nosigma = get_stages( model_n, graph_n, init_n, past_traj, traj_mask, betas, alphas, abs_sqrt, oma_sqrt, False, traj_scale, init_pos_np) stages_sigma = get_stages( model_s, graph_s, init_s, past_traj, traj_mask, betas, alphas, abs_sqrt, oma_sqrt, True, traj_scale, init_pos_np) # Compute zoom region around the action all_pos = np.concatenate([past_abs_np.reshape(-1, 2), fut_abs_np.reshape(-1, 2)]) cx, cy = all_pos.mean(axis=0) span = max(all_pos.max(axis=0) - all_pos.min(axis=0)) * 0.6 zoom = ([cx - span, cx + span], [cy + span, cy - span]) make_figure(stages_nosigma, stages_sigma, past_abs_np, fut_abs_np, sample_idx, zoom_region=zoom) print(f'\nAll saved to {OUT_DIR}/') if __name__ == '__main__': main()