pxdepth/evaluation/metrics.py

"""Depth, point-cloud, local-structure, and boundary metrics.

Raw model outputs are aligned with the same low-resolution robust affine
procedures used by the MoGe evaluation protocol. Depth-space, log-depth-space,
and disparity-space predictions are converted to positive depth before common
metrics and point-cloud reconstruction are evaluated.
"""

from typing import Dict, Literal, Tuple, Union
from numbers import Number

import cv2
import torch
import numpy as np
import utils3d

from ..utils.alignment import (
    align_affine_lstsq,
    align_depth_affine,
    align_points_scale_xyz_shift,
)
from ..utils.tools import key_average


ALIGN_MIN_VALID_PIXELS = 16


def rel_depth(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6):
    """Compute mean absolute relative depth error.

    Args:
        pred: Positive predicted depths at selected pixels, tensor ``[N]``.
        gt: Positive ground-truth depths at the same pixels, tensor ``[N]``.
        eps: Denominator stabilizer for near-zero GT values.

    Returns:
        Python float containing ``mean(abs(pred-gt)/(gt+eps))``.
    """
    rel = (torch.abs(pred - gt) / (gt + eps)).mean()
    return rel.item()


def delta1_depth(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6):
    """Compute the fraction of depth ratios below ``1.25``.

    Args:
        pred: Positive predicted depth tensor ``[N]``.
        gt: Positive ground-truth depth tensor ``[N]``.
        eps: Compatibility argument retained by the public metric API.

    Returns:
        Python float in ``[0,1]``; larger is better.
    """
    delta1 = (torch.maximum(gt / pred, pred / gt) < 1.25).float().mean()
    return delta1.item()


def rel_point(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6):
    """Compute 3D endpoint error relative to GT camera-space radius.

    Args:
        pred: Predicted camera-space points ``[N,3]``.
        gt: Corresponding ground-truth points ``[N,3]``.
        eps: Stabilizer added to each GT point radius.

    Returns:
        Python float mean relative Euclidean point error.
    """
    dist_gt = torch.norm(gt, dim=-1)
    dist_err = torch.norm(pred - gt, dim=-1)
    rel = (dist_err / (dist_gt + eps)).mean()
    return rel.item()


def delta1_point(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6):
    """Compute the MoGe point accuracy under a 25% radial tolerance.

    Args:
        pred: Predicted camera-space points ``[N,3]``.
        gt: Corresponding ground-truth points ``[N,3]``.
        eps: Compatibility argument retained by the metric API.

    Returns:
        Python float fraction whose 3D error is below 25% of the smaller
        predicted/GT camera-space radius.
    """
    dist_pred = torch.norm(pred, dim=-1)
    dist_gt = torch.norm(gt, dim=-1)
    dist_err = torch.norm(pred - gt, dim=-1)

    delta1 = (dist_err < 0.25 * torch.minimum(dist_gt, dist_pred)).float().mean()
    return delta1.item()


def rel_point_local(pred: torch.Tensor, gt: torch.Tensor, diameter: torch.Tensor):
    """Normalize local 3D endpoint error by an object's GT diameter.

    Args:
        pred: Locally aligned predicted points ``[N,3]``.
        gt: Ground-truth points ``[N,3]`` for the same region.
        diameter: Scalar tensor containing the largest GT bounding-box extent.

    Returns:
        Python float mean error divided by ``diameter``.
    """
    dist_err = torch.norm(pred - gt, dim=-1)
    rel = (dist_err / diameter).mean()
    return rel.item()


def delta1_point_local(pred: torch.Tensor, gt: torch.Tensor, diameter: torch.Tensor):
    """Compute local point accuracy at one quarter of object diameter.

    Args:
        pred: Locally aligned predicted points ``[N,3]``.
        gt: Ground-truth points ``[N,3]``.
        diameter: Scalar GT region diameter.

    Returns:
        Python float fraction with Euclidean error below ``0.25*diameter``.
    """
    dist_err = torch.norm(pred - gt, dim=-1)
    delta1 = (dist_err < 0.25 * diameter).float().mean()
    return delta1.item()


def _nan_boundary_metrics() -> Dict[str, float]:
    """Create a complete boundary metric record for invalid edge samples.

    A stable key set keeps aggregation and JSON schemas consistent.

    Returns:
        Dictionary whose boundary accuracy and Chamfer distance are both NaN.
    """
    return {
        'acc': float('nan'),
        'cd': float('nan'),
    }


def _mda_boundary_mask(gt_depth: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Extract the Canny GT depth edge mask used for boundary evaluation.

    Args:
        gt_depth: Ground-truth depth map ``[H,W]`` in meters.
        mask: Boolean valid-depth mask ``[H,W]``.

    Returns:
        Boolean edge tensor ``[H,W]`` on ``gt_depth.device``. Depth is clipped
        to ``[0.1,65]`` meters and invalid support is dilated by a 2x2 kernel
        before Canny thresholds 100/200 are applied.
    """
    depth_np = gt_depth.detach().float().cpu().numpy()
    valid_np = mask.detach().cpu().numpy().astype(bool)
    depth_np = np.nan_to_num(depth_np, nan=0.0, posinf=65.0, neginf=0.0)

    depth_gt_clamp = np.clip(depth_np, 0.1, 65.0)
    min_val = depth_gt_clamp.min()
    max_val = depth_gt_clamp.max()
    norm_depth = (depth_gt_clamp - min_val) / (max_val - min_val + 1e-5)
    norm_depth = np.clip(norm_depth, 0.0, 1.0)
    depth_uint8 = (norm_depth * 255).astype(np.uint8)

    edge = cv2.Canny(depth_uint8, 100, 200) > 0.5
    kernel = np.ones((2, 2), np.uint8)
    valid_np = cv2.dilate(1 - valid_np.astype(np.uint8), kernel, iterations=1) < 0.5
    edge = edge & valid_np
    return torch.from_numpy(edge).to(device=gt_depth.device, dtype=torch.bool)


def _as_o3d_point_cloud(points: np.ndarray):
    """Convert an XYZ NumPy array to an Open3D point cloud.

    Args:
        points: Finite point array ``float [N,3]``.

    Returns:
        ``open3d.geometry.PointCloud`` containing the supplied XYZ positions.
    """
    import open3d as o3d

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    return pcd


def boundary_edge_metrics(
    pred_depth: torch.Tensor,
    gt_depth: torch.Tensor,
    mask: torch.Tensor,
    intrinsics: torch.Tensor,
    return_misc: bool = False,
    edge_mode: Literal['mda'] = 'mda',
) -> Union[Dict[str, float], Tuple[Dict[str, float], Dict[str, torch.Tensor]]]:
    """Evaluate edge depth and 3D boundary point-cloud quality.

    GT Canny edges select the boundary point clouds. The predicted cloud is
    rigidly refined to GT with point-to-point ICP, then bidirectional
    nearest-neighbor distances produce accuracy and symmetric Chamfer distance
    in millimeters.

    Args:
        pred_depth: Globally aligned predicted depth ``[H,W]`` in meters.
        gt_depth: Ground-truth depth ``[H,W]`` in meters.
        mask: Boolean GT valid-depth mask ``[H,W]``.
        intrinsics: Normalized camera matrix ``[3,3]``.
        return_misc: Also return edge masks, aligned clouds, and ICP transform.
        edge_mode: Boundary extraction protocol. The release supports ``'mda'``.

    Returns:
        metrics: Dictionary containing ``acc`` and ``cd`` in millimeters.
        misc: Returned only when requested. Contains ``edge_mask`` ``[H,W]``,
            edge point arrays ``[N,3]``, and ``icp_transform`` ``[4,4]``.
    """
    from scipy.spatial import cKDTree as KDTree
    import open3d as o3d

    metrics = _nan_boundary_metrics()
    misc: Dict[str, torch.Tensor] = {}

    def _finish():
        """Package the current metric state according to ``return_misc``.

        The closure captures the partially populated dictionaries by reference.

        Returns:
            Metrics dictionary alone, or ``(metrics,misc)`` when requested.
        """
        return (metrics, misc) if return_misc else metrics

    valid = mask & torch.isfinite(gt_depth) & (gt_depth > 0)
    pred_valid = torch.isfinite(pred_depth) & (pred_depth > 0)
    if edge_mode == 'mda':
        edge = _mda_boundary_mask(gt_depth, valid)
    else:
        raise ValueError(f"Unknown boundary edge mode: {edge_mode}")
    gt_edge_mask = edge & valid
    pred_edge_mask = gt_edge_mask & pred_valid
    if return_misc:
        misc['edge_mask'] = edge
    if gt_edge_mask.sum().item() < 10 or pred_edge_mask.sum().item() < 10:
        return _finish()

    pred_depth_clean = pred_depth.float().clone()
    pred_depth_clean[~pred_valid] = 1.0
    gt_depth_clean = gt_depth.float().clone()
    gt_depth_clean[~valid] = 1.0

    pred_points_full = utils3d.pt.depth_map_to_point_map(pred_depth_clean, intrinsics=intrinsics)
    gt_points_full = utils3d.pt.depth_map_to_point_map(gt_depth_clean, intrinsics=intrinsics)

    pred_points = pred_points_full[pred_edge_mask].detach().float().cpu().numpy()
    gt_points = gt_points_full[gt_edge_mask].detach().float().cpu().numpy()
    pred_points = pred_points[np.isfinite(pred_points).all(axis=1)]
    gt_points = gt_points[np.isfinite(gt_points).all(axis=1)]
    if pred_points.shape[0] < 10 or gt_points.shape[0] < 10:
        return _finish()

    pcd = _as_o3d_point_cloud(pred_points)
    pcd_gt = _as_o3d_point_cloud(gt_points)
    reg_p2p = o3d.pipelines.registration.registration_icp(
        pcd,
        pcd_gt,
        0.1,
        np.eye(4),
        o3d.pipelines.registration.TransformationEstimationPointToPoint(),
    )
    transform = reg_p2p.transformation
    pcd.transform(transform)
    pred_points_aligned = np.asarray(pcd.points)

    gt_tree = KDTree(gt_points)
    acc_distances, _ = gt_tree.query(pred_points_aligned, workers=-1)
    pred_tree = KDTree(pred_points_aligned)
    comp_distances, _ = pred_tree.query(gt_points, workers=-1)

    acc = float(np.mean(acc_distances))
    comp = float(np.mean(comp_distances))
    cd = (acc + comp) / 2.0
    if np.isfinite(acc) and np.isfinite(cd):
        metrics['acc'] = acc * 1000.0
        metrics['cd'] = cd * 1000.0

    if return_misc:
        misc['pred_edge_points'] = torch.from_numpy(pred_points_aligned).to(device=gt_depth.device, dtype=torch.float32)
        misc['gt_edge_points'] = torch.from_numpy(gt_points).to(device=gt_depth.device, dtype=torch.float32)
        misc['icp_transform'] = torch.from_numpy(transform.copy()).to(device=gt_depth.device, dtype=torch.float32)
    return _finish()


def _moge_lowres_affine(
    pred: torch.Tensor,
    target: torch.Tensor,
    mask: torch.Tensor,
    weight_depth: torch.Tensor,
) -> Tuple[torch.Tensor, bool]:
    """Fit MoGe-style weighted affine alignment on a 64x64 valid subset.

    Args:
        pred: Raw prediction map ``[H,W]`` in depth or log-depth space.
        target: GT target map ``[H,W]`` in the same affine space.
        mask: Boolean candidate-fit mask ``[H,W]``.
        weight_depth: Positive GT depth ``[H,W]`` used for inverse-depth weights.

    Returns:
        aligned: Full-resolution floating prediction ``[H,W]``.
        success: Boolean indicating whether finite affine parameters were found.
    """
    valid = (
        mask
        & torch.isfinite(pred)
        & torch.isfinite(target)
        & torch.isfinite(weight_depth)
        & (weight_depth > 0)
    )
    if valid.sum().item() < ALIGN_MIN_VALID_PIXELS:
        return pred.float(), False

    pred_clean = torch.where(valid, pred.float(), torch.zeros_like(pred, dtype=torch.float32))
    target_clean = torch.where(valid, target.float(), torch.zeros_like(target, dtype=torch.float32))
    weight_depth_clean = torch.where(valid, weight_depth.float(), torch.ones_like(weight_depth, dtype=torch.float32))
    try:
        pred_lr, target_lr, weight_depth_lr, mask_lr = utils3d.pt.masked_nearest_resize(
            pred_clean,
            target_clean,
            weight_depth_clean,
            mask=valid,
            size=(64, 64),
        )
        weight = mask_lr.flatten(-2, -1).float() / weight_depth_lr.flatten(-2, -1).clamp_min(1e-3)
        if (weight > 0).sum().item() < ALIGN_MIN_VALID_PIXELS:
            return pred.float(), False
        scale, shift = align_depth_affine(
            pred_lr.flatten(-2, -1),
            target_lr.flatten(-2, -1),
            weight,
        )
        scale = scale.squeeze()
        shift = shift.squeeze()
        ok = torch.isfinite(scale) & torch.isfinite(shift)
        if not bool(ok.item() if ok.ndim == 0 else ok.all().item()):
            return pred.float(), False
        return pred.float() * scale + shift, True
    except Exception:
        return pred.float(), False


def _moge_disparity_affine(
    pred_disparity: torch.Tensor,
    gt_disparity: torch.Tensor,
    mask: torch.Tensor,
) -> Tuple[torch.Tensor, bool]:
    """Fit least-squares scale and shift in disparity space.

    Args:
        pred_disparity: Raw predicted disparity ``[H,W]``.
        gt_disparity: Ground-truth reciprocal depth ``[H,W]``.
        mask: Boolean fit mask ``[H,W]``.

    Returns:
        aligned: Full-resolution disparity ``[H,W]``.
        success: Boolean indicating a finite affine fit.
    """
    valid = mask & torch.isfinite(pred_disparity) & torch.isfinite(gt_disparity) & (gt_disparity > 0)
    if valid.sum().item() < ALIGN_MIN_VALID_PIXELS:
        return pred_disparity.float(), False
    try:
        scale, shift = align_affine_lstsq(pred_disparity[valid].float(), gt_disparity[valid].float())
        ok = torch.isfinite(scale) & torch.isfinite(shift)
        if not bool(ok.item() if ok.ndim == 0 else ok.all().item()):
            return pred_disparity.float(), False
        return pred_disparity.float() * scale + shift, True
    except Exception:
        return pred_disparity.float(), False


def _moge_points_affine(
    pred_points: torch.Tensor,
    gt_points: torch.Tensor,
    mask: torch.Tensor,
) -> Tuple[torch.Tensor, bool]:
    """Fit one global scale and XYZ translation to a predicted point map.

    Args:
        pred_points: Predicted camera-space point map ``[H,W,3]``.
        gt_points: Ground-truth point map ``[H,W,3]``.
        mask: Boolean valid correspondence mask ``[H,W]``.

    Returns:
        aligned: Full-resolution point map ``[H,W,3]``.
        success: Boolean indicating whether robust alignment succeeded.
    """
    valid = mask & torch.isfinite(pred_points).all(dim=-1) & torch.isfinite(gt_points).all(dim=-1)
    if valid.sum().item() < ALIGN_MIN_VALID_PIXELS:
        return pred_points.float(), False

    pred_clean = torch.where(valid[..., None], pred_points.float(), torch.zeros_like(pred_points, dtype=torch.float32))
    gt_clean = torch.where(valid[..., None], gt_points.float(), torch.zeros_like(gt_points, dtype=torch.float32))
    try:
        pred_lr, gt_lr, mask_lr = utils3d.pt.masked_nearest_resize(
            pred_clean,
            gt_clean,
            mask=valid,
            size=(64, 64),
        )
        weight = mask_lr.flatten(-2, -1).float() / gt_lr.norm(dim=-1).flatten(-2, -1).clamp_min(1e-6)
        if (weight > 0).sum().item() < ALIGN_MIN_VALID_PIXELS:
            return pred_points.float(), False
        scale, shift = align_points_scale_xyz_shift(
            pred_lr.flatten(-3, -2),
            gt_lr.flatten(-3, -2),
            weight,
        )
        scale = scale.squeeze()
        shift = shift.squeeze()
        ok = torch.isfinite(scale) & torch.isfinite(shift).all()
        if not bool(ok.item() if ok.ndim == 0 else ok.all().item()):
            return pred_points.float(), False
        return pred_points.float() * scale + shift, True
    except Exception:
        return pred_points.float(), False


def compute_metrics(
    pred: Dict[str, torch.Tensor],
    gt: Dict[str, torch.Tensor],
    vis: bool = False,
    compute_boundary: bool = True,
) -> Tuple[Dict[str, Dict[str, Number]], Dict[str, torch.Tensor]]:
    """Align one prediction and compute all applicable benchmark metrics.

    Args:
        pred: Prediction dictionary. It may contain raw ``depth_affine_invariant``
            ``[H,W]`` plus ``depth_affine_space`` (``'depth'``/``'log'``), raw
            ``disparity_affine_invariant`` ``[H,W]``, optional point map
            ``points_affine_invariant`` ``[H,W,3]``, and predicted ``mask``
            ``[H,W]``.
        gt: Ground-truth sample containing depth/mask ``[H,W]``, point map
            ``[H,W,3]``, normalized intrinsics ``[3,3]``, metric/boundary flags,
            and optional segmentation annotations.
        vis: Include aligned depth/points and boundary visualization tensors in the
            auxiliary output.
        compute_boundary: Evaluate boundary metrics when the dataset is marked
            ``has_sharp_boundary``.

    Returns:
        metrics: Nested Python-number dictionary for depth, points, local points,
            and optional boundary quality.
        misc: Tensor dictionary containing aligned maps and optional boundary
            visualization data when ``vis=True``.
    """
    metrics = {}
    misc = {}

    mask = gt['depth_mask']
    gt_depth = gt['depth']
    gt_points = gt['points']

    valid_depth = mask & torch.isfinite(gt_depth) & (gt_depth > 0)
    pred_depth_aligned = None
    pred_points_aligned = None

    if 'depth_affine_invariant' in pred:
        raw_depth = pred['depth_affine_invariant'].float()
        fit_mask = valid_depth & torch.isfinite(raw_depth)
        affine_space = str(pred.get('depth_affine_space', 'depth')).lower()
        if affine_space == 'log':
            target_log = torch.log1p(gt_depth)
            aligned_log, ok = _moge_lowres_affine(raw_depth, target_log, fit_mask, gt_depth)
            pred_depth_aligned = torch.expm1(aligned_log if ok else raw_depth)
        elif affine_space == 'depth':
            aligned_depth, ok = _moge_lowres_affine(raw_depth, gt_depth, fit_mask, gt_depth)
            pred_depth_aligned = aligned_depth if ok else raw_depth
        else:
            raise ValueError(f"Unsupported depth_affine_space={affine_space!r}")

        metric_mask = fit_mask
        if metric_mask.any():
            metrics['depth_affine_invariant'] = {
                'rel': rel_depth(pred_depth_aligned[metric_mask], gt_depth[metric_mask]),
                'delta1': delta1_depth(pred_depth_aligned[metric_mask], gt_depth[metric_mask]),
            }

    elif 'disparity_affine_invariant' in pred:
        raw_disparity = pred['disparity_affine_invariant'].float()
        fit_mask = valid_depth & torch.isfinite(raw_disparity)
        gt_disparity = torch.where(valid_depth, gt_depth.reciprocal(), torch.zeros_like(gt_depth))
        aligned_disparity, ok = _moge_disparity_affine(raw_disparity, gt_disparity, fit_mask)
        aligned_disparity = aligned_disparity if ok else raw_disparity
        if fit_mask.any():
            max_depth = gt_depth[fit_mask].max()
            pred_depth_metric = aligned_disparity.clamp_min(max_depth.reciprocal()).reciprocal()
        else:
            pred_depth_metric = aligned_disparity.clamp_min(1e-6).reciprocal()
        pred_depth_aligned = pred_depth_metric
        metric_mask = fit_mask & torch.isfinite(pred_depth_metric)
        if metric_mask.any():
            metrics['depth_affine_invariant'] = {
                'rel': rel_depth(pred_depth_metric[metric_mask], gt_depth[metric_mask]),
                'delta1': delta1_depth(pred_depth_metric[metric_mask], gt_depth[metric_mask]),
            }

    pred_points_affine_invariant = pred.get('points_affine_invariant', None)
    if pred_points_affine_invariant is None and pred_depth_aligned is not None:
        point_intrinsics = gt['intrinsics'].to(
            device=pred_depth_aligned.device,
            dtype=pred_depth_aligned.dtype,
        )
        pred_points_affine_invariant = utils3d.pt.depth_map_to_point_map(
            pred_depth_aligned,
            intrinsics=point_intrinsics,
        )

    if pred_points_affine_invariant is not None:
        point_mask = (
            valid_depth
            & torch.isfinite(pred_points_affine_invariant).all(dim=-1)
            & torch.isfinite(gt_points).all(dim=-1)
        )
        if point_mask.any():
            aligned_points, ok = _moge_points_affine(pred_points_affine_invariant, gt_points, point_mask)
            pred_points_aligned = aligned_points if ok else pred_points_affine_invariant
            metrics['points_affine_invariant'] = {
                'rel': rel_point(pred_points_aligned[point_mask], gt_points[point_mask]),
                'delta1': delta1_point(pred_points_aligned[point_mask], gt_points[point_mask]),
            }

    # Local points
    if 'segmentation_mask' in gt and 'points' in gt and pred_points_affine_invariant is not None:
        pred_points = pred_points_affine_invariant
        gt_points = gt['points']
        segmentation_mask = gt['segmentation_mask']
        segmentation_labels = gt['segmentation_labels']
        local_points_metrics = []
        for _, seg_id in segmentation_labels.items():
            valid_mask = (
                (segmentation_mask == seg_id)
                & valid_depth
                & torch.isfinite(pred_points).all(dim=-1)
                & torch.isfinite(gt_points).all(dim=-1)
            )
            if valid_mask.sum().item() < 10:
                continue

            try:
                pred_lr, gt_lr, mask_lr = utils3d.pt.masked_nearest_resize(
                    torch.where(valid_mask[..., None], pred_points.float(), torch.zeros_like(pred_points, dtype=torch.float32)),
                    torch.where(valid_mask[..., None], gt_points.float(), torch.zeros_like(gt_points, dtype=torch.float32)),
                    mask=valid_mask,
                    size=(64, 64),
                )
                pred_points_masked = pred_lr[mask_lr]
                gt_points_masked = gt_lr[mask_lr]
                if pred_points_masked.shape[0] < 10:
                    continue
                diameter = (gt_points_masked.max(dim=0).values - gt_points_masked.min(dim=0).values).max()
                scale, shift = align_points_scale_xyz_shift(
                    pred_points_masked.unsqueeze(0),
                    gt_points_masked.unsqueeze(0),
                    diameter.clamp_min(1e-6).reciprocal().expand(1, gt_points_masked.shape[0]),
                )
                pred_points_masked = pred_points[valid_mask] * scale.squeeze() + shift.squeeze()
                gt_points_masked = gt_points[valid_mask]
            except Exception:
                pred_points_masked = pred_points[valid_mask]
                gt_points_masked = gt_points[valid_mask]
                diameter = (gt_points_masked.max(dim=0).values - gt_points_masked.min(dim=0).values).max()

            local_points_metrics.append({
                'rel': rel_point_local(pred_points_masked, gt_points_masked, diameter),
                'delta1': delta1_point_local(pred_points_masked, gt_points_masked, diameter),
            })

        metrics['local_points'] = key_average(local_points_metrics)

    # Boundary Acc/CD with the MDA/Canny edge.
    boundary_depth = pred_depth_aligned
    if compute_boundary and boundary_depth is not None and gt['has_sharp_boundary']:
        if vis:
            boundary_metrics, boundary_misc = boundary_edge_metrics(
                boundary_depth,
                gt_depth,
                mask,
                gt['intrinsics'],
                return_misc=True,
                edge_mode='mda',
            )
        else:
            boundary_metrics = boundary_edge_metrics(
                boundary_depth,
                gt_depth,
                mask,
                gt['intrinsics'],
                edge_mode='mda',
            )
            boundary_misc = {}
        metrics['boundary'] = boundary_metrics
        if vis:
            if 'edge_mask' in boundary_misc:
                misc['boundary_edge_mask'] = boundary_misc['edge_mask']
            if 'pred_edge_points' in boundary_misc:
                misc['boundary_pred_edge_points'] = boundary_misc['pred_edge_points']
            if 'gt_edge_points' in boundary_misc:
                misc['boundary_gt_edge_points'] = boundary_misc['gt_edge_points']
            if 'icp_transform' in boundary_misc:
                misc['boundary_icp_transform'] = boundary_misc['icp_transform']

    if vis:
        if pred_points_aligned is not None:
            misc['pred_points'] = pred_points_aligned
        elif pred_depth_aligned is not None:
            misc['pred_points'] = utils3d.pt.depth_map_to_point_map(pred_depth_aligned, intrinsics=gt['intrinsics'])
        if pred_depth_aligned is not None:
            misc['pred_depth'] = pred_depth_aligned

    return metrics, misc