test_npsn.py

import os
import glob
import pickle
import argparse
import numpy as np
import torch
from tqdm import tqdm
from torch.utils.data.dataloader import DataLoader
from baselines.converter import get_identity
from npsn import *

parser = argparse.ArgumentParser()
parser.add_argument('--baseline', default='sgcn', help='baseline network ["sgcn","stgcnn"]')
parser.add_argument('--method', default='npsn', help='sampling method ["mc","qmc","npsn"]')
parser.add_argument('--tag', default='npsn', help='personal tag for the model')
parser.add_argument('--gpu_num', default='0', type=str)

test_args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = test_args.gpu_num

if test_args.baseline == 'stgcnn':
    from baselines.stgcnn import *
elif test_args.baseline == 'sgcn':
    from baselines.sgcn import *


@torch.no_grad()
def test(model, model_npsn, loader_test, method='npsn', samples=20, trials=1):
    model.eval()
    model_npsn.eval()
    ade_all, fde_all, tcc_all = [], [], []

    if method == 'qmc':
        sobol_generator = torch.quasirandom.SobolEngine(dimension=2, scramble=True, seed=0)

    for batch in tqdm(loader_test, desc=loader_test.dataset.data_dir):
        obs_traj, pred_traj_gt = [tensor.cuda() for tensor in batch[:2]]

        if test_args.baseline == 'stgcnn':
            V_obs, A_obs, V_tr, A_tr = [tensor.cuda() for tensor in batch[-4:]]
            V_obs_tmp = V_obs.permute(0, 3, 1, 2)
            V_pred, _ = model(V_obs_tmp, A_obs.squeeze())
            V_pred = V_pred.permute(0, 2, 3, 1).detach()
        elif test_args.baseline == 'sgcn':
            V_obs, V_tr = [tensor.cuda() for tensor in batch[-2:]]
            identity = get_identity(V_obs.shape)
            V_pred = model(V_obs, identity)
            V_obs = V_obs[..., 1:]

        mu, cov = generate_statistics_matrices(V_pred.squeeze(dim=0))

        if method == 'npsn':
            loc = model_npsn(V_obs.permute(0, 2, 3, 1))

        V_obs_traj = obs_traj.permute(0, 3, 1, 2).squeeze(dim=0)
        V_pred_traj_gt = pred_traj_gt.permute(0, 3, 1, 2).squeeze(dim=0)

        ade_stack, fde_stack, tcc_stack = [], [], []

        for trial in range(trials):
            if method == 'mc':
                V_pred_sample = mc_sample(mu, cov, samples)
            elif method == 'qmc':
                V_pred_sample = qmc_sample(mu, cov, samples, sobol_generator)
            elif method == 'npsn':
                V_pred_sample = purposive_sample(mu, cov, samples, loc)
            else:
                raise NotImplementedError

            # Evaluate trajectories
            V_absl = V_pred_sample.cumsum(dim=1) + V_obs_traj[[-1], :, :]
            ADEs, FDEs, TCCs = compute_batch_metric(V_absl, V_pred_traj_gt)

            ade_stack.append(ADEs.detach().cpu().numpy())
            fde_stack.append(FDEs.detach().cpu().numpy())
            tcc_stack.append(TCCs.detach().cpu().numpy())

        ade_all.append(np.array(ade_stack))
        fde_all.append(np.array(fde_stack))
        tcc_all.append(np.array(tcc_stack))

    ade_all = np.concatenate(ade_all, axis=1)
    fde_all = np.concatenate(fde_all, axis=1)
    tcc_all = np.concatenate(tcc_all, axis=1)

    mean_ade, mean_fde, mean_tcc = ade_all.mean(axis=0).mean(), fde_all.mean(axis=0).mean(), tcc_all.mean(axis=0).mean()
    return mean_ade, mean_fde, mean_tcc


def main():
    ADE_ls, FDE_ls, TCC_ls = [], [], []
    print("*" * 50)
    root_ = './checkpoints/' + test_args.tag + '-' + test_args.baseline + '/'
    dataset = ['eth', 'hotel', 'univ', 'zara1', 'zara2']

    paths = list(map(lambda x: root_ + x, dataset))

    for feta in range(len(paths)):
        path = paths[feta]
        exps = glob.glob(path)
        print('Model being tested are:', exps)
        for exp_path in exps:
            print("*" * 50)
            print("Evaluating model:", exp_path)

            args_path = exp_path + '/args.pkl'
            with open(args_path, 'rb') as f:
                args = pickle.load(f)

            data_set = './dataset/' + args.dataset + '/'
            model_path = './pretrained/' + test_args.baseline + '/' + args.dataset + '/val_best.pth'
            model_npsn_path = exp_path + '/val_best.pth'

            dset_test = TrajectoryDataset(data_set + 'test/', obs_len=args.obs_len, pred_len=args.pred_len, skip=1)
            loader_test = DataLoader(dset_test, batch_size=1, shuffle=False, num_workers=0)

            if test_args.baseline == 'stgcnn':
                model = STGCNN(n_stgcnn=1, n_txpcnn=5, output_feat=5, kernel_size=3, seq_len=8, pred_seq_len=12).cuda()
            elif test_args.baseline == 'sgcn':
                model = SGCN(number_asymmetric_conv_layer=7, embedding_dims=64, number_gcn_layers=1, dropout=0,
                             obs_len=8, pred_len=12, n_tcn=5, out_dims=5).cuda()
            else:
                raise NotImplementedError

            model.load_state_dict(torch.load(model_path))
            model_npsn = NPSN(t_obs=args.obs_len, s=2, n=args.num_samples).cuda()
            model_npsn.load_state_dict(torch.load(model_npsn_path))

            ADE, FDE, TCC = test(model, model_npsn, loader_test, test_args.method.lower(), args.num_samples)
            ADE_ls.append(ADE), FDE_ls.append(FDE), TCC_ls.append(TCC)
            print("Method: {} N: {} ADE: {:.8f} FDE: {:.8f} TCC: {:.8f}".format(test_args.method, args.num_samples,
                                                                                ADE, FDE, TCC))

    print("*" * 50)
    print("AVG ADE: {:.8f} AVG FDE: {:.8f} AVG TCC: {:.8f}".format(sum(ADE_ls) / 5, sum(FDE_ls) / 5, sum(TCC_ls) / 5))
    print("*" * 50)


if __name__ == '__main__':
    main()