depth_validation.py

#!/usr/bin/env python3
"""Validate vision masks against depth/LiDAR as an independent geometric check.

A segmentation model can hallucinate a wall where there is none, or miss a
protruding pipe near a corner. A depth sensor sees the same scene in metric 3D,
so it can independently confirm or refute the vision output:

* ``fit_plane`` finds the dominant wall plane by RANSAC, so we know the scene is
  a planar wall and how well it fits;
* ``metric_scale_mm_per_pixel`` derives the true wall scale from depth, replacing
  the guessed ``--mm-per-unit`` the robot planner otherwise needs;
* the signed residual to the plane separates **protrusions** (closer than the
  wall: pipes, vents, frames) from **recesses** (farther: windows, doors);
* ``validate`` cross-checks those against the vision masks and reports the
  geometry the model missed, plus recess/window agreement.

Works with a depth image (``.npy`` float metres, or 16-bit PNG) or a dense
unprojected point cloud. numpy + scipy only. Depth can come from a LiDAR,
structured-light, or stereo source.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


def load_depth(path, unit="m"):
    path = Path(path)
    if path.suffix == ".npy":
        depth = np.load(path).astype(np.float32)
    else:
        from PIL import Image
        depth = np.asarray(Image.open(path)).astype(np.float32)
    return depth * (1.0 if unit == "m" else 0.001)


def unproject(depth, fx, fy, cx, cy):
    """Depth map -> camera-frame XYZ in metres."""
    height, width = depth.shape
    ys, xs = np.mgrid[0:height, 0:width]
    z = depth
    x = (xs - cx) * z / fx
    y = (ys - cy) * z / fy
    return np.stack([x, y, z], axis=-1).astype(np.float32)


def fit_plane(points, distance_threshold=0.02, iterations=300, seed=0, refit=True):
    """RANSAC plane fit; returns ``(normal, offset, inliers)`` with plane = n·p + d.

    The normal is oriented so the camera origin sits on the negative side
    (``offset <= 0``), which makes the residual sign meaningful: positive is
    farther than the wall (a recess), negative is closer (a protrusion).
    """
    points = np.asarray(points, dtype=np.float64).reshape(-1, 3)
    points = points[np.isfinite(points).all(axis=1)]
    if len(points) < 3:
        raise ValueError("need at least 3 finite points to fit a plane")
    rng = np.random.default_rng(seed)
    best_normal = best_offset = None
    best_inliers = np.zeros(len(points), dtype=bool)
    for _ in range(max(1, iterations)):
        sample = points[rng.choice(len(points), 3, replace=False)]
        normal = np.cross(sample[1] - sample[0], sample[2] - sample[0])
        norm = np.linalg.norm(normal)
        if norm < 1e-9:
            continue
        normal /= norm
        offset = -float(normal @ sample[0])
        inliers = np.abs(points @ normal + offset) < distance_threshold
        if inliers.sum() > best_inliers.sum():
            best_inliers, best_normal, best_offset = inliers, normal, offset
    if best_normal is None:
        raise ValueError("failed to find a plane (degenerate depth)")
    if refit:
        inlier_points = points[best_inliers]
        centroid = inlier_points.mean(axis=0)
        # Smallest eigenvector of the 3x3 covariance is the plane normal. Do NOT
        # call np.linalg.svd on the (N, 3) points directly: with full_matrices=True
        # it allocates an N x N array and explodes on large clouds.
        centered = inlier_points - centroid
        _, eigenvectors = np.linalg.eigh(centered.T @ centered)
        best_normal = eigenvectors[:, 0]
        best_offset = -float(best_normal @ centroid)
        best_inliers = np.abs(points @ best_normal + best_offset) < distance_threshold
    if best_offset > 0:
        best_normal, best_offset = -best_normal, -best_offset
    return best_normal, float(best_offset), best_inliers


def plane_residuals(points, normal, offset):
    return np.asarray(points, dtype=np.float64) @ np.asarray(normal) + offset


def metric_scale_mm_per_pixel(points, fx, fy):
    """Approximate physical millimetres per pixel on the observed surface."""
    points = np.asarray(points, dtype=np.float64).reshape(-1, 3)
    depth = points[:, 2]
    depth = depth[depth > 0]
    if depth.size == 0:
        return {"x": None, "y": None}
    return {"x": round(float(np.median(depth) * 1000.0 / fx), 3),
            "y": round(float(np.median(depth) * 1000.0 / fy), 3)}


def _mask_agreement(left, right):
    union = int((left | right).sum())
    return round(int((left & right).sum()) / union, 4) if union else 1.0


def validate(depth, intrinsics, wall_mask=None, window_mask=None, keepout_mask=None,
             plane_distance_threshold=0.02, protrusion_threshold=0.05, recess_threshold=0.05, seed=0):
    """Compare depth geometry with vision masks; returns ``(report, layers)``."""
    fx, fy, cx, cy = intrinsics
    points = unproject(depth, fx, fy, cx, cy)
    valid = np.isfinite(points).all(axis=-1) & (depth > 0)
    if valid.sum() < 3:
        raise ValueError("depth is empty or all-invalid")
    normal, offset, inliers = fit_plane(points[valid], distance_threshold=plane_distance_threshold, seed=seed)
    residual = np.full(depth.shape, np.nan, dtype=np.float32)
    residual[valid] = plane_residuals(points[valid], normal, offset)

    protrusions = valid & (residual < -protrusion_threshold)
    recesses = valid & (residual > recess_threshold)
    inlier_map = valid & (np.abs(residual) <= plane_distance_threshold)

    report = {
        "points": int(valid.sum()),
        "plane": {"normal": [round(float(v), 5) for v in normal], "offset": round(float(offset), 5),
                  "inlier_fraction": round(float(inliers.sum()) / int(valid.sum()), 4)},
        "metric_scale_mm_per_pixel": metric_scale_mm_per_pixel(points[valid], fx, fy),
        "protrusion_pixels": int(protrusions.sum()),
        "recess_pixels": int(recesses.sum()),
        "wall_fraction_on_plane": None,
        "missed_protrusions_pixels": None,
        "missed_recesses_pixels": None,
        "recess_vs_window_iou": None,
        "warnings": [],
    }
    if wall_mask is not None:
        wall = np.asarray(wall_mask, bool) & valid
        if wall.any():
            report["wall_fraction_on_plane"] = round(float(inlier_map[wall].mean()), 4)
            if report["wall_fraction_on_plane"] < 0.6:
                report["warnings"].append(
                    f"only {report['wall_fraction_on_plane']:.0%} of predicted wall pixels lie on the depth "
                    f"wall plane; the vision wall may be wrong or the scene is not planar")
    if keepout_mask is not None:
        keepout = np.asarray(keepout_mask, bool)
        report["missed_protrusions_pixels"] = int((protrusions & ~keepout).sum())
        report["missed_recesses_pixels"] = int((recesses & ~keepout).sum())
        if report["missed_protrusions_pixels"] > 0.002 * int(valid.sum()):
            report["warnings"].append("depth shows protrusions (obstacles) that the vision keep-out does not cover")
        if report["missed_recesses_pixels"] > 0.002 * int(valid.sum()):
            report["warnings"].append("depth shows recesses (windows/doors) that the vision keep-out does not cover")
    if window_mask is not None:
        report["recess_vs_window_iou"] = _mask_agreement(recesses, np.asarray(window_mask, bool))

    return report, {"residual": residual, "normal": normal, "offset": offset, "valid": valid,
                    "protrusions": protrusions, "recesses": recesses, "inliers": inlier_map}


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--depth", type=Path, required=True, help="depth map (.npy metres or 16-bit PNG)")
    parser.add_argument("--depth-unit", choices=("m", "mm"), default="m", help="unit of an image depth map")
    parser.add_argument("--intrinsics", required=True, metavar="fx,fy,cx,cy")
    parser.add_argument("--mask", type=Path, help="semantic mask PNG to derive wall/window keep-out from")
    parser.add_argument("--wall-mask", type=Path, help="binary wall mask PNG (overrides --mask)")
    parser.add_argument("--window-mask", type=Path, help="binary window mask PNG")
    parser.add_argument("--keepout-mask", type=Path, help="binary keep-out mask PNG")
    parser.add_argument("--protrusion-threshold", type=float, default=0.05, help="metres closer than the wall")
    parser.add_argument("--recess-threshold", type=float, default=0.05, help="metres farther than the wall")
    parser.add_argument("--report", type=Path, help="write the validation report JSON")
    parser.add_argument("--overlay", type=Path, help="write a residual overlay PNG")
    args = parser.parse_args()

    try:
        intrinsics = [float(v) for v in args.intrinsics.split(",")]
        if len(intrinsics) != 4:
            raise ValueError
    except ValueError:
        parser.error("--intrinsics must be four numbers: fx,fy,cx,cy")

    depth = load_depth(args.depth, args.depth_unit)
    from PIL import Image

    def load_binary(path):
        return np.asarray(Image.open(path).convert("L")) > 0

    wall_mask = window_mask = keepout_mask = None
    if args.wall_mask:
        wall_mask = load_binary(args.wall_mask)
    elif args.mask:
        semantic = np.asarray(Image.open(args.mask))
        wall_mask = np.isin(semantic, (1, 2, 3))
        window_mask = semantic == 8
    if args.window_mask:
        window_mask = load_binary(args.window_mask)
    if args.keepout_mask:
        keepout_mask = load_binary(args.keepout_mask)

    report, layers = validate(depth, intrinsics, wall_mask, window_mask, keepout_mask,
                              protrusion_threshold=args.protrusion_threshold, recess_threshold=args.recess_threshold)
    rendered = json.dumps(report, indent=2)
    print(rendered)
    if args.report:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(rendered + "\n", encoding="utf-8")
    if args.overlay:
        residual = layers["residual"]
        finite = residual[np.isfinite(residual)]
        span = float(np.percentile(np.abs(finite), 98)) if finite.size else 1.0
        normalized = np.zeros(depth.shape, dtype=np.uint8)
        clipped = np.clip(residual / max(span, 1e-6), -1, 1)
        normalized[np.isfinite(residual)] = ((clipped[np.isfinite(residual)] + 1) * 127).astype(np.uint8)
        args.overlay.parent.mkdir(parents=True, exist_ok=True)
        Image.fromarray(normalized).save(args.overlay)
        print(f"overlay: {args.overlay}")


if __name__ == "__main__":
    main()