LED/viz_denoising_steps.py

"""
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()