LED/eval_sdd_led_mid_protocol.py

"""
Re-evaluate an LED SDD checkpoint using the MID evaluation protocol:
  per-pedestrian full-horizon (12-frame) ADE / final-frame FDE,
  best_of_20 per pedestrian, averaged across all target pedestrians,
  ×50 to report in pixels.

Usage:
  python eval_sdd_led_mid_protocol.py --exp baseline_v2 --epoch 60
  python eval_sdd_led_mid_protocol.py --exp graph_sigma_v2 --epoch 60 --use_graph
"""
import argparse, os, sys, random, torch, numpy as np
from torch.utils.data import DataLoader
from data.dataloader_sdd import SDDDataset, sdd_seq_collate
from models.model_led_initializer import LEDInitializer as InitializationModel
from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
from trainer.train_sdd_led import NUM_Tau
from utils.config import Config


def build_models(cfg, use_graph, use_v6_graph, ckpt_path, device):
    model = CoreDenoisingModel(past_len=cfg.past_frames).to(device)
    core_ckpt = torch.load(cfg.pretrained_core_denoising_model, map_location='cpu')
    model.load_state_dict(core_ckpt['model_dict'])
    model.eval()

    init = InitializationModel(
        t_h=cfg.past_frames, d_h=6,
        t_f=cfg.future_frames, d_f=2, k_pred=20).to(device)

    ckpt = torch.load(ckpt_path, map_location='cpu')
    init.load_state_dict(ckpt['model_initializer_dict'])
    init.eval()

    graph = None
    if use_graph:
        from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper
        graph = FutureInteractionGraphV6Wrapper(
            num_agents=64, future_steps=cfg.future_frames,
            past_steps=cfg.past_frames, past_channels=6,
            node_dim=128, top_n=5, num_denoise_steps=NUM_Tau).to(device)
        sd = {k: v for k, v in ckpt['interaction_graph_dict'].items()
              if '_single_edge_index' not in k}
        graph.load_state_dict(sd, strict=False)
        graph.eval()

    return model, init, graph


def make_beta_schedule(n=100, start=1e-4, end=5e-2):
    return torch.linspace(start, end, n)


def extract(a, t, x):
    out = torch.gather(a, 0, t.to(a.device))
    return out.reshape(t.shape[0], *([1] * (len(x.shape) - 1)))


@torch.no_grad()
def p_sample_accelerate(x, mask, cur_y, t, model, graph, use_v6_graph, sigma,
                        betas, alphas, alphas_bar_sqrt, one_minus_alphas_bar_sqrt):
    t_tensor = torch.tensor([int(t)]).to(x.device)
    eps_factor = ((1 - extract(alphas, t_tensor, cur_y))
                  / extract(one_minus_alphas_bar_sqrt, t_tensor, cur_y))
    beta = extract(betas, t_tensor.repeat(x.shape[0]), cur_y)
    eps_theta = model.generate_accelerate(cur_y, beta, x, mask)

    if graph is not None:
        abs_t = extract(alphas_bar_sqrt, t_tensor, cur_y)
        am1_t = extract(one_minus_alphas_bar_sqrt, t_tensor, cur_y)
        y0_hat = (cur_y - am1_t * eps_theta) / abs_t
        delta = graph(y0_hat, x, int(t), sigma=sigma, A_override=x.size(0))
        eps_theta = eps_theta - (abs_t / am1_t) * delta

    mean = (1 / extract(alphas, t_tensor, cur_y).sqrt()) \
           * (cur_y - eps_factor * eps_theta)
    z = torch.randn_like(cur_y)
    sigma_t = extract(betas, t_tensor, cur_y).sqrt()
    return mean + sigma_t * z * 0.00001


@torch.no_grad()
def run(args):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    cfg = Config(args.cfg, args.exp)

    test_dset = SDDDataset(obs_len=cfg.past_frames,
                           pred_len=cfg.future_frames, split='test')
    loader = DataLoader(test_dset, batch_size=1, shuffle=False,
                        num_workers=2, collate_fn=sdd_seq_collate)

    exp_dir = cfg.model_dir  # results_sdd/sdd/sdd/<info>/models
    ckpt_path = cfg.model_path % args.epoch
    print(f'Loading checkpoint: {ckpt_path}')

    model, init, graph = build_models(
        cfg, use_graph=args.use_graph, use_v6_graph=args.use_v6_graph,
        ckpt_path=ckpt_path, device=device)

    betas = make_beta_schedule().to(device)
    alphas = 1 - betas
    alphas_prod = torch.cumprod(alphas, 0)
    abs_sqrt = torch.sqrt(alphas_prod)
    one_minus_abs_sqrt = torch.sqrt(1 - alphas_prod)

    traj_mean = torch.FloatTensor(cfg.traj_mean).to(device).view(1, 1, 1, 2)
    traj_scale = float(cfg.traj_scale)

    np.random.seed(0); random.seed(0)
    torch.manual_seed(0); torch.cuda.manual_seed_all(0)

    total_ade, total_fde, n = 0.0, 0.0, 0
    T = cfg.future_frames
    for data in loader:
        pre = data['pre_motion_3D'].to(device)       # [1, A, 8, 2]
        fut = data['fut_motion_3D'].to(device)
        A = pre.size(1)
        initial_pos = pre[:, :, -1:]
        past_abs = ((pre - traj_mean) / traj_scale).contiguous().view(-1, cfg.past_frames, 2)
        past_rel = ((pre - initial_pos) / traj_scale).contiguous().view(-1, cfg.past_frames, 2)
        past_vel = torch.cat([past_rel[:, 1:] - past_rel[:, :-1],
                              torch.zeros_like(past_rel[:, -1:])], dim=1)
        past = torch.cat([past_abs, past_rel, past_vel], dim=-1)
        fut_rel = ((fut - initial_pos) / traj_scale).contiguous().view(-1, T, 2)
        mask = torch.ones(A, A).to(device)

        sp, me, ve = init(past, mask)
        ve = ve.clamp(min=-5, max=5)
        sp = torch.exp(ve / 2)[..., None, None] * sp \
             / (sp.std(dim=1).mean(dim=(1, 2))[:, None, None, None] + 1e-6)
        loc = sp + me[:, None]
        sigma_in = ve if args.use_v6_graph else None

        # leapfrog: 20 modes = 10+10 two halves, each 5 reverse steps
        cur_y = loc[:, :10]
        for i in reversed(range(NUM_Tau)):
            cur_y = p_sample_accelerate(
                past, mask, cur_y, i, model, graph, args.use_v6_graph, sigma_in,
                betas, alphas, abs_sqrt, one_minus_abs_sqrt)
        cur_y_ = loc[:, 10:]
        for i in reversed(range(NUM_Tau)):
            cur_y_ = p_sample_accelerate(
                past, mask, cur_y_, i, model, graph, args.use_v6_graph, sigma_in,
                betas, alphas, abs_sqrt, one_minus_abs_sqrt)
        pred = torch.cat((cur_y_, cur_y), dim=1)   # [A, 20, T, 2]

        # target only (index 0) per scene, full-horizon ADE / final FDE, best_of_20
        pred_0 = pred[0:1]
        fut_0 = fut_rel[0:1]
        dist = torch.norm(fut_0.unsqueeze(1) - pred_0, dim=-1) * traj_scale
        ade = dist.mean(dim=-1).min(dim=-1)[0]
        fde = dist[:, :, -1].min(dim=-1)[0]
        total_ade += ade.sum().item()
        total_fde += fde.sum().item()
        n += 1

    ade_px = total_ade / n * 50.0
    fde_px = total_fde / n * 50.0
    print(f'Epoch {args.epoch} MID-protocol: ADE={ade_px:.4f}  FDE={fde_px:.4f}  (n={n})')


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--cfg', default='sdd/sdd')
    p.add_argument('--exp', required=True, help='info tag e.g. baseline_v2')
    p.add_argument('--epoch', type=int, required=True)
    p.add_argument('--use_graph', action='store_true')
    p.add_argument('--use_v6_graph', action='store_true')
    args = p.parse_args()
    run(args)