test.py

import argparse
import logging
import os
import time

import gin
import numpy as np
import pytorch_lightning as pl
import torch
from rich.console import Console
from rich.progress import track
from rich.table import Table

import src.data
import src.feature
import src.models
from src.dhvr import DHVR
from src.utils.file import ensure_dir
from src.utils.logger import setup_logger
from src.utils.misc import count_parameters


def print_table(subset_names, stats, rte_ths, rre_ths):
    console = Console()
    table = Table(show_header=True, header_style="bold")

    columns = ["scene", "recall", "rte", "rre", "time"]
    for col in columns:
        table.add_column(col)

    if stats.ndim == 3:
        stats = stats[-1, :, :]

    stats[:, 0] = (stats[:, 1] < rte_ths) * (stats[:, 2] < rre_ths)
    scene_vals = np.zeros((len(subset_names), 4))
    for sid, _ in enumerate(subset_names):
        curr_scene = stats[:, -1] == sid
        if curr_scene.sum() > 0:
            curr_scene_stats = stats[curr_scene]
            success = curr_scene_stats[:, 0] > 0
            recall = success.mean()
            scene_vals[sid][0] = recall
            scene_vals[sid][1:4] = curr_scene_stats[success, 1:4].mean(0)
        else:
            scene_vals[sid] = None

    for sid, vals in zip(subset_names, scene_vals):
        table.add_row(sid, *[f"{v:.4f}" for v in vals])

    success = stats[:, 0] > 0
    recall = success.mean()
    metrics = stats[success, :4].mean(0)
    metrics[0] = recall
    table.add_row("avg", *[f"{m:.4f}" for m in metrics])
    console.print(table)


def rte_rre(T_pred, T_gt, rte_thresh, rre_thresh, eps=1e-16):
    if T_pred is None:
        return np.array([0, np.inf, np.inf])

    rte = np.linalg.norm(T_pred[:3, 3] - T_gt[:3, 3]) * 100
    rre = (
        np.arccos(
            np.clip(
                (np.trace(T_pred[:3, :3].T @ T_gt[:3, :3]) - 1) / 2, -1 + eps, 1 - eps
            )
        )
        * 180
        / np.pi
    )
    return np.array([rte < rte_thresh and rre < rre_thresh, rte, rre])


def run_benchmark(
    data_loader,
    method,
    TE_THRESH,
    RE_THRESH,
    log_interval=100,
):
    tot_num_data = len(data_loader)
    data_loader_iter = iter(data_loader)

    dataset = data_loader.dataset
    subset_names = dataset.subset_names

    stats = np.zeros((tot_num_data, 5))
    stats[:, -1] = -1
    poses = []

    with torch.no_grad():
        for batch_idx in track(range(tot_num_data)):
            batch = data_loader_iter.next()
            sname, xyz0, xyz1, trans = batch[0]
            sid = subset_names.index(sname)
            T_gt = np.linalg.inv(trans)

            start = time.time()
            T = method.register(xyz0, xyz1)
            end = time.time()

            result = rte_rre(T, T_gt, TE_THRESH, RE_THRESH)
            stats[batch_idx, :3] = result
            stats[batch_idx, 3] = end - start
            stats[batch_idx, 4] = sid
            poses.append(T.numpy())

            if batch_idx % log_interval == 0 and batch_idx > 0:
                cur_stats = stats[:batch_idx]
                cur_recall = cur_stats[:, 0].mean() * 100
                cur_rte = cur_stats[cur_stats[:, 0] > 0, 1].mean()
                cur_rre = cur_stats[cur_stats[:, 0] > 0, 2].mean()
                print(
                    f"recall: {cur_recall:.2f}, rte: {cur_rte:.2f}, rre: {cur_rre:.2f}"
                )

    return subset_names, stats, np.stack(poses, axis=0)


@gin.configurable()
def test(
    out_dir,
    run_name,
    checkpoint_path,
    feature_class,
    model_class,
    dataset_class,
    log_interval,
):
    # initialize data loader
    dataset = dataset_class()
    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=1, num_workers=1, shuffle=False, collate_fn=lambda x: x
    )
    TE_THRESH = dataset.TE_THRESH
    RE_THRESH = dataset.RE_THRESH

    # initialize device
    device = torch.device("cuda")

    # initialize feature extractor
    feature_extractor = feature_class(device=device)

    # initialize refinement model
    refine_model = model_class().to(device)
    ckpt = torch.load(checkpoint_path)

    def remove_prefix(k, prefix):
        return k[len(prefix) :] if k.startswith(prefix) else k

    state_dict = {remove_prefix(k, "model."): v for k, v in ckpt["state_dict"].items()}
    refine_model.load_state_dict(state_dict)
    logging.info(f"Load refine model from checkpoint {checkpoint_path}")
    logging.info(f"number of parameters: {count_parameters(refine_model)}")
    refine_model.eval()
    dhvr = DHVR(
        device=device, feature_extractor=feature_extractor, refine_model=refine_model
    )

    # run benchmark
    subset_names, stats, poses = run_benchmark(
        dataloader,
        method=dhvr,
        TE_THRESH=TE_THRESH,
        RE_THRESH=RE_THRESH,
        log_interval=log_interval,
    )

    # print_table
    print_table(subset_names, stats, TE_THRESH, RE_THRESH)

    # save results
    exp_dir = os.path.join(out_dir, run_name)
    ensure_dir(exp_dir)
    stat_filename = os.path.join(exp_dir, "stats.npz")
    conf_filename = os.path.join(exp_dir, "config.gin")
    np.savez(stat_filename, stats=stats, names=["dhvr"], poses=poses)
    with open(conf_filename, "w") as f:
        f.write(gin.operative_config_str())
    logging.info(f"Saved results to {stat_filename}, {conf_filename}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("config", type=str, help="path to config file")
    parser.add_argument("--run_name", type=str, required=True, help="experiment title")
    parser.add_argument(
        "--load_path", type=str, required=True, help="path to checkpoint"
    )
    parser.add_argument(
        "--out_dir",
        type=str,
        default="experiments",
        help="path to save benchmark results",
    )
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--seed", type=int, default=1234)
    args = parser.parse_args()

    # random seed
    pl.seed_everything(args.seed)

    # setup config and logger
    gin.parse_config_file(args.config)
    setup_logger(args.run_name, args.debug)

    # start test
    test(args.out_dir, args.run_name, args.load_path)