LED/viz_uncertainty_denoising.py

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