inference.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

import numpy as np
import open3d as o3d
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
from vis_utils import (
    visualize_2d,
    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


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).",
)


def inference(args):
    setup_default_logging()

    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

    # Pre ICP
    add_meter_pre_score = utils.AverageMeter()
    add_meter_pre_precision = utils.AverageMeter()
    add_meter_pre_recall = utils.AverageMeter()
    t_err_meter_pre = utils.AverageMeter()
    rot_err_meter_pre = utils.AverageMeter()
    max_t_err_meter_pre = np.array([0.0, 0.0, 0.0])
    max_rot_err_meter_pre = np.array([0.0, 0.0, 0.0])

    # Post ICP
    add_meter_post_score = utils.AverageMeter()
    add_meter_post_precision = utils.AverageMeter()
    add_meter_post_recall = utils.AverageMeter()
    t_err_meter_post = utils.AverageMeter()
    rot_err_meter_post = utils.AverageMeter()
    max_t_err_meter_post = np.array([0.0, 0.0, 0.0])
    max_rot_err_meter_post = np.array([0.0, 0.0, 0.0])

    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.

    bench.eval()
    with torch.inference_mode():
        tqdm_postfixes = ["mADD", "Precision", "Recall", "t_err", "rot_err"]

        if dataset.NAME == "depthpose" and args.refine:
            tqdm_postfixes += ["mADDr", "Precisionr", "Recallr"]

        tqdm_loader = tqdm.tqdm(
            loader,
            unit_scale=args.batch_size,
            postfix={p: 0 for p in tqdm_postfixes},
        )

        targets = []
        outputs = []
        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)

            if args.dump:
                targets.append(target)
                outputs.append(output.cpu().numpy())
                # Skip the rest
                continue

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

                targets = postprocess_targets(batch_target, has_groundtruth)

                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,
                )

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

                (
                    add_metric,
                    add_precision,
                    add_recall,
                    mean_t_error,
                    mean_rot_error,
                    max_t_error,
                    max_rot_error
                ) = evaluate_metrics(
                    base_mesh, predictions, targets, dataset.MODEL_DIAMETER
                )
                if add_metric is not None:
                    add_meter_pre_score.update(add_metric)
                if add_precision is not None:
                    add_meter_pre_precision.update(add_precision)
                if add_recall is not None:
                    add_meter_pre_recall.update(add_recall)
                if mean_t_error is not None:
                    t_err_meter_pre.update(mean_t_error)
                if mean_rot_error is not None:
                    rot_err_meter_pre.update(mean_rot_error)
                if max_t_error is not None:
                    max_t_err_meter_pre = np.max([max_t_error, max_t_err_meter_pre], axis=0)
                if max_rot_error is not None:
                    max_rot_err_meter_pre = np.max([max_rot_error, max_rot_err_meter_pre], axis=0)

                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 dataset.NAME == "depthpose" and args.refine:
                    refine_predictions(base_mesh, predictions, depth_pcd, icp_threshold)

                    (
                        add_metric,
                        add_precision,
                        add_recall,
                        mean_t_error,
                        mean_rot_error,
                        max_t_error,
                        max_rot_error
                    ) = evaluate_metrics(
                        base_mesh, predictions, targets, dataset.MODEL_DIAMETER
                    )
                    if add_metric is not None:
                        add_meter_post_score.update(add_metric)
                    if add_precision is not None:
                        add_meter_post_precision.update(add_precision)
                    if add_recall is not None:
                        add_meter_post_recall.update(add_recall)
                    if mean_t_error is not None:
                        t_err_meter_post.update(mean_t_error)
                    if mean_rot_error is not None:
                        rot_err_meter_post.update(mean_rot_error)
                    if max_t_error is not None:
                        max_t_err_meter_post = np.max([max_t_error, max_t_err_meter_post], axis=0)
                    if max_rot_error is not None:
                        max_rot_err_meter_post = np.max([max_rot_error, max_rot_err_meter_post], axis=0)

                    if add_meter_post_score.val > add_meter_pre_score.val:
                        print(
                            "ICP worsened the ADD metric: "
                            f"{add_meter_pre_score.val} -> {add_meter_post_score.val}"
                        )
                        print(dataset.parser.img_infos[img_id]["file_name"])

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

                    tqdm_loader.set_postfix(
                        {
                            "mADD": add_meter_pre_score.avg,
                            "Precision": add_meter_pre_precision.avg,
                            "Recall": add_meter_pre_recall.avg,
                            "t_err": t_err_meter_pre.avg,
                            "rot_err": rot_err_meter_pre.avg,
                            "max_t_err": max_t_err_meter_pre,
                            "max_rot_err": max_rot_err_meter_pre,
                            "mADDr": add_meter_post_score.avg,
                            "Precisionr": add_meter_post_precision.avg,
                            "Recallr": add_meter_post_recall.avg,
                            "t_err_r": t_err_meter_post.avg,
                            "rot_err_r": rot_err_meter_post.avg,
                            "max_t_err_r": max_t_err_meter_post,
                            "max_rot_err_r": max_rot_err_meter_post,
                        }
                    )
                else:
                    tqdm_loader.set_postfix(
                        {
                            "mADD": add_meter_pre_score.avg,
                            "Precision": add_meter_pre_precision.avg,
                            "Recall": add_meter_pre_recall.avg,
                            "t_err": t_err_meter_pre.avg,
                            "rot_err": rot_err_meter_pre.avg,
                            "max_t_err": max_t_err_meter_pre,
                            "max_rot_err": max_rot_err_meter_pre,
                        }
                    )

                if args.export:
                    suffix = "export"
                    if args.z_aug > 0:
                        suffix += f"_z_aug_{args.z_aug}"
                    export_predictions(
                        img_id,
                        predictions,
                        ortho_scale,
                        K,
                        dataset,
                        data_folder,
                        checkpoint_path,
                        z_aug=args.z_aug,
                        suffix=suffix,
                    )

    if args.dump:
        dump(targets, outputs, checkpoint_path, args.split)
    else:
        results_path = checkpoint_path.parent / f"{checkpoint_path.stem}_results.json"
        if results_path.is_file():
            with results_path.open("r") as file:
                results = json.load(file)
        else:
            results = {}

        key = (
            f"top{args.topk}-{args.split}"
            if args.topk > 0
            else f"min{args.min_score}-{args.split}"
        )
        results[key] = (
            {
                "score_threshold": score_threshold,
                "topk": args.topk,
                "split": args.split,
                "mADD": add_meter_pre_score.avg,
                "Precision": add_meter_pre_precision.avg,
                "Recall": add_meter_pre_recall.avg,
                "mean_t_err": t_err_meter_pre.avg.tolist(),
                "mean_rot_err": rot_err_meter_pre.avg.tolist(),
                "max_t_err": max_t_err_meter_pre.tolist(),
                "max_rot_err": max_rot_err_meter_pre.tolist(),
                # "icp_threshold": icp_threshold,
                # "mADDr": add_meter_post_score.avg,
                # "Precisionr": add_meter_post_precision.avg,
                # "Recallr": add_meter_post_recall.avg,
                # "t_err_r": t_err_meter_post.avg.tolist(),
                # "rot_err_r": rot_err_meter_post.avg.tolist(),
                # "max_t_err_r": max_t_err_meter_post.tolist(),
                # "max_rot_err_r": max_rot_err_meter_post.tolist(),
            },
        )

        with results_path.open("w") as file:
            json.dump(results, file)


def postprocess_predictions(
    output,
    threshold,
    dataset,
    ortho_scale,
    img_size,
    T_blender_world,
    T_world_cam,
    K,
    z_offset,
    top_k: int = 0,
):
    # 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()
    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()
    # centerpoint_preds = centerpoint_preds[:, [1, 0]]  # yx -> xy
    translation_z_preds = output[solid_detections, 14].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):
    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)
        rotation_targets_from6d = rotation6d.compute_rotation_matrix_from_ortho6d(
            target["rotation6d"][solid_targets]
        )
        rotation_targets_from6d = rotation_targets_from6d.cpu().numpy()

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

        translation_targets = target["translation"][solid_targets].cpu().numpy()
        targets["bbox"] = bbox_targets
        targets["rotation"] = rotation_targets_from6d
        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 evaluate_metrics(
    base_mesh: o3d.t.geometry.TriangleMesh,
    predictions: dict,
    targets: dict,
    diameter: float,
    threshold: float = 0.1,
):
    """Calculate the ADD metric and the rotation/translation error for the
    predicted transformations.

    A Pose is correct if the mean distance is smaller than 10% of the object diameter.

    Returns:
        tuple: score, precision, recall
    """
    model_points = base_mesh.vertex.positions[::1000].cpu().numpy()

    pred_centerpoints = predictions["centerpoint"]
    target_centerpoints = targets.get("centerpoint", [])

    if len(target_centerpoints) == 0:
        return None, None, None, None, None, None, None
    elif len(pred_centerpoints) == 0:
        # No predictions, so no mADD, no Precision but a Recall of 0.
        return None, None, 0.0, None, None, None, None

    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.
    too_far_indices = np.nonzero(knn_distances > threshold * diameter * 2)[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 by distance.
    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 prediction.
    # TODO: Is this the correct way to handle this?
    #       This does **not** penalize multiple predictions for the same target.
    # print("WARNING: Currently not removing duplicate target assignments during ADD.")
    print("WARNING: Currently removing duplicate target assignments during ADD, change this.")
    _, unique_indices = np.unique(target_indices, return_index=True)
    target_indices = target_indices[unique_indices]
    pred_indices = pred_indices[unique_indices]

    true_positives = 0
    false_negatives = len(target_centerpoints) - len(target_indices)
    false_positives = len(too_far_indices)

    distances = []
    t_errors = []
    rot_errors = []
    for pred_idx, target_idx in zip(pred_indices, target_indices):
        pred_t = predictions["translation"][pred_idx]
        pred_R = predictions["rotation"][pred_idx]
        pred_points = model_points @ pred_R + pred_t

        targ_t = targets["translation"][target_idx]
        targ_R = targets["rotation"][target_idx]
        targ_points = model_points @ targ_R + targ_t

        t_error = np.abs(pred_t - targ_t)
        t_errors.append(t_error)
        rot_error = np.abs((cv2.Rodrigues(pred_R)[0] - cv2.Rodrigues(targ_R)[0]).flatten())
        rot_errors.append(rot_error)

        mean_point_distance = np.linalg.norm(pred_points - targ_points, axis=1).mean()
        if mean_point_distance < diameter * threshold:
            true_positives += 1
        else:
            false_positives += 1
        distances.append(mean_point_distance)

    if len(distances) == 0:
        # No valid predictions, so no mADD, no Precision but a Recall of 0.
        return None, None, 0.0, None, None, None, None

    mean_t_error = np.mean(t_errors, axis=0)
    max_t_error = np.max(t_errors, axis=0)
    mean_rot_error = np.mean(rot_errors, axis=0)
    max_rot_error = np.max(rot_errors, axis=0)

    precision = true_positives / (true_positives + false_positives)

    # NOTE: DepthPose might have missed an object during annotation of the
    #       ImagePose dataset, in that case the Recall would be a bit off.
    try:
        recall = true_positives / (true_positives + false_negatives)
    except ZeroDivisionError:
        recall = 0.0

    return np.mean(distances), precision, recall, mean_t_error, mean_rot_error, max_t_error, max_rot_error


def export_predictions(
    img_id: int,
    predictions: dict,
    ortho_scale: np.ndarray,
    K: np.ndarray,
    dataset,
    data_folder: Path,
    checkpoint: Path,
    z_aug: float = 0.0,
    suffix: str = "",
):
    image_name = dataset.parser.img_infos[img_id]["file_name"]
    if suffix:
        json_name = f"{image_name}_{suffix}.json"
    else:
        json_name = f"{image_name}.json"
    json_path = data_folder / json_name

    checkpoint_name = f"{checkpoint.parent.name}/{checkpoint.name}"

    shapes = []
    for idx in range(len(predictions["scores"])):
        shape = {
            "label": "fleckenzwerg",
            "shape_type": "rectangle",
            "points": predictions["bbox"][idx].reshape(2, 2).tolist(),
            # Rotation matrix to axis angle.
            "rotation": cv2.Rodrigues(predictions["rotation"][idx])[0].tolist(),
            "translation": predictions["translation"][idx].tolist(),
            "centerpoint": predictions["centerpoint"][idx].tolist(),
        }
        shapes.append(shape)

    anno = {
        "shapes": shapes,
        "comment": "inferred and refined predictions",
        "depthPoseCheckpoint": checkpoint_name,
        "imagePath": image_name,
        "imageData": None,
        "imageWidth": dataset.parser.img_infos[img_id]["width"],
        "imageHeight": dataset.parser.img_infos[img_id]["height"],
        "K": K.tolist(),
        "ortho_scale": ortho_scale.tolist(),
        "z_aug": z_aug,
    }

    with json_path.open("w") as f:
        json.dump(anno, f, indent=2)


def dump(targets, predictions, checkpoint_path: Path, name: str) -> None:
    concat_targets = {}
    for key in targets[0].keys():
        concat_targets[key] = np.concatenate([target[key].cpu().numpy() for target in targets])
    concat_preds = np.vstack(predictions)
    np.save(checkpoint_path.parent / f"{name}_targets.npy", concat_targets)
    np.save(checkpoint_path.parent / f"{name}_predictions.npy", concat_preds)


def main():
    args = parser.parse_args()
    inference(args)


if __name__ == "__main__":
    main()