evaluate.py

#!/usr/bin/env python
""" Inference script

Hacked together by Ross Wightman (https://github.com/rwightman)
"""
import argparse
import json
import os
from contextlib import suppress
from pathlib import Path
from sys import prefix
from types import SimpleNamespace
from attr import dataclass

import numpy as np
import open3d as o3d
from sympy import false
import timm.utils as utils
import torch
import torch.nn.parallel
import tqdm
from sklearn.neighbors import NearestNeighbors
from timm.utils import setup_default_logging

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"

import cv2  # noqa: E402

import dataset.transformations as transformations
from effdet import create_dataset, create_loader, create_model, rotation6d
from effdet.data import resolve_input_config, transforms
import pose_metrics
from vis_utils import (
    visualize_2d,
    save_2d_vis,
    save_3d_vis,
    visualize_3d,
    get_base_mesh,
    get_depth_pcd,
    get_camera_facing_pcd,
    transform_meshes,
)

try:
    from timm.layers import set_layer_config
except ImportError:
    from timm.models.layers import set_layer_config

has_apex = False
try:
    from apex import amp

    has_apex = True
except ImportError:
    pass

has_native_amp = False
try:
    if getattr(torch.cuda.amp, "autocast") is not None:
        has_native_amp = True
except AttributeError:
    pass

torch.backends.cudnn.benchmark = True

setup_default_logging()


def add_bool_arg(parser, name, default=False, help=""):  # FIXME move to utils
    dest_name = name.replace("-", "_")
    group = parser.add_mutually_exclusive_group(required=False)
    group.add_argument("--" + name, dest=dest_name, action="store_true", help=help)
    group.add_argument("--no-" + name, dest=dest_name, action="store_false", help=help)
    parser.set_defaults(**{dest_name: default})


parser = argparse.ArgumentParser(description="PyTorch ImageNet Validation")
parser.add_argument("root", metavar="DIR", help="path to dataset root")
parser.add_argument(
    "--dataset",
    default="coco",
    type=str,
    metavar="DATASET",
    help='Name of dataset (default: "coco"',
)
parser.add_argument("--split", default="val", help="validation split")
parser.add_argument(
    "--model",
    "-m",
    metavar="MODEL",
    default="tf_efficientdet_d1",
    help="model architecture (default: tf_efficientdet_d1)",
)
add_bool_arg(
    parser,
    "redundant-bias",
    default=None,
    help="override model config for redundant bias layers",
)
add_bool_arg(
    parser, "soft-nms", default=None, help="override model config for soft-nms"
)
add_bool_arg(
    parser, "visualize", default=False, help="Visualize detections and ground truth"
)
add_bool_arg(parser, "refine", default=True, help="ICP Refinement of predictions")
add_bool_arg(parser, "export", default=False, help="Export predictions to JSON file")
add_bool_arg(
    parser, "dump", default=False, help="Only dump targets and outputs to file."
)
parser.add_argument(
    "--num-classes",
    type=int,
    default=None,
    metavar="N",
    help="Override num_classes in model config if set. For fine-tuning from pretrained.",
)
parser.add_argument(
    "-j",
    "--workers",
    default=4,
    type=int,
    metavar="N",
    help="number of data loading workers (default: 4)",
)
add_bool_arg(parser, "persistent-workers", default=False, help="Keep workers alive.")
parser.add_argument(
    "-b",
    "--batch-size",
    default=4,
    type=int,
    metavar="N",
    help="mini-batch size (default: 128)",
)
parser.add_argument(
    "--img-size",
    default=None,
    type=int,
    metavar="N",
    help="Input image dimension, uses model default if empty",
)
parser.add_argument(
    "--mean",
    type=float,
    nargs="+",
    default=None,
    metavar="MEAN",
    help="Override mean pixel value of dataset",
)
parser.add_argument(
    "--std",
    type=float,
    nargs="+",
    default=None,
    metavar="STD",
    help="Override std deviation of of dataset",
)
parser.add_argument(
    "--interpolation",
    default="bilinear",
    type=str,
    metavar="NAME",
    help="Image resize interpolation type (overrides model)",
)
parser.add_argument(
    "--fill-color",
    default=None,
    type=str,
    metavar="NAME",
    help='Image augmentation fill (background) color ("mean" or int)',
)
parser.add_argument(
    "--log-freq",
    default=10,
    type=int,
    metavar="N",
    help="batch logging frequency (default: 10)",
)
parser.add_argument(
    "--checkpoint",
    default="",
    type=str,
    metavar="PATH",
    help="path to latest checkpoint (default: none)",
)
parser.add_argument(
    "--pretrained", dest="pretrained", action="store_true", help="use pre-trained model"
)
parser.add_argument("--num-gpu", type=int, default=1, help="Number of GPUS to use")
parser.add_argument(
    "--no-prefetcher",
    action="store_true",
    default=False,
    help="disable fast prefetcher",
)
parser.add_argument(
    "--pin-mem",
    action="store_true",
    default=False,
    help="Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.",
)
parser.add_argument(
    "--use-ema",
    dest="use_ema",
    action="store_true",
    help="use ema version of weights if present",
)
parser.add_argument(
    "--amp",
    action="store_true",
    default=False,
    help="Use AMP mixed precision. Defaults to Apex, fallback to native Torch AMP.",
)
parser.add_argument(
    "--apex-amp",
    action="store_true",
    default=False,
    help="Use NVIDIA Apex AMP mixed precision",
)
parser.add_argument(
    "--native-amp",
    action="store_true",
    default=False,
    help="Use Native Torch AMP mixed precision",
)
parser.add_argument(
    "--torchscript",
    dest="torchscript",
    action="store_true",
    help="convert model torchscript for inference",
)
parser.add_argument(
    "--torchcompile",
    nargs="?",
    type=str,
    default=None,
    const="inductor",
    help="Enable compilation w/ specified backend (default: inductor).",
)
parser.add_argument(
    "--results",
    default="",
    type=str,
    metavar="FILENAME",
    help="JSON filename for evaluation results",
)
parser.add_argument(
    "--z-offset",
    default=0.0,
    type=float,
    metavar="Z_OFFSET",
    help="Offset the z-predictions in mm (ImagePose only!).",
)
parser.add_argument(
    "--min-score",
    default=0.9,
    type=float,
    help="Minimum score for predictions to be considered valid.",
)
parser.add_argument(
    "--topk", default=0, type=int, help="Only use the top k predictions per image."
)

# EfficientPose architecture options
utils.add_bool_arg(
    parser,
    "combined-t-head",
    default=False,
    help="Use combined head for xy and z regression",
)
utils.add_bool_arg(
    parser,
    "iterative-pose-heads",
    default=False,
    help="Use iterative heads for rotation and translation.",
)
utils.add_bool_arg(
    parser,
    "axis-angle",
    default=False,
    help="Use axis angle instead of 6D rotation representation.",
)

# EfficientPose loss options
utils.add_bool_arg(
    parser,
    "mean-dist-z",
    default=False,
    help="Use mean distance weighted z loss.",
)
utils.add_bool_arg(
    parser,
    "mean-dist-rot",
    default=False,
    help="Use mean distance weighted z loss.",
)

# EfficientPose IoU options
parser.add_argument(
    "--detection-iou",
    type=float,
    default=0.5,
    help="IoU threshold for the NMS of the detections (default: 0.5).",
)

# EfficientPose Augmentations
utils.add_bool_arg(
    parser,
    "mask-aug",
    default=False,
    help="Enable Mask and Threshold augmentations (train and eval).",
)
utils.add_bool_arg(
    parser,
    "color-aug",
    default=False,
    help="Enable color space augmentations (only train).",
)
utils.add_bool_arg(
    parser,
    "rot-aug",
    default=False,
    help="6D Augmentation - Enable stepped rot augmentation (only train).",
)
utils.add_bool_arg(
    parser,
    "z-aug",
    default=False,
    help="6D Augmentation - Enable Z augmentation (only train).",
)
parser.add_argument(
    "--z-aug-min",
    type=float,
    default=0.9,
    help="6D Augmentation - MIN scale factor (default: 0.9).",
)
parser.add_argument(
    "--z-aug-max",
    type=float,
    default=1.2,
    help="6D Augmentation - MAX scale factor (default: 1.2).",
)
add_bool_arg(
    parser, "save-vis2d", default=False, help="Save images with predictions to disk."
)
add_bool_arg(
    parser, "save-vis3d", default=False, help="Render scenes with predictions to disk."
)


def evaluate(args):
    if isinstance(args, dict):
        # This makes the dict items accessible via dot notation.
        # Only really needed if train is called with a dictionary instead
        # of the parser output directly.
        args = SimpleNamespace(**args)

    if args.amp:
        if has_native_amp:
            args.native_amp = True
        elif has_apex:
            args.apex_amp = True
    assert not args.apex_amp or not args.native_amp, "Only one AMP mode should be set."
    args.pretrained = (
        args.pretrained or not args.checkpoint
    )  # might as well try to validate something
    args.prefetcher = not args.no_prefetcher

    # create model
    with set_layer_config(scriptable=args.torchscript):
        extra_args = {}
        if args.img_size is not None:
            extra_args = dict(image_size=(args.img_size, args.img_size))
        bench = create_model(
            args.model,
            bench_task="predict",
            num_classes=args.num_classes,
            pretrained=args.pretrained,
            redundant_bias=args.redundant_bias,
            soft_nms=args.soft_nms,
            checkpoint_path=args.checkpoint,
            checkpoint_ema=args.use_ema,
            combined_t_head=args.combined_t_head,
            iterative_pose_heads=args.iterative_pose_heads,
            axis_angle=args.axis_angle,
            mean_dist_z=args.mean_dist_z,
            mean_dist_rot=args.mean_dist_rot,
            detection_iou=args.detection_iou,
            **extra_args,
        )
    model_config = bench.config

    param_count = sum([m.numel() for m in bench.parameters()])
    print("Model %s created, param count: %d" % (args.model, param_count))

    bench = bench.cuda()

    if args.torchscript:
        assert (
            not args.apex_amp
        ), "Cannot use APEX AMP with torchscripted model, force native amp with `--native-amp` flag"
        bench = torch.jit.script(bench)
    elif args.torchcompile:
        bench = torch.compile(bench, backend=args.torchcompile)

    amp_autocast = suppress
    if args.apex_amp:
        bench = amp.initialize(bench, opt_level="O1")
        print("Using NVIDIA APEX AMP. Validating in mixed precision.")
    elif args.native_amp:
        amp_autocast = torch.cuda.amp.autocast
        print("Using native Torch AMP. Validating in mixed precision.")
    else:
        print("AMP not enabled. Validating in float32.")

    if args.num_gpu > 1:
        bench = torch.nn.DataParallel(bench, device_ids=list(range(args.num_gpu)))

    transforms_eval = transforms.build_imagepose_transform(args.mask_aug)

    dataset = create_dataset(args.dataset, args.root, args.split)
    input_config = resolve_input_config(args, model_config)

    loader = create_loader(
        dataset,
        input_size=input_config["input_size"],
        batch_size=args.batch_size,
        use_prefetcher=args.prefetcher,
        interpolation=input_config["interpolation"],
        fill_color=input_config["fill_color"],
        mean=input_config["mean"],
        std=input_config["std"],
        num_workers=args.workers,
        persistent_workers=args.persistent_workers,
        pin_mem=args.pin_mem,
        transform_fn=transforms_eval,
        axis_angle=args.axis_angle,
    )

    data_folder = Path(args.root) / args.split
    base_mesh = get_base_mesh()
    has_groundtruth = dataset.parser.has_labels
    score_threshold = args.min_score
    icp_threshold = 0.7

    checkpoint_path = Path(args.checkpoint)
    if dataset.NAME == "imagepose":
        z_offset = args.z_offset
    elif dataset.NAME == "depthpose":
        z_offset = dataset.MEAN[0] - args.mean[0] if args.mean else 0
        z_offset *= 1000  # convert to mm.

    total_metrics = {"metrics": {}, "details": {}}
    total_metrics_icp = {"metrics": {}, "details": {}}

    bench.eval()
    with torch.inference_mode():
        tqdm_loader = tqdm.tqdm(loader, unit_scale=args.batch_size)

        for input, target in tqdm_loader:
            # image_grid = torchvision.utils.make_grid(input, nrow=2, normalize=True, scale_each=True)
            # cv2.imshow("train_grid", image_grid.numpy().transpose(1, 2, 0))
            # cv2.waitKey(1)

            with amp_autocast():
                output = bench(input, img_info=target)

            for batch_idx in range(input.size(0)):
                batch_target = {key: value[batch_idx] for key, value in target.items()}
                batch_input = input[batch_idx]
                batch_output = output[batch_idx]

                targets = postprocess_targets(batch_target, has_groundtruth, args.axis_angle)

                ortho_scale = targets["ortho_scale"]
                K = targets["K"]
                img_size = targets["img_size"]
                img_id = targets["img_idx"]

                T_blender_world = transformations.get_T_blender_world(ortho_scale)
                T_world_cam = transformations.get_T_world_cam(
                    targets["cam_R"], targets["cam_t"]
                )

                predictions = postprocess_predictions(
                    batch_output,
                    score_threshold,
                    dataset,
                    ortho_scale,
                    img_size,
                    T_blender_world,
                    T_world_cam,
                    K,
                    z_offset,
                    top_k=args.topk,
                    axis_angle=args.axis_angle
                )

                if dataset.NAME == "depthpose":
                    file_name = dataset.parser.img_infos[img_id]["file_name"]
                    depth_path = data_folder / file_name
                    depth_pcd = get_depth_pcd(
                        depth_path,
                        ortho_scale,
                    )
                else:
                    # NOTE: We could try to map the image to the depth image here,
                    #       but we would need to pass the path to the depthpose
                    #       dataset as well. Implement only if necessary.
                    # HACK: This is only for debugging purposes.
                    # TODO: Use an absolute path for the depth_name, then this
                    #       is no longer a hack but a feature (if the file is found).
                    depth_name = dataset.parser.img_infos[img_id]["depth_name"]
                    depth_path = (
                        data_folder.parent.parent
                        / "fleckenzwerg_dataset_synth"
                        / "real"
                        / depth_name
                    )
                    if depth_path.is_file():
                        depth_pcd = get_depth_pcd(
                            depth_path,
                            ortho_scale,
                        )
                    else:
                        # print("WARNING: Hardcoded depth image path hack enabled!")
                        depth_pcd = o3d.t.geometry.PointCloud()

                metrics = evaluate_image_metrics(
                    base_mesh, predictions, targets, dataset.MODEL_DIAMETER
                )

                update_total_metrics(total_metrics, metrics)
                update_tqdm_metrics(tqdm_loader, total_metrics["metrics"])

                if args.visualize:
                    visualize_2d(
                        batch_input,
                        predictions,
                        targets,
                        T_blender_world,
                        T_world_cam,
                        dataset,
                        checkpoint_path.stem,
                    )
                    visualize_3d(base_mesh, predictions, targets, dataset, depth_pcd)

                if args.save_vis2d:
                    save_2d_vis(
                        batch_input,
                        predictions,
                        targets,
                        T_blender_world,
                        T_world_cam,
                        dataset,
                        checkpoint_path,
                        score_threshold,
                    )

                if args.save_vis3d:
                    save_3d_vis(
                        base_mesh,
                        predictions,
                        targets,
                        dataset,
                        depth_pcd,
                        checkpoint_path,
                        score_threshold,
                    )

                if dataset.NAME == "depthpose" and args.refine:
                    refine_predictions(base_mesh, predictions, depth_pcd, icp_threshold)

                    metrics_icp = evaluate_image_metrics(
                        base_mesh, predictions, targets, dataset.MODEL_DIAMETER
                    )

                    update_total_metrics(total_metrics_icp, metrics_icp)
                    update_tqdm_metrics(tqdm_loader, total_metrics_icp["metrics"])

                    if args.visualize:
                        visualize_3d(
                            base_mesh,
                            predictions,
                            targets,
                            dataset,
                            depth_pcd,
                            show_groundtruth=False,
                        )

    save_metrics(total_metrics, checkpoint_path.parent, args.split, score_threshold)


def update_tqdm_metrics(tqdm_loader, metrics, prefix=""):
    postfixes = {prefix + key: value for key, value in metrics.items()}
    tqdm_loader.set_postfix(postfixes)


def postprocess_predictions(
    output,
    threshold,
    dataset,
    ortho_scale,
    img_size,
    T_blender_world,
    T_world_cam,
    K,
    z_offset,
    top_k: int = 0,
    axis_angle: bool = False,
):
    # output format: [y, x, y, x, score, class, r1, r2, r3, r4, r5, r6, ty, tx, tz]

    # Filter by score.
    scores = output[:, 4].cpu().numpy()
    if top_k > 0:
        # Scores were already sorted, so just use the top k.
        solid_detections = np.arange(min(top_k, len(scores)))
    else:
        solid_detections = scores > threshold

    scores = scores[solid_detections]

    # Extract predictions.
    bbox_preds = output[solid_detections, :4].cpu().numpy()  # 2D bounding boxes
    # bbox_preds = bbox_preds[:, [1, 0, 3, 2]]  # yxyx -> xyxy
    class_preds = output[solid_detections, 5].cpu().numpy()
    if not axis_angle:
        rotation6d_preds = output[solid_detections, 6:12]
        rotation_preds = (
            rotation6d.compute_rotation_matrix_from_ortho6d(rotation6d_preds).cpu().numpy()
        )
        centerpoint_preds = output[solid_detections, 12:14].cpu().numpy()
        translation_z_preds = output[solid_detections, 14].cpu().numpy()
    else:
        rotation_aa_preds = output[solid_detections, 6:9].cpu().numpy()
        rotation_preds = np.array([cv2.Rodrigues(aa)[0] for aa in rotation_aa_preds])
        centerpoint_preds = output[solid_detections, 9:11].cpu().numpy()
        translation_z_preds = output[solid_detections, 11].cpu().numpy()

    if dataset.NAME == "depthpose":
        translation_z_preds -= z_offset
        translation_preds = transformations.orthographic_unproject(
            centerpoint_preds,
            translation_z_preds,
            img_size,
            ortho_scale,
        )
    elif dataset.NAME == "imagepose":
        translation_z_preds += z_offset
        translation_preds = transformations.perspective_unproject(
            centerpoint_preds,
            translation_z_preds,
            T_blender_world,
            T_world_cam,
            K,
        )

    return {
        "scores": scores,
        "bbox": bbox_preds,
        "rotation": rotation_preds,  # Rotation Matrix
        "centerpoint": centerpoint_preds,
        "translation": translation_preds,
    }


def postprocess_targets(target, has_groundtruth=True, axis_angle=False):
    targets = {
        "img_idx": target["img_idx"].cpu().numpy(),
        "img_size": target["img_size"][0].cpu().numpy(),
        "ortho_scale": target["ortho_scale"].cpu().numpy(),
        "cam_R": target["cam_R"].cpu().numpy(),
        "cam_t": target["cam_t"].cpu().numpy(),
        "K": target["K"].cpu().numpy().reshape(3, 3),
    }

    if has_groundtruth:
        centerpoint_targets = target["centerpoint"].cpu().numpy()
        solid_targets = centerpoint_targets[:, 0] != -1
        centerpoint_targets = centerpoint_targets[solid_targets][:, [1, 0]]  # yx -> xy
        bbox_targets = target["bbox"][solid_targets].cpu().numpy()
        bbox_targets = bbox_targets[:, [1, 0, 3, 2]]  # yxyx -> xyxy
        rotation_targets = target["rotation"][solid_targets].cpu().numpy()
        rotation_targets = rotation_targets.reshape(-1, 3, 3)
        if not axis_angle:
            rotation_targets_out = rotation6d.compute_rotation_matrix_from_ortho6d(
                target["rotation6d"][solid_targets]
            )
            rotation_targets_out = rotation_targets_out.cpu().numpy()

            # Quick sanity check to see whether the conversion is working correctly.
            assert np.allclose(
                rotation_targets_out, rotation_targets
            ), "6D Rotation conversion failed."
        else:
            rotation_targets_out = rotation_targets

        translation_targets = target["translation"][solid_targets].cpu().numpy()
        targets["bbox"] = bbox_targets
        targets["rotation"] = rotation_targets_out
        targets["centerpoint"] = centerpoint_targets
        targets["translation"] = translation_targets

    return targets


def get_icp_transformation(
    mesh: o3d.geometry.TriangleMesh,
    pcd: o3d.geometry.PointCloud,
):
    mesh_pcd = get_camera_facing_pcd(mesh)  # This will be around 25-50k points.
    bbox = mesh_pcd.get_axis_aligned_bounding_box()
    bbox.scale(1.3, center=bbox.get_center())

    # Cut the pcd around the model's bounding box.
    sub_pcd = pcd.crop(bbox)

    # Skip if the pcd is too small.
    if len(sub_pcd.point.positions) < 100:
        print("Skipping ICP, PCD is too small.")
        return np.eye(4), 0, 0

    icp = o3d.t.pipelines.registration.icp(
        mesh_pcd,
        sub_pcd,
        4,
        np.eye(4),
        o3d.t.pipelines.registration.TransformationEstimationPointToPoint(),
        o3d.t.pipelines.registration.ICPConvergenceCriteria(
            relative_fitness=1e-4,  # default 1e-6
            relative_rmse=1e-4,  # default 1e-6
            max_iteration=30,  # default 30
        ),
    )
    return icp.transformation.cpu().numpy(), icp.fitness, icp.inlier_rmse


def refine_predictions(
    base_mesh: o3d.t.geometry.TriangleMesh,
    predictions: dict,
    depth_pcd: o3d.t.geometry.PointCloud,
    fitness_threshold: float = 0.7,
):
    meshes_pred = transform_meshes(
        base_mesh, predictions["rotation"], predictions["translation"]
    )

    icp_transformations = []
    valid_predictions = []
    for idx, mesh in enumerate(meshes_pred):
        transformation, fitness, rmse = get_icp_transformation(mesh, depth_pcd)
        if fitness > fitness_threshold:
            icp_transformations.append(transformation)
            valid_predictions.append(idx)

    # Filter out invalid predictions.
    for key in ["scores", "bbox", "rotation", "centerpoint", "translation"]:
        predictions[key] = predictions[key][valid_predictions]

    for idx in range(len(valid_predictions)):
        R = predictions["rotation"][idx]
        t = predictions["translation"][idx]
        transformation = np.eye(4)
        transformation[:3, :3] = R
        transformation[:3, 3] = t
        result = icp_transformations[idx] @ transformation
        predictions["rotation"][idx] = result[:3, :3]
        predictions["translation"][idx] = result[:3, 3]


def match_targets(
    predictions: dict,
    targets: dict,
    diameter: float,
    threshold: float = 0.1,
):
    """Match up predictions and targets via distance of their centerpoints.

    Each target is to be used at most once, duplicate assignments are discarded.
    """
    pred_centerpoints = predictions["centerpoint"]
    target_centerpoints = targets.get("centerpoint", [])

    if len(target_centerpoints) == 0:
        raise NotImplementedError
    elif len(pred_centerpoints) == 0:
        raise NotImplementedError

    # Build a KDTree containing the targets.
    # (Actually set to "auto" so for small numbers of targets it uses brute force.)
    # Then query the tree for each prediction.
    knn = NearestNeighbors(n_neighbors=1).fit(target_centerpoints)
    knn_distances, knn_indices = knn.kneighbors(pred_centerpoints)
    pred_indices = np.arange(len(pred_centerpoints))

    # Too far away to be considered a match.
    # Uses a "magic number" as maximum distance
    max_distance = threshold * diameter * 2  # (= 28px)
    too_far_indices = np.nonzero(knn_distances > max_distance)[0]
    pred_indices = np.delete(pred_indices, too_far_indices)
    knn_indices = np.delete(knn_indices, too_far_indices)
    knn_distances = np.delete(knn_distances, too_far_indices)

    # Sort the pairs by distance and therefore confidence.
    sorting_indices = np.argsort(knn_distances.flatten())
    target_indices = knn_indices[sorting_indices]
    pred_indices = pred_indices[sorting_indices]

    # Make sure that each target is only matched to its first/closest pred.
    # Duplicate assignments of a pred to a target are discarded.
    _, unique_indices = np.unique(target_indices, return_index=True)
    target_indices = target_indices[unique_indices]
    pred_indices = pred_indices[unique_indices]

    # Number of targets used multiple times
    too_far_count = len(too_far_indices)

    # Mask which targets have been used.
    target_mask = np.zeros(len(target_centerpoints), dtype=bool)
    target_mask[target_indices] = True

    # Number of targets missed (includes the too_far_count).
    false_negative = len(target_centerpoints) - len(target_indices)

    return pred_indices, target_indices, false_negative, too_far_count, target_mask


def evaluate_image_metrics(
    base_mesh: o3d.t.geometry.TriangleMesh,
    predictions: dict,
    targets: dict,
    diameter: float,
    threshold: float = 0.1,
):
    results = {
        # To be used for presentation of current image.
        "metrics": {
            "Precision": 0,
            "Recall": 0,
            "F1": 0,
            "ADD Mean": 0,
            "ADD Std": 0,
            "T Err Mean": 0,
            "T Err Std": 0,
            "Rot Err Mean": 0,
            "Rot Err Std": 0,
        },
        # To be used for re-calculation of metrics over the whole batch/dataset.
        "details": {
            "TP": 0,
            "FP": 0,
            "FN": 0,
            "ADD": [],
            "T Err": [],
            "Rot Err": [],
            "Target Mask": [],  # A mask of the targets that have been matched.
            "Target Count": 0,
        },
    }

    if len(targets["centerpoint"]) == 0:
        print("No targets? No Metrics!")
        return results

    if len(predictions["centerpoint"]) == 0:
        print("No predictions? No Candies!")
        results["details"]["FN"] = len(targets["centerpoint"])
        results["details"]["Target Count"] = len(targets["centerpoint"])
        results["details"]["Target Mask"] = np.zeros(len(targets["centerpoint"]), dtype=bool)
        return results

    model_points = base_mesh.vertex.positions[::1000].cpu().numpy()  # 171 points.

    pred_indices, target_indices, false_negative, too_far_count, target_mask = match_targets(
        predictions, targets, diameter, threshold
    )
    results["details"]["Target Count"] = len(targets["centerpoint"])
    results["details"]["Target Mask"] = target_mask

    target_t = targets["translation"][target_indices]
    target_R = targets["rotation"][target_indices]
    pred_t = predictions["translation"][pred_indices]
    pred_R = predictions["rotation"][pred_indices]

    # Average Distance of Model Points (ADD)
    add_metric, add_true_positive_mask = pose_metrics.calculate_add(
        pred_t, pred_R, target_t, target_R, model_points, diameter, threshold
    )
    add_metric_mean = np.mean(add_metric, axis=-1)  # mean per image
    add_metric_std = np.std(add_metric, axis=-1)  # std per image
    add_true_positive = np.sum(add_true_positive_mask)
    add_false_positive = len(add_true_positive_mask) - add_true_positive

    # Translation Error Metric
    t_metric = pose_metrics.calculate_translation_error(pred_t, target_t)
    t_metric_mean = np.mean(np.abs(t_metric), axis=-1)  # mean per image
    t_metric_std = np.std(np.abs(t_metric), axis=-1)  # std per image

    # XY Error Metric
    xy_metric = pose_metrics.calculate_xy_error(pred_t, target_t)
    xy_metric_mean = np.mean(xy_metric, axis=-1)  # mean per image
    xy_metric_std = np.std(xy_metric, axis=-1)  # std per image

    # Depth Error Metric
    z_metric = pose_metrics.calculate_depth_error(pred_t, target_t)
    z_metric_mean = np.mean(z_metric, axis=-1)  # mean per image
    z_metric_std = np.std(z_metric, axis=-1)  # std per image

    # Rotation Error Metric
    r_metric = pose_metrics.calculate_rotation_error(pred_R, target_R)
    r_metric_mean = np.mean(r_metric, axis=-1)  # mean per image
    r_metric_std = np.std(r_metric, axis=-1)  # std per image

    # NOTE: Accuracy requires True Negatives (TN) which are undefined for detection.
    precision = add_true_positive / (add_true_positive + add_false_positive)
    recall = add_true_positive / (add_true_positive + false_negative)
    f1_score = 2 * (precision * recall) / (precision + recall)

    # To be used for presentation of current image.
    results["metrics"]["Precision"] = precision
    results["metrics"]["Recall"] = recall
    results["metrics"]["F1"] = f1_score
    results["metrics"]["ADD Mean"] = add_metric_mean
    results["metrics"]["ADD Std"] = add_metric_std
    results["metrics"]["T Err Mean"] = t_metric_mean
    results["metrics"]["T Err Std"] = t_metric_std
    results["metrics"]["XY Err Mean"] = xy_metric_mean
    results["metrics"]["XY Err Std"] = xy_metric_std
    results["metrics"]["Z Err Mean"] = z_metric_mean
    results["metrics"]["Z Err Std"] = z_metric_std
    results["metrics"]["Rot Err Mean"] = r_metric_mean
    results["metrics"]["Rot Err Std"] = r_metric_std

    # To be used for re-calculation of metrics over the whole batch/dataset.
    results["details"]["TP"] = add_true_positive
    results["details"]["FP"] = add_false_positive
    results["details"]["FN"] = false_negative
    results["details"]["ADD"] = add_metric
    results["details"]["T Err"] = t_metric
    results["details"]["XY Err"] = xy_metric
    results["details"]["Z Err"] = z_metric
    results["details"]["Rot Err"] = r_metric

    return results


def update_total_metrics(total_metrics, image_metrics):
    """Update the total metrics with the current image metrics in place.

    This is used to calculate the mean and std of the metrics over the whole
    batch and dataset.
    """
    # Concatenate the new details.
    for metric, value in image_metrics["details"].items():
        if metric in total_metrics["details"]:
            if not isinstance(value, np.ndarray):
                value = np.array([value])
                if value.shape[-1] == 0:
                    continue
            total_metrics["details"][metric] = np.concatenate(
                (total_metrics["details"][metric], value), axis=-1
            )
        else:
            if not isinstance(value, np.ndarray):
                value = np.array([value])
                if value.shape[-1] == 0:
                    continue
            total_metrics["details"][metric] = np.array(value)

    # Calculate the mean and std of the metrics.
    mean_std_metric_keys = ["ADD", "T Err", "XY Err", "Rot Err"]
    for metric in mean_std_metric_keys:
        # -1 is fine as backup, since it will be re-calculated each frame anyways.
        total_metrics["metrics"][metric + " Mean"] = np.mean(
            total_metrics["details"].get(metric, -1)
        )
        total_metrics["metrics"][metric + " Std"] = np.std(
            total_metrics["details"].get(metric, -1)
        )

    # Calculating Z separately since we need absolute values to make sense here.
    # -1 is fine as backup, since it will be re-calculated each frame anyways.
    total_metrics["metrics"]["Z Err Mean"] = np.mean(
        np.abs(total_metrics["details"].get("Z Err", -1))
    )
    total_metrics["metrics"]["Z Err Std"] = np.std(
        np.abs(total_metrics["details"].get("Z Err", -1))
    )

    # Calculate Precision, Recall and F1.
    TP = np.sum(total_metrics["details"]["TP"])
    FP = np.sum(total_metrics["details"]["FP"])
    FN = np.sum(total_metrics["details"]["FN"])
    precision = TP / (TP + FP)
    recall = TP / (TP + FN)
    f1_score = 2 * (precision * recall) / (precision + recall)

    total_metrics["metrics"]["Precision"] = precision
    total_metrics["metrics"]["Recall"] = recall
    total_metrics["metrics"]["F1"] = f1_score


def save_metrics(total_metrics: dict, output_dir: Path, split: str, min_score: float):
    """Save the metrics to a JSON file."""
    metrics = total_metrics["metrics"]
    details = total_metrics["details"]

    # Convert numpy arrays to lists.
    for key, value in metrics.items():
        if isinstance(value, np.ndarray):
            metrics[key] = value.tolist()

    for key, value in details.items():
        if isinstance(value, np.ndarray):
            details[key] = value.tolist()

    metrics_path = output_dir / f"{split}_{min_score}_metrics.json"
    details_path = output_dir / f"{split}_{min_score}_details.json"

    with metrics_path.open("w") as f:
        json.dump(metrics, f, indent=4)

    with details_path.open("w") as f:
        json.dump(details, f, indent=4)


if __name__ == "__main__":
    args = parser.parse_args()
    evaluate(args)