AnyImage/src/anyimage/operators/cutout_tool/depth_calibration.py

"""Calibrate metric reference depth and the Depth Symmetry base plane."""

PLANE_VARIANCE_FLOOR_RATIO = 0.025
DEPTH_REFERENCE_PERCENTILE = 95.0
ADAPTIVE_DEPTH_TRANSPARENCY_RATIO = 0.05
DEFAULT_DEPTH_BASELINE = 8.0
DEFAULT_DEPTH_DIRECTION = (0.0, 0.0, 1.0)


def use_adaptive_depth_calibration(selection_transparent_fraction):
    """Return whether the Selection input has enough transparent context."""
    return float(selection_transparent_fraction) >= ADAPTIVE_DEPTH_TRANSPARENCY_RATIO


def metric_depth_reference(camera_depth, valid, uniform_scale):
    """Return the shared upper-percentile scene depth for both Depth shapes."""
    import numpy as np

    camera_depth = np.asarray(camera_depth, dtype=np.float32)
    valid = np.asarray(valid, dtype=bool)
    if valid.shape != camera_depth.shape:
        raise ValueError("Depth values and validity mask dimensions do not match")
    usable = valid & np.isfinite(camera_depth) & (camera_depth > 0.0)
    if not usable.any():
        raise ValueError("The depth field contains no valid metric depth")
    if not uniform_scale > 0.0:
        raise ValueError("Depth Uniform Scale must be positive")
    reference = np.percentile(
        camera_depth[usable],
        DEPTH_REFERENCE_PERCENTILE,
    )
    return float(reference) * float(uniform_scale)


def _sample_valid_bilinear(values, valid, x, y):
    import numpy as np

    height, width = values.shape
    x = np.clip(x, 0.0, max(width - 1, 0))
    y = np.clip(y, 0.0, max(height - 1, 0))
    x0 = np.floor(x).astype(np.int64)
    y0 = np.floor(y).astype(np.int64)
    x1 = np.minimum(x0 + 1, width - 1)
    y1 = np.minimum(y0 + 1, height - 1)
    fx = x - x0
    fy = y - y0
    samples = np.stack(
        (
            values[y0, x0],
            values[y0, x1],
            values[y1, x0],
            values[y1, x1],
        ),
        axis=1,
    )
    weights = np.stack(
        (
            (1.0 - fx) * (1.0 - fy),
            fx * (1.0 - fy),
            (1.0 - fx) * fy,
            fx * fy,
        ),
        axis=1,
    )
    sample_valid = np.stack(
        (
            valid[y0, x0],
            valid[y0, x1],
            valid[y1, x0],
            valid[y1, x1],
        ),
        axis=1,
    )
    weights *= sample_valid
    totals = weights.sum(axis=1)
    result = np.zeros(len(x), dtype=np.float32)
    usable = totals > 1e-8
    result[usable] = (samples[usable] * weights[usable]).sum(axis=1) / totals[usable]
    return result, usable


def _weighted_plane(coordinates, values, weights):
    import numpy as np

    total = float(np.sum(weights))
    if total <= 1e-12:
        return np.zeros(3, dtype=np.float32)
    center = np.sum(coordinates * weights[:, None], axis=0) / total
    value_center = float(np.sum(values * weights) / total)
    centered = coordinates - center
    covariance = (centered * weights[:, None]).T @ centered / total
    cross = centered.T @ (weights * (values - value_center)) / total
    eigenvalues, eigenvectors = np.linalg.eigh(covariance)
    largest = max(float(eigenvalues[-1]), 0.0)
    if largest <= 1e-12:
        slopes = np.zeros(2, dtype=np.float64)
    else:
        supported = np.maximum(eigenvalues, largest * PLANE_VARIANCE_FLOOR_RATIO)
        slopes = eigenvectors @ ((eigenvectors.T @ cross) / supported)
    intercept = value_center - float(center @ slopes)
    return np.asarray((*slopes, intercept), dtype=np.float32)


def _stable_plane(coordinates, values):
    import numpy as np

    coordinates = np.asarray(coordinates, dtype=np.float64)
    values = np.asarray(values, dtype=np.float64)
    if len(coordinates) < 2:
        return np.asarray((0.0, 0.0, np.median(values)), dtype=np.float32)
    weights = np.ones(len(values), dtype=np.float64)
    plane = np.zeros(3, dtype=np.float32)
    for _iteration in range(4):
        plane = _weighted_plane(coordinates, values, weights)
        residual = values - coordinates @ plane[:2] - plane[2]
        centered = residual - np.median(residual)
        scale = max(float(np.median(np.abs(centered))) * 1.4826, 1e-6)
        weights = np.minimum(
            1.0,
            (2.5 * scale) / np.maximum(np.abs(centered), 1e-12),
        )
    return plane


def depth_symmetry_projection_parameters(
    vertices,
    vertex_uv,
    camera_depth,
    valid,
    uniform_scale,
):
    """Fit a metric base plane and return its center depth and direction."""
    import numpy as np

    vertices = np.asarray(vertices, dtype=np.float32)
    uv = np.asarray(vertex_uv, dtype=np.float32)
    if vertices.ndim != 2 or vertices.shape[1] != 3:
        raise ValueError("Depth Symmetry vertices must have three components")
    if uv.shape != (len(vertices), 2):
        raise ValueError("Depth Symmetry UV coordinates do not match its vertices")
    camera_depth = np.asarray(camera_depth, dtype=np.float32)
    valid = np.asarray(valid, dtype=bool)
    if valid.shape != camera_depth.shape:
        raise ValueError("Depth Symmetry depth and mask dimensions do not match")
    if not uniform_scale > 0.0:
        raise ValueError("Depth Symmetry Uniform Scale must be positive")
    height, width = camera_depth.shape
    sample_x = uv[:, 0] * max(width - 1, 0)
    sample_y = (1.0 - uv[:, 1]) * max(height - 1, 0)
    calibration_valid = valid & np.isfinite(camera_depth) & (camera_depth > 0.0)
    sampled, usable = _sample_valid_bilinear(
        camera_depth,
        calibration_valid,
        sample_x,
        sample_y,
    )
    if not usable.any():
        raise ValueError("MoGe-2 found no valid Depth Symmetry samples")
    local_coordinates = vertices[:, :2]
    metric_depth = sampled[usable] * float(uniform_scale)
    plane = _stable_plane(local_coordinates[usable], metric_depth)
    reference = metric_depth_reference(
        camera_depth,
        calibration_valid,
        uniform_scale,
    )
    direction = np.asarray((plane[0], plane[1], 1.0), dtype=np.float64)
    direction /= np.linalg.norm(direction)
    return reference, tuple(float(value) for value in direction)