accessibilityamodal/reconstruct.py

"""Category-aware geometry-first depth completion for accessibility surfaces.

This script is intentionally separate from inference.py. Accessibility3D is a learned
single-object generator; stair/ramp/sidewalk completion needs explicit scene
geometry. The pipeline here is:

1. Load RGB image plus optional obstacle mask, amodal stair mask, and depth map.
2. Use stair bands only for stairs; fit a continuous surface for walkways and
   ramps.
3. Extrapolate visible target depth only into the reviewed hidden region.
4. Export completed depth, confidence, point cloud, mesh, and debug overlays.

Mask convention:
- obstacle mask: white/nonzero marks the object hiding the stair.
- amodal mask: white/nonzero marks the whole stair/ramp target region, including
  the occluded part.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Iterable

import cv2
import numpy as np
from PIL import Image, ImageOps

from accessibilityamodal.geometry_analysis import build_accessibility_geometry_analysis


DEFAULT_OBSTACLE_BOX = (0.36, 0.32, 0.68, 0.88)


def parse_box(text: str) -> tuple[float, float, float, float]:
    values = tuple(float(v) for v in text.split(','))
    if len(values) != 4:
        raise argparse.ArgumentTypeError('box must be x1,y1,x2,y2')
    x1, y1, x2, y2 = values
    if x2 <= x1 or y2 <= y1:
        raise argparse.ArgumentTypeError('box must satisfy x2>x1 and y2>y1')
    return values


def read_rgb(path: str | Path, max_size: int | None = None) -> np.ndarray:
    image = ImageOps.exif_transpose(Image.open(path)).convert('RGB')
    if max_size and max(image.size) > max_size:
        scale = max_size / max(image.size)
        new_size = (round(image.size[0] * scale), round(image.size[1] * scale))
        image = image.resize(new_size, Image.Resampling.LANCZOS)
    return np.array(image)


def display_shape(path: str | Path) -> tuple[int, int]:
    """Return the canonical EXIF-normalized source HxW without resampling it."""
    image = ImageOps.exif_transpose(Image.open(path)).convert('RGB')
    return image.height, image.width


def _aspect_ratio_matches(source_shape: tuple[int, int], target_shape: tuple[int, int]) -> bool:
    source_h, source_w = source_shape
    target_h, target_w = target_shape
    return abs(source_w * target_h - target_w * source_h) <= max(source_w, target_w)


def resize_mask(mask: np.ndarray, shape: tuple[int, int]) -> np.ndarray:
    h, w = shape
    if mask.shape[:2] != (h, w):
        if not _aspect_ratio_matches(mask.shape[:2], (h, w)):
            raise ValueError(
                f'Mask/RGB raster mismatch: mask={mask.shape[:2]}, rgb={(h, w)}. '
                'Refusing to resize across incompatible aspect ratios because this usually indicates EXIF-orientation misalignment.'
            )
        mask = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
    return mask > 0


def read_mask(path: str | Path, shape: tuple[int, int]) -> np.ndarray:
    mask = np.array(ImageOps.exif_transpose(Image.open(path)).convert('L'))
    return resize_mask(mask > 127, shape)


def read_source_grid_mask(
    path: str | Path,
    source_shape: tuple[int, int],
    geometry_shape: tuple[int, int],
) -> np.ndarray:
    """Read a source-grid mask, then perform the one known shared downscale.

    The normal pipeline creates target/obstacle masks on the normalized source
    RGB grid.  Requiring that exact grid here prevents a stale same-aspect mask
    from being silently accepted merely because it can be resized.
    """
    mask = np.array(ImageOps.exif_transpose(Image.open(path)).convert('L')) > 127
    if mask.shape != source_shape:
        raise ValueError(
            f'Source-grid mask mismatch for {path}: mask={mask.shape}, source_rgb={source_shape}. '
            'Expected a mask drawn on the canonical EXIF-normalized source grid before geometry downscaling.'
        )
    return resize_mask(mask, geometry_shape)


def read_source_grid_rgb(
    path: str | Path,
    source_shape: tuple[int, int],
    geometry_shape: tuple[int, int],
) -> np.ndarray:
    """Read an RGB completion on the exact source grid, then downscale once.

    A generated completion is allowed to provide texture only when it is
    aligned with the canonical, EXIF-normalized source image.  This avoids
    silently painting a mesh from a stale or rotated completion.
    """
    image = ImageOps.exif_transpose(Image.open(path)).convert('RGB')
    if (image.height, image.width) != source_shape:
        raise ValueError(
            f'Source-grid completed RGB mismatch for {path}: '
            f'completed_rgb={(image.height, image.width)}, source_rgb={source_shape}. '
            'Expected a completion on the canonical EXIF-normalized source grid.'
        )
    target_h, target_w = geometry_shape
    if image.size != (target_w, target_h):
        image = image.resize((target_w, target_h), Image.Resampling.LANCZOS)
    return np.asarray(image)


def boxes_to_mask(boxes: Iterable[tuple[float, float, float, float]], shape: tuple[int, int]) -> np.ndarray:
    h, w = shape
    mask = np.zeros((h, w), dtype=bool)
    for x1, y1, x2, y2 in boxes:
        if max(x1, y1, x2, y2) <= 1.0:
            left, top, right, bottom = x1 * w, y1 * h, x2 * w, y2 * h
        else:
            left, top, right, bottom = x1, y1, x2, y2
        left = int(np.clip(round(left), 0, w - 1))
        right = int(np.clip(round(right), left + 1, w))
        top = int(np.clip(round(top), 0, h - 1))
        bottom = int(np.clip(round(bottom), top + 1, h))
        mask[top:bottom, left:right] = True
    return mask


def default_amodal_mask(shape: tuple[int, int]) -> np.ndarray:
    h, w = shape
    mask = np.zeros((h, w), dtype=bool)
    top = int(h * 0.08)
    mask[top:, :] = True
    return mask


def save_mask(path: Path, mask: np.ndarray) -> None:
    Image.fromarray((mask.astype(np.uint8) * 255)).save(path)


def visualize_depth(depth: np.ndarray, valid: np.ndarray) -> np.ndarray:
    vis = np.zeros_like(depth, dtype=np.float32)
    values = depth[valid & np.isfinite(depth) & (depth > 0)]
    if values.size == 0:
        return np.zeros((*depth.shape, 3), dtype=np.uint8)
    lo, hi = np.percentile(values, [2, 98])
    if hi <= lo:
        hi = lo + 1.0
    vis = np.clip((depth - lo) / (hi - lo), 0, 1)
    vis[~valid] = 0
    colored = cv2.applyColorMap((vis * 255).astype(np.uint8), cv2.COLORMAP_TURBO)
    colored[~valid] = 0
    return cv2.cvtColor(colored, cv2.COLOR_BGR2RGB)


def line_endpoints_on_mask(
    y_center: float,
    slope: float,
    mask: np.ndarray,
) -> tuple[tuple[int, int], tuple[int, int]] | None:
    h, w = mask.shape
    x_center = (w - 1) / 2.0
    xs = np.arange(w, dtype=np.float32)
    ys = np.rint(y_center + slope * (xs - x_center)).astype(np.int32)
    in_frame = (ys >= 0) & (ys < h)
    if not np.any(in_frame):
        return None
    valid_xs = xs[in_frame].astype(np.int32)
    valid_ys = ys[in_frame]
    on_mask = mask[valid_ys, valid_xs]
    if np.any(on_mask):
        valid_xs = valid_xs[on_mask]
        valid_ys = valid_ys[on_mask]
    left_index = int(np.argmin(valid_xs))
    right_index = int(np.argmax(valid_xs))
    return (
        (int(valid_xs[left_index]), int(valid_ys[left_index])),
        (int(valid_xs[right_index]), int(valid_ys[right_index])),
    )


def save_overlay(
    path: Path,
    rgb: np.ndarray,
    obstacle_mask: np.ndarray,
    amodal_mask: np.ndarray,
    edges_y: list[int],
    edge_slope: float = 0.0,
) -> None:
    overlay = rgb.copy()
    overlay[amodal_mask] = (0.65 * overlay[amodal_mask] + 0.35 * np.array([0, 160, 255])).astype(np.uint8)
    overlay[obstacle_mask] = (0.45 * overlay[obstacle_mask] + 0.55 * np.array([255, 60, 20])).astype(np.uint8)
    for y in edges_y:
        y = int(np.clip(y, 0, overlay.shape[0] - 1))
        endpoints = line_endpoints_on_mask(float(y), edge_slope, amodal_mask)
        if endpoints is None:
            endpoints = ((0, y), (overlay.shape[1] - 1, y))
        cv2.line(overlay, endpoints[0], endpoints[1], (80, 255, 80), 2)
    Image.fromarray(overlay).save(path)


def read_depth(path: str | Path, shape: tuple[int, int], depth_scale: float) -> np.ndarray:
    path = Path(path)
    if path.suffix.lower() == '.npy':
        depth = np.load(path).astype(np.float32)
    elif path.suffix.lower() == '.npz':
        data = np.load(path)
        depth = data[data.files[0]].astype(np.float32)
    else:
        raw = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
        if raw is None:
            raise FileNotFoundError(path)
        if raw.ndim == 3:
            raw = cv2.cvtColor(raw, cv2.COLOR_BGR2GRAY)
        depth = raw.astype(np.float32)
    h, w = shape
    if depth.shape[:2] != (h, w):
        if not _aspect_ratio_matches(depth.shape[:2], (h, w)):
            raise ValueError(
                f'Depth/RGB raster mismatch: depth={depth.shape[:2]}, rgb={(h, w)}. '
                'Refusing to resize across incompatible aspect ratios because this usually indicates orientation or source mismatch.'
            )
        depth = cv2.resize(depth, (w, h), interpolation=cv2.INTER_LINEAR)
    depth *= depth_scale
    return depth


def detect_stair_edge_model(
    rgb: np.ndarray,
    roi: np.ndarray,
    max_edges: int,
    min_gap_ratio: float = 0.025,
    min_line_ratio: float = 0.10,
    maximum_angle_degrees: float = 25.0,
    hough_threshold_ratio: float = 0.02,
    max_line_gap_ratio: float = 0.04,
) -> tuple[list[int], float]:
    h, w = roi.shape
    x_center = (w - 1) / 2.0
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    median = float(np.median(gray[roi])) if np.any(roi) else float(np.median(gray))
    lower = int(max(20, 0.66 * median))
    upper = int(min(220, 1.33 * median + 20))
    edges = cv2.Canny(gray, lower, upper)
    edges[~roi] = 0
    min_len = max(25, int(w * min_line_ratio))
    lines = cv2.HoughLinesP(
        edges,
        rho=1,
        theta=np.pi / 180,
        threshold=max(20, int(w * hough_threshold_ratio)),
        minLineLength=min_len,
        maxLineGap=max(12, int(w * max_line_gap_ratio)),
    )
    if lines is None:
        return [], 0.0

    candidates: list[tuple[float, float, float]] = []
    for line in lines[:, 0, :]:
        x1, y1, x2, y2 = line.astype(float)
        dx = x2 - x1
        dy = y2 - y1
        length = float(np.hypot(dx, dy))
        if length < min_len:
            continue
        if abs(dx) < 1e-6:
            continue
        angle = abs(np.degrees(np.arctan2(dy, dx)))
        angle = min(angle, abs(180 - angle))
        if angle > maximum_angle_degrees:
            continue
        slope = dy / dx
        y_at_center = y1 + slope * (x_center - x1)
        if not (0.06 * h <= y_at_center <= 0.97 * h):
            continue
        candidates.append((y_at_center, slope, length))

    if not candidates:
        return [], 0.0

    candidates.sort(key=lambda item: item[0])
    min_gap = max(12, int(h * min_gap_ratio))
    clusters: list[list[tuple[float, float, float]]] = []
    for item in candidates:
        if not clusters or item[0] - clusters[-1][-1][0] > min_gap:
            clusters.append([item])
        else:
            clusters[-1].append(item)

    merged: list[tuple[int, float, float]] = []
    for cluster in clusters:
        weights = np.array([item[2] for item in cluster], dtype=np.float32)
        ys = np.array([item[0] for item in cluster], dtype=np.float32)
        slopes = np.array([item[1] for item in cluster], dtype=np.float32)
        merged.append((
            int(round(float(np.average(ys, weights=weights)))),
            float(np.average(slopes, weights=weights)),
            float(weights.sum()),
        ))

    merged.sort(key=lambda item: item[2], reverse=True)
    selected = sorted(merged[:max_edges], key=lambda item: item[0])
    if not selected:
        return [], 0.0
    weights = np.array([item[2] for item in selected], dtype=np.float32)
    slopes = np.array([item[1] for item in selected], dtype=np.float32)
    maximum_slope = float(np.tan(np.radians(maximum_angle_degrees)))
    dominant_slope = float(np.clip(np.average(slopes, weights=weights), -maximum_slope, maximum_slope))
    return [int(item[0]) for item in selected], dominant_slope


def detect_horizontal_edges(
    rgb: np.ndarray,
    roi: np.ndarray,
    max_edges: int,
    min_gap_ratio: float = 0.025,
    min_line_ratio: float = 0.10,
    maximum_angle_degrees: float = 25.0,
    hough_threshold_ratio: float = 0.02,
    max_line_gap_ratio: float = 0.04,
) -> list[int]:
    edges_y, _ = detect_stair_edge_model(
        rgb,
        roi,
        max_edges,
        min_gap_ratio=min_gap_ratio,
        min_line_ratio=min_line_ratio,
        maximum_angle_degrees=maximum_angle_degrees,
        hough_threshold_ratio=hough_threshold_ratio,
        max_line_gap_ratio=max_line_gap_ratio,
    )
    return edges_y


def detect_horizontal_gradient_peaks(
    rgb: np.ndarray,
    roi: np.ndarray,
    max_edges: int,
    min_gap_ratio: float = 0.018,
) -> list[int]:
    """Recover stair tread rows when short perspective lines defeat Hough voting.

    This detector is deliberately a secondary fallback: it aggregates vertical
    image gradients only inside the reviewed visible stair mask and is used only
    when the line detector found too few edges.
    """
    h, w = roi.shape
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    vertical_gradient = np.abs(cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3))
    scores = np.zeros(h, dtype=np.float32)
    widths = roi.sum(axis=1)
    minimum_row_pixels = max(8, int(w * 0.012))
    for y in range(h):
        values = vertical_gradient[y, roi[y]]
        if values.size < minimum_row_pixels:
            continue
        strongest_count = max(1, int(values.size * 0.30))
        scores[y] = float(np.partition(values, -strongest_count)[-strongest_count:].mean())
    scores = cv2.GaussianBlur(scores[:, None], (1, 7), 0)[:, 0]
    eligible = widths >= minimum_row_pixels
    valid_scores = scores[eligible]
    if valid_scores.size == 0 or float(valid_scores.max()) <= 0:
        return []
    threshold = max(
        float(np.percentile(valid_scores, 60)),
        float(np.median(valid_scores) + 0.25 * np.std(valid_scores)),
    )
    candidates = [
        y
        for y in range(2, h - 2)
        if eligible[y]
        and scores[y] >= threshold
        and scores[y] >= float(scores[y - 2 : y + 3].max())
    ]
    minimum_gap = max(8, int(h * min_gap_ratio))
    selected: list[int] = []
    for y in sorted(candidates, key=lambda row: float(scores[row]), reverse=True):
        if all(abs(y - other) > minimum_gap for other in selected):
            selected.append(y)
        if len(selected) >= max_edges:
            break
    return sorted(selected)


def detect_depth_gradient_stair_edges(
    depth: np.ndarray,
    roi: np.ndarray,
    max_edges: int,
    min_gap_ratio: float = 0.020,
) -> list[int]:
    """Detect stair edges from depth map vertical gradient peaks.

    This is a tertiary fallback for scenes where RGB Hough lines are
    insufficient (e.g. distant/crowded stairs). Depth maps from
    Depth Anything V2 often show clear discontinuities at step boundaries.
    """
    h, w = roi.shape
    if not np.any(roi) or not np.any(np.isfinite(depth[roi])):
        return []
    depth_f = depth.astype(np.float32).copy()
    depth_f[~np.isfinite(depth_f)] = 0.0
    depth_smooth = cv2.GaussianBlur(depth_f, (7, 7), 0)
    vertical_grad = np.abs(cv2.Sobel(depth_smooth, cv2.CV_32F, 0, 1, ksize=5))
    scores = np.zeros(h, dtype=np.float32)
    widths = roi.sum(axis=1)
    min_row_px = max(6, int(w * 0.01))
    for y in range(h):
        vals = vertical_grad[y, roi[y]]
        if vals.size < min_row_px:
            continue
        top_count = max(1, int(vals.size * 0.25))
        scores[y] = float(np.partition(vals, -top_count)[-top_count:].mean())
    scores = cv2.GaussianBlur(scores[:, None], (1, 9), 0)[:, 0]
    eligible = widths >= min_row_px
    valid_scores = scores[eligible]
    if valid_scores.size == 0 or float(valid_scores.max()) <= 0:
        return []
    threshold = max(
        float(np.percentile(valid_scores, 65)),
        float(np.median(valid_scores) + 0.3 * np.std(valid_scores)),
    )
    candidates = [
        y for y in range(2, h - 2)
        if eligible[y]
        and scores[y] >= threshold
        and scores[y] >= float(scores[y - 2 : y + 3].max())
    ]
    min_gap = max(8, int(h * min_gap_ratio))
    selected: list[int] = []
    for y in sorted(candidates, key=lambda row: float(scores[row]), reverse=True):
        if all(abs(y - other) > min_gap for other in selected):
            selected.append(y)
        if len(selected) >= max_edges:
            break
    return sorted(selected)


def stair_edge_coverage_ratio(edges_y: list[int], mask: np.ndarray) -> float:
    rows = np.where(mask)[0]
    if len(edges_y) < 2 or rows.size == 0:
        return 0.0
    target_span = max(int(rows.max() - rows.min()), 1)
    return float((max(edges_y) - min(edges_y)) / target_span)


def fallback_edges(mask: np.ndarray, count: int) -> list[int]:
    ys = np.where(mask)[0]
    h = mask.shape[0]
    if ys.size == 0:
        top, bottom = int(0.15 * h), int(0.92 * h)
    else:
        top, bottom = int(np.percentile(ys, 8)), int(np.percentile(ys, 96))
    if bottom <= top:
        return []
    return [int(round(v)) for v in np.linspace(top, bottom, count + 2)[1:-1]]


def merge_edge_positions(edges: Iterable[int], min_gap: int) -> list[int]:
    values = sorted(int(v) for v in edges)
    if not values:
        return []
    clusters: list[list[int]] = []
    for value in values:
        if not clusters or value - clusters[-1][-1] > min_gap:
            clusters.append([value])
        else:
            clusters[-1].append(value)
    return [int(round(float(np.mean(cluster)))) for cluster in clusters]


def make_step_prior_depth(shape: tuple[int, int], edges_y: list[int], near: float, far: float) -> np.ndarray:
    h, w = shape
    yy = np.linspace(0, 1, h, dtype=np.float32)[:, None]
    base = near + (1.0 - yy) * (far - near)
    depth = np.repeat(base, w, axis=1)

    boundaries = [0] + sorted(int(y) for y in edges_y if 0 < y < h - 1) + [h]
    if len(boundaries) > 2:
        n_bands = len(boundaries) - 1
        for band_idx, (top, bottom) in enumerate(zip(boundaries[:-1], boundaries[1:])):
            band_height = max(1, bottom - top)
            local = np.linspace(0, 1, band_height, dtype=np.float32)[:, None]
            # Top bands are farther away; lower bands are closer. Keeping a
            # slight within-band slope avoids fully flat cardboard strips.
            band_far = far - (far - near) * band_idx / n_bands
            band_near = far - (far - near) * (band_idx + 0.72) / n_bands
            depth[top:bottom, :] = band_far + local * (band_near - band_far)
    return depth.astype(np.float32)


def make_continuous_prior_depth(shape: tuple[int, int], near: float, far: float) -> np.ndarray:
    h, w = shape
    yy = np.linspace(0.0, 1.0, h, dtype=np.float32)[:, None]
    depth = far - yy * (far - near)
    return np.repeat(depth, w, axis=1).astype(np.float32)


def camera_intrinsics(w: int, h: int, fx: float | None, fy: float | None, cx: float | None, cy: float | None):
    default_f = float(max(w, h))
    fx = default_f if fx is None else fx
    fy = default_f if fy is None else fy
    cx = (w - 1) / 2.0 if cx is None else cx
    cy = (h - 1) / 2.0 if cy is None else cy
    return fx, fy, cx, cy


def pixels_to_points(depth: np.ndarray, fx: float, fy: float, cx: float, cy: float):
    h, w = depth.shape
    xs, ys = np.meshgrid(np.arange(w, dtype=np.float32), np.arange(h, dtype=np.float32))
    z = depth.astype(np.float32)
    x = (xs - cx) * z / fx
    y = (ys - cy) * z / fy
    return np.stack([x, y, z], axis=-1)


def fit_plane(points: np.ndarray):
    if points.shape[0] < 50:
        return None
    centroid = points.mean(axis=0)
    centered = points - centroid
    try:
        _, _, vh = np.linalg.svd(centered, full_matrices=False)
    except np.linalg.LinAlgError:
        return None
    normal = vh[-1]
    norm = np.linalg.norm(normal)
    if norm < 1e-6:
        return None
    normal = normal / norm
    d = -float(np.dot(normal, centroid))
    return normal.astype(np.float32), d


def intersect_plane_for_pixels(shape: tuple[int, int], plane, fx: float, fy: float, cx: float, cy: float) -> np.ndarray:
    h, w = shape
    normal, d = plane
    xs, ys = np.meshgrid(np.arange(w, dtype=np.float32), np.arange(h, dtype=np.float32))
    rays = np.stack([(xs - cx) / fx, (ys - cy) / fy, np.ones((h, w), dtype=np.float32)], axis=-1)
    denom = rays @ normal
    with np.errstate(divide='ignore', invalid='ignore'):
        t = -d / denom
    t[~np.isfinite(t)] = 0
    t[t <= 0] = 0
    return t.astype(np.float32)


def inpaint_depth(depth: np.ndarray, fill_mask: np.ndarray, valid: np.ndarray) -> np.ndarray:
    values = depth[valid & np.isfinite(depth) & (depth > 0)]
    if values.size == 0:
        return depth.copy()
    lo, hi = np.percentile(values, [1, 99])
    if hi <= lo:
        hi = lo + 1.0
    normalized = np.clip((depth - lo) / (hi - lo), 0, 1)
    normalized[~np.isfinite(normalized)] = float(np.median(normalized[valid]))
    filled = cv2.inpaint((normalized * 255).astype(np.uint8), fill_mask.astype(np.uint8) * 255, 5, cv2.INPAINT_TELEA)
    return filled.astype(np.float32) / 255.0 * (hi - lo) + lo


def inpaint_rgb(rgb: np.ndarray, fill_mask: np.ndarray, radius: float = 5.0) -> np.ndarray:
    if not np.any(fill_mask):
        return rgb.copy()
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    inpainted = cv2.inpaint(bgr, fill_mask.astype(np.uint8) * 255, radius, cv2.INPAINT_TELEA)
    return cv2.cvtColor(inpainted, cv2.COLOR_BGR2RGB)


def stair_band_mask(
    shape: tuple[int, int],
    top: float,
    bottom: float,
    stair_edge_slope: float = 0.0,
) -> np.ndarray:
    """Return one tread/riser band aligned with the observed image perspective.

    ``edges_y`` are measured at the image centre.  A non-zero slope means that
    the same edge moves vertically across the frame, so horizontal image rows
    would mix different physical steps.  This helper preserves the legacy
    horizontal behaviour when the slope is zero.
    """
    h, w = shape
    ys, xs = np.indices((h, w), dtype=np.float32)
    center_x = (w - 1) * 0.5
    centerline_y = ys - float(stair_edge_slope) * (xs - center_x)
    return (centerline_y >= float(top)) & (centerline_y < float(bottom))


def complete_depth_by_planes(
    depth: np.ndarray,
    obstacle_mask: np.ndarray,
    amodal_mask: np.ndarray,
    edges_y: list[int],
    fx: float,
    fy: float,
    cx: float,
    cy: float,
    min_fit_points: int,
    stair_edge_slope: float = 0.0,
):
    h, w = depth.shape
    completed = depth.copy()
    visible = amodal_mask & ~obstacle_mask & np.isfinite(depth) & (depth > 0)
    inpainted = inpaint_depth(depth, obstacle_mask & amodal_mask, visible)
    points = pixels_to_points(depth, fx, fy, cx, cy)
    confidence = np.zeros(depth.shape, dtype=np.float32)
    distance_to_visible = cv2.distanceTransform((~visible).astype(np.uint8), cv2.DIST_L2, 5)
    plane_blend_scale = max(8.0, 0.025 * float(np.hypot(*depth.shape)))
    plane_weight = np.clip(distance_to_visible / plane_blend_scale, 0.0, 0.85)

    boundaries = [0] + sorted(int(y) for y in edges_y if 0 < y < h - 1) + [h]
    fitted_planes = []
    for top, bottom in zip(boundaries[:-1], boundaries[1:]):
        band = stair_band_mask((h, w), top, bottom, stair_edge_slope)
        fit_mask = band & visible
        fill_mask = band & obstacle_mask & amodal_mask
        plane = None
        if int(fit_mask.sum()) >= min_fit_points:
            sample = points[fit_mask]
            if sample.shape[0] > 30000:
                sample = sample[np.linspace(0, sample.shape[0] - 1, 30000).astype(np.int64)]
            plane = fit_plane(sample)
        if plane is not None:
            plane_depth = intersect_plane_for_pixels((h, w), plane, fx, fy, cx, cy)
            ok = fill_mask & (plane_depth > 0) & np.isfinite(plane_depth)
            completed[ok] = (
                plane_weight[ok] * plane_depth[ok]
                + (1.0 - plane_weight[ok]) * inpainted[ok]
            )
            missing = fill_mask & ~ok
            completed[missing] = inpainted[missing]
            residual = np.abs(sample @ plane[0] + plane[1])
            scale = max(float(np.median(sample[:, 2])), 1e-6)
            quality = float(np.exp(-12.0 * np.median(residual) / scale))
            confidence[ok] = np.clip(quality, 0.35, 0.95)
            confidence[missing] = 0.25
            mode = 'plane_diagonal' if abs(stair_edge_slope) > 1e-6 else 'plane'
            fitted_planes.append((top, bottom, int(fit_mask.sum()), mode))
        else:
            completed[fill_mask] = inpainted[fill_mask]
            confidence[fill_mask] = 0.20
            mode = 'inpaint_fallback_diagonal' if abs(stair_edge_slope) > 1e-6 else 'inpaint_fallback'
            fitted_planes.append((top, bottom, int(fit_mask.sum()), mode))

    return completed, fitted_planes, confidence


def robust_fit_plane(points: np.ndarray, min_points: int):
    if points.shape[0] < min_points:
        return None, 0.0, 0
    sample = points
    if sample.shape[0] > 50000:
        sample = sample[np.linspace(0, sample.shape[0] - 1, 50000).astype(np.int64)]
    plane = None
    for _ in range(4):
        plane = fit_plane(sample)
        if plane is None:
            return None, 0.0, int(sample.shape[0])
        residual = np.abs(sample @ plane[0] + plane[1])
        median = float(np.median(residual))
        mad = float(np.median(np.abs(residual - median)))
        threshold = max(median + 2.5 * max(mad, 1e-6), float(np.percentile(residual, 70)))
        retained = sample[residual <= threshold]
        if retained.shape[0] < min_points or retained.shape[0] >= sample.shape[0] * 0.98:
            break
        sample = retained
    if plane is None:
        return None, 0.0, int(sample.shape[0])
    residual = np.abs(sample @ plane[0] + plane[1])
    scale = max(float(np.median(sample[:, 2])), 1e-6)
    quality = float(np.exp(-10.0 * float(np.median(residual)) / scale))
    return plane, float(np.clip(quality, 0.0, 1.0)), int(sample.shape[0])


def complete_depth_continuous_surface(
    depth: np.ndarray,
    completion_mask: np.ndarray,
    amodal_mask: np.ndarray,
    fx: float,
    fy: float,
    cx: float,
    cy: float,
    min_fit_points: int,
):
    completed = depth.copy()
    confidence = np.zeros(depth.shape, dtype=np.float32)
    visible = amodal_mask & ~completion_mask & np.isfinite(depth) & (depth > 0)
    inpainted = inpaint_depth(depth, completion_mask, visible)
    points = pixels_to_points(depth, fx, fy, cx, cy)
    plane, quality, sample_count = robust_fit_plane(points[visible], min_fit_points)
    method = 'edge_aware_inpaint'
    if plane is None:
        completed[completion_mask] = inpainted[completion_mask]
        distance_to_visible = cv2.distanceTransform((~visible).astype(np.uint8), cv2.DIST_L2, 5)
        distance_scale = max(8.0, 0.06 * float(np.hypot(*depth.shape)))
        distance_confidence = np.exp(-distance_to_visible / distance_scale)
        confidence[completion_mask] = np.clip(
            0.20 + 0.25 * distance_confidence[completion_mask],
            0.20,
            0.45,
        )
        return completed, [(0, depth.shape[0], sample_count, method)], confidence

    plane_depth = intersect_plane_for_pixels(depth.shape, plane, fx, fy, cx, cy)
    visible_values = depth[visible]
    low, high = np.percentile(visible_values, [1, 99])
    plausible = (
        completion_mask
        & np.isfinite(plane_depth)
        & (plane_depth > max(1e-6, 0.50 * low))
        & (plane_depth < 1.50 * high)
    )
    # Keep local boundary detail from inpainting while the robust plane supplies
    # a stable global surface through larger occluders.
    distance_to_visible = cv2.distanceTransform((~visible).astype(np.uint8), cv2.DIST_L2, 5)
    plane_blend_scale = max(8.0, 0.025 * float(np.hypot(*depth.shape)))
    plane_weight = np.clip(distance_to_visible / plane_blend_scale, 0.0, 0.85)
    completed[plausible] = (
        plane_weight[plausible] * plane_depth[plausible]
        + (1.0 - plane_weight[plausible]) * inpainted[plausible]
    )
    fallback = completion_mask & ~plausible
    completed[fallback] = inpainted[fallback]

    distance_scale = max(8.0, 0.06 * float(np.hypot(*depth.shape)))
    distance_confidence = np.exp(-distance_to_visible / distance_scale)
    confidence[plausible] = np.clip(quality * distance_confidence[plausible], 0.25, 0.95)
    confidence[fallback] = np.clip(
        0.20 + 0.35 * distance_confidence[fallback],
        0.20,
        0.55,
    )
    method = 'robust_continuous_plane_plus_edge_aware_inpaint'
    return completed, [(0, depth.shape[0], sample_count, method)], confidence


def complete_depth_generic_inpaint(
    depth: np.ndarray, completion_mask: np.ndarray, amodal_mask: np.ndarray
):
    visible = amodal_mask & ~completion_mask & np.isfinite(depth) & (depth > 0)
    completed = depth.copy()
    inpainted = inpaint_depth(depth, completion_mask, visible)
    completed[completion_mask] = inpainted[completion_mask]
    confidence = np.zeros(depth.shape, dtype=np.float32)
    confidence[completion_mask] = 0.20
    return completed, [(0, depth.shape[0], int(visible.sum()), 'edge_aware_inpaint')], confidence


def enforce_hidden_boundary_continuity(
    completed: np.ndarray,
    source_depth: np.ndarray,
    visible: np.ndarray,
    hidden: np.ndarray,
    radius: int,
    anchor_power: float = 0.45,
) -> np.ndarray:
    if radius <= 0 or not np.any(hidden) or not np.any(visible):
        return completed
    size = radius * 2 + 1
    sums = cv2.boxFilter(
        np.where(visible, source_depth, 0.0).astype(np.float32),
        -1,
        (size, size),
        normalize=False,
        borderType=cv2.BORDER_CONSTANT,
    )
    counts = cv2.boxFilter(
        visible.astype(np.float32),
        -1,
        (size, size),
        normalize=False,
        borderType=cv2.BORDER_CONSTANT,
    )
    boundary = hidden & (counts > 0)
    if not np.any(boundary):
        return completed
    neighbor_depth = np.zeros_like(completed, dtype=np.float32)
    neighbor_depth[boundary] = sums[boundary] / counts[boundary]
    distance = cv2.distanceTransform((~visible).astype(np.uint8), cv2.DIST_L2, 5)
    raw_anchor = 1.0 - np.clip((distance - 1.0) / max(float(radius), 1.0), 0.0, 1.0)
    anchor_weight = np.power(raw_anchor, max(float(anchor_power), 1e-3))
    output = completed.copy()
    output[boundary] = (
        (1.0 - anchor_weight[boundary]) * completed[boundary]
        + anchor_weight[boundary] * neighbor_depth[boundary]
    )
    return output


DEPTH_PLAUSIBILITY_POLICIES = {
    # Stairs legitimately span several depth layers, so retain a wider target
    # envelope.  A plane/ray intersection must still stay inside the robust
    # numerical domain observed in the source depth map.
    'stairs': {
        'target_span_margin': 1.00,
        'maximum_correction_ratio': 0.05,
    },
    # A reviewed walkable surface or ramp should be locally continuous.  The
    # margin remains wide enough for perspective extrapolation already accepted
    # by complete_depth_continuous_surface.
    'walkable': {
        'target_span_margin': 0.75,
        'maximum_correction_ratio': 0.02,
    },
    'ramp': {
        'target_span_margin': 0.75,
        'maximum_correction_ratio': 0.02,
    },
    # Generic mode has the weakest category prior, but it may not invent an
    # unbounded scale outside the observed depth domain.
    'generic': {
        'target_span_margin': 1.50,
        'maximum_correction_ratio': 0.05,
    },
}


def enforce_hidden_depth_plausibility(
    completed: np.ndarray,
    source_depth: np.ndarray,
    visible: np.ndarray,
    hidden: np.ndarray,
    geometry_mode: str,
) -> tuple[np.ndarray, np.ndarray, dict]:
    """Reject numerically implausible hidden depths and return an audit record.

    Plane/ray intersections can explode when a fitted plane is nearly parallel
    to a camera ray.  This guard is deliberately conservative:

    * only hidden pixels may be changed;
    * the accepted range is inferred from the observed source-depth scale, not
      assumed to be metric;
    * category-specific target margins allow layered stairs more variation than
      continuous walkable surfaces and ramps;
    * rejected values use bounded edge-aware inpainting instead of clipping a
      singular plane to a hard wall.

    The returned correction mask is suitable for lowering confidence and for a
    review overlay.  A large correction ratio is marked as withheld in the
    audit rather than silently presented as a trustworthy geometry candidate.
    """
    if completed.shape != source_depth.shape:
        raise ValueError(
            f'Completed/source depth shape mismatch: {completed.shape} != {source_depth.shape}'
        )
    if visible.shape != completed.shape or hidden.shape != completed.shape:
        raise ValueError('Visible/hidden masks must share the completed-depth raster')
    if geometry_mode not in DEPTH_PLAUSIBILITY_POLICIES:
        raise ValueError(f'Unsupported geometry mode for depth plausibility: {geometry_mode!r}')

    hidden = hidden.astype(bool)
    visible = visible.astype(bool) & ~hidden
    policy = DEPTH_PLAUSIBILITY_POLICIES[geometry_mode]
    hidden_count = int(hidden.sum())
    empty_corrections = np.zeros(completed.shape, dtype=bool)
    if hidden_count == 0:
        return completed.copy(), empty_corrections, {
            'policy': f'{geometry_mode}_observed_depth_domain_v1',
            'status': 'not_applicable_empty_hidden_region',
            'candidate_eligible_for_review': True,
            'hidden_pixel_count': 0,
            'corrected_pixel_count': 0,
            'corrected_ratio': 0.0,
            'maximum_correction_ratio': float(policy['maximum_correction_ratio']),
            'source_depth_is_treated_as_metric_truth': False,
        }

    source_valid = np.isfinite(source_depth) & (source_depth > 0)
    source_values = source_depth[source_valid]
    visible_valid = visible & source_valid
    visible_values = source_depth[visible_valid]
    if source_values.size == 0 or visible_values.size == 0:
        invalid = hidden & (
            ~np.isfinite(completed)
            | (completed <= 0)
        )
        return completed.copy(), invalid, {
            'policy': f'{geometry_mode}_observed_depth_domain_v1',
            'status': 'withheld_no_visible_depth_support',
            'candidate_eligible_for_review': False,
            'hidden_pixel_count': hidden_count,
            'corrected_pixel_count': 0,
            'corrected_ratio': 0.0,
            'maximum_correction_ratio': float(policy['maximum_correction_ratio']),
            'source_depth_is_treated_as_metric_truth': False,
            'reason': 'No positive finite visible target depth was available for a bounded fallback.',
        }

    source_q001, source_q999 = np.percentile(source_values, [0.1, 99.9])
    source_span = max(
        float(source_q999 - source_q001),
        0.05 * abs(float(np.median(source_values))),
        1e-6,
    )
    # Ignore isolated source-map outliers while never extending beyond the
    # finite positive range actually emitted by the depth estimator.
    source_lower = max(
        float(source_values.min()),
        float(source_q001 - 0.10 * source_span),
    )
    source_upper = min(
        float(source_values.max()),
        float(source_q999 + 0.10 * source_span),
    )

    target_q01, target_q99 = np.percentile(visible_values, [1.0, 99.0])
    target_span = max(
        float(target_q99 - target_q01),
        0.05 * abs(float(np.median(visible_values))),
        1e-6,
    )
    target_margin = float(policy['target_span_margin']) * target_span
    accepted_lower = max(source_lower, float(target_q01 - target_margin))
    accepted_upper = min(source_upper, float(target_q99 + target_margin))
    if not np.isfinite(accepted_lower) or not np.isfinite(accepted_upper) or accepted_upper < accepted_lower:
        accepted_lower, accepted_upper = source_lower, source_upper

    corrections = hidden & (
        ~np.isfinite(completed)
        | (completed < accepted_lower)
        | (completed > accepted_upper)
    )
    corrected_count = int(corrections.sum())
    output = completed.copy()
    if corrected_count:
        fallback = inpaint_depth(source_depth, hidden, visible_valid)
        fallback = np.nan_to_num(
            fallback,
            nan=float(np.median(visible_values)),
            posinf=accepted_upper,
            neginf=accepted_lower,
        )
        fallback = np.clip(fallback, accepted_lower, accepted_upper)
        output[corrections] = fallback[corrections]

    # Preserve source depth bit-for-bit outside the reviewed hidden region.
    output[~hidden] = source_depth[~hidden]
    corrected_ratio = float(corrected_count / hidden_count)
    maximum_correction_ratio = float(policy['maximum_correction_ratio'])
    candidate_eligible = bool(
        corrected_ratio <= maximum_correction_ratio
        and np.all(np.isfinite(output[hidden]))
        and np.all(output[hidden] >= accepted_lower)
        and np.all(output[hidden] <= accepted_upper)
    )
    if not candidate_eligible:
        status = 'withheld_excessive_depth_corrections'
    elif corrected_count:
        status = 'corrected_sparse_depth_outliers'
    else:
        status = 'within_observed_depth_domain'
    audit = {
        'policy': f'{geometry_mode}_observed_depth_domain_v1',
        'status': status,
        'candidate_eligible_for_review': candidate_eligible,
        'hidden_pixel_count': hidden_count,
        'corrected_pixel_count': corrected_count,
        'corrected_ratio': corrected_ratio,
        'maximum_correction_ratio': maximum_correction_ratio,
        'accepted_depth_lower': float(accepted_lower),
        'accepted_depth_upper': float(accepted_upper),
        'observed_source_depth_min': float(source_values.min()),
        'observed_source_depth_max': float(source_values.max()),
        'visible_target_depth_p01': float(target_q01),
        'visible_target_depth_p99': float(target_q99),
        'source_depth_is_treated_as_metric_truth': False,
    }
    return output, corrections, audit


def write_point_cloud_ply(path: Path, points: np.ndarray, colors: np.ndarray, mask: np.ndarray, stride: int) -> int:
    sampled = np.zeros(mask.shape, dtype=bool)
    sampled[::stride, ::stride] = True
    use = mask & sampled & np.all(np.isfinite(points), axis=-1) & (points[..., 2] > 0)
    pts = points[use]
    cols = colors[use]
    with open(path, 'w', encoding='ascii') as f:
        f.write('ply\nformat ascii 1.0\n')
        f.write(f'element vertex {len(pts)}\n')
        f.write('property float x\nproperty float y\nproperty float z\n')
        f.write('property uchar red\nproperty uchar green\nproperty uchar blue\n')
        f.write('end_header\n')
        for p, c in zip(pts, cols):
            f.write(f'{p[0]:.6f} {p[1]:.6f} {p[2]:.6f} {int(c[0])} {int(c[1])} {int(c[2])}\n')
    return int(len(pts))


def write_mesh_ply(path: Path, points: np.ndarray, colors: np.ndarray, mask: np.ndarray, stride: int, max_depth_jump: float) -> tuple[int, int]:
    h, w = mask.shape
    ys = np.arange(0, h, stride)
    xs = np.arange(0, w, stride)
    vertex_id = -np.ones((len(ys), len(xs)), dtype=np.int64)
    vertices = []
    vertex_colors = []
    for iy, y in enumerate(ys):
        for ix, x in enumerate(xs):
            if mask[y, x] and np.isfinite(points[y, x]).all() and points[y, x, 2] > 0:
                vertex_id[iy, ix] = len(vertices)
                vertices.append(points[y, x])
                vertex_colors.append(colors[y, x])

    faces = []
    for iy in range(len(ys) - 1):
        for ix in range(len(xs) - 1):
            ids = [vertex_id[iy, ix], vertex_id[iy, ix + 1], vertex_id[iy + 1, ix], vertex_id[iy + 1, ix + 1]]
            if min(ids) < 0:
                continue
            z = np.array([vertices[i][2] for i in ids], dtype=np.float32)
            if float(z.max() - z.min()) > max_depth_jump:
                continue
            faces.append((ids[0], ids[2], ids[1]))
            faces.append((ids[1], ids[2], ids[3]))

    with open(path, 'w', encoding='ascii') as f:
        f.write('ply\nformat ascii 1.0\n')
        f.write(f'element vertex {len(vertices)}\n')
        f.write('property float x\nproperty float y\nproperty float z\n')
        f.write('property uchar red\nproperty uchar green\nproperty uchar blue\n')
        f.write(f'element face {len(faces)}\n')
        f.write('property list uchar int vertex_indices\n')
        f.write('end_header\n')
        for p, c in zip(vertices, vertex_colors):
            f.write(f'{p[0]:.6f} {p[1]:.6f} {p[2]:.6f} {int(c[0])} {int(c[1])} {int(c[2])}\n')
        for face in faces:
            f.write(f'3 {face[0]} {face[1]} {face[2]}\n')
    return int(len(vertices)), int(len(faces))


CATEGORY_GEOMETRY_MODES = {
    'stairs': 'stairs',
    'ramp': 'ramp',
    'walkway': 'walkable',
    'walkable': 'walkable',
    'curb_cut': 'walkable',
    'raised_curb': 'walkable',
    'tactile_paving': 'walkable',
    'unknown': 'generic',
}


def resolve_geometry_mode(category, requested_mode):
    """Infer a safe mode from category and reject contradictory geometry priors."""
    if category is None:
        return requested_mode or 'stairs'
    expected_mode = CATEGORY_GEOMETRY_MODES[category]
    if requested_mode is not None and requested_mode != expected_mode:
        raise ValueError(
            f'Category {category!r} requires geometry mode {expected_mode!r}; '
            f'got {requested_mode!r}'
        )
    return expected_mode


def run(args):
    if not args.image:
        raise ValueError("--image is required")
    args.geometry_mode = resolve_geometry_mode(
        getattr(args, 'category', None),
        getattr(args, 'geometry_mode', None),
    )
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    source_shape = display_shape(args.image)
    rgb = read_rgb(args.image, args.max_size)
    h, w = rgb.shape[:2]
    shape = (h, w)

    if args.amodal_mask:
        amodal_mask = read_source_grid_mask(args.amodal_mask, source_shape, shape)
        amodal_source = args.amodal_mask
    else:
        amodal_mask = default_amodal_mask(shape)
        amodal_source = 'default lower-scene mask'

    if args.obstacle_mask:
        obstacle_mask = read_source_grid_mask(args.obstacle_mask, source_shape, shape)
        obstacle_source = args.obstacle_mask
    else:
        boxes = args.occlusion_box
        if not boxes and not args.no_default_obstacle:
            boxes = [DEFAULT_OBSTACLE_BOX]
        obstacle_mask = boxes_to_mask(boxes, shape) if boxes else np.zeros(shape, dtype=bool)
        obstacle_source = 'occlusion boxes' if boxes else 'empty obstacle mask'

    obstacle_mask &= amodal_mask
    if args.visible_mask:
        visible_mask = read_source_grid_mask(args.visible_mask, source_shape, shape) & amodal_mask & ~obstacle_mask
        visible_source = args.visible_mask
        completion_mask = amodal_mask & ~visible_mask
    else:
        visible_mask = amodal_mask & ~obstacle_mask
        visible_source = 'amodal mask minus obstacle mask'
        completion_mask = obstacle_mask & amodal_mask

    geometry_mode = args.geometry_mode
    reference_edges: list[int] = []
    target_edges: list[int] = []
    edges_y: list[int] = []
    reference_edge_slope = 0.0
    target_edge_slope = 0.0
    stair_edge_slope = 0.0
    edge_source = 'not_applicable_for_continuous_surface'
    used_regular_fallback_edges = False
    used_perspective_edge_expansion = False
    used_gradient_edge_expansion = False
    used_depth_gradient_edge_expansion = False
    stair_edge_confidence = 1.0 if geometry_mode != 'stairs' else 0.0
    stair_edge_coverage = 1.0 if geometry_mode != 'stairs' else 0.0
    if geometry_mode == 'stairs':
        reference_rgb = None
        if args.stair_edge_y:
            edges_y = sorted(set(int(y) for y in args.stair_edge_y if 0 < int(y) < h - 1))
            stair_edge_slope = float(np.tan(np.radians(args.stair_edge_angle_degrees)))
            edge_source = 'reviewed stair edge y positions'
            stair_edge_confidence = 1.0
            stair_edge_coverage = stair_edge_coverage_ratio(edges_y, amodal_mask)
        elif args.reference_image and Path(args.reference_image).exists():
            reference_rgb = read_rgb(args.reference_image, max_size=max(h, w))
            if reference_rgb.shape[:2] != shape:
                if not _aspect_ratio_matches(reference_rgb.shape[:2], shape):
                    raise ValueError(
                        f'Reference RGB/target raster mismatch: reference={reference_rgb.shape[:2]}, target={shape}. '
                        'Refusing to resize across incompatible aspect ratios because the reference supplies stair-edge evidence.'
                    )
                reference_rgb = cv2.resize(reference_rgb, (w, h), interpolation=cv2.INTER_AREA)
            reference_edges, reference_edge_slope = detect_stair_edge_model(
                reference_rgb,
                amodal_mask,
                args.max_step_edges,
            )

        if not args.stair_edge_y:
            target_edges, target_edge_slope = detect_stair_edge_model(rgb, visible_mask, args.max_step_edges)
            if reference_edges and len(reference_edges) >= len(target_edges):
                edges_y = reference_edges
                stair_edge_slope = reference_edge_slope
                edge_source = args.reference_image
            elif target_edges:
                edges_y = target_edges
                stair_edge_slope = target_edge_slope
                edge_source = args.image
            else:
                edge_source = 'regular fallback edges'

            initial_coverage = stair_edge_coverage_ratio(edges_y, amodal_mask)
            if initial_coverage < args.minimum_auto_stair_edge_coverage:
                relaxed_target, relaxed_target_slope = detect_stair_edge_model(
                    rgb,
                    visible_mask,
                    args.max_step_edges,
                    min_gap_ratio=0.025,
                    min_line_ratio=0.08,
                    maximum_angle_degrees=30.0,
                    hough_threshold_ratio=0.015,
                    max_line_gap_ratio=0.05,
                )
                relaxed_reference = []
                relaxed_reference_slope = 0.0
                if reference_rgb is not None:
                    relaxed_reference, relaxed_reference_slope = detect_stair_edge_model(
                        reference_rgb,
                        amodal_mask,
                        args.max_step_edges,
                        min_gap_ratio=0.025,
                        min_line_ratio=0.08,
                        maximum_angle_degrees=30.0,
                        hough_threshold_ratio=0.015,
                        max_line_gap_ratio=0.05,
                    )
                expanded = merge_edge_positions(
                    edges_y + relaxed_target + relaxed_reference,
                    max(8, int(h * 0.018)),
                )
                expanded_coverage = stair_edge_coverage_ratio(expanded, amodal_mask)
                if expanded_coverage > initial_coverage:
                    edges_y = expanded[: args.max_step_edges]
                    if abs(stair_edge_slope) < 1e-6:
                        stair_edge_slope = relaxed_target_slope or relaxed_reference_slope
                    edge_source = f'{edge_source} + perspective-tolerant expansion'
                    used_perspective_edge_expansion = True

            if len(edges_y) < args.minimum_detected_step_edges:
                gradient_edges = detect_horizontal_gradient_peaks(
                    rgb,
                    visible_mask,
                    args.max_step_edges,
                )
                expanded = merge_edge_positions(
                    edges_y + gradient_edges,
                    max(8, int(h * 0.018)),
                )
                if len(expanded) > len(edges_y):
                    edges_y = expanded[: args.max_step_edges]
                    edge_source = f'{edge_source} + visible-mask row-gradient expansion'
                    used_gradient_edge_expansion = True

            fallback = fallback_edges(amodal_mask, args.fallback_step_edges)
            if len(edges_y) < args.minimum_detected_step_edges:
                edges_y = merge_edge_positions(edges_y + fallback, max(10, int(h * 0.035)))
                edge_source = (
                    f'{edge_source} + regular fallback edges'
                    if edge_source != 'regular fallback edges'
                    else edge_source
                )
                used_regular_fallback_edges = True
                stair_edge_confidence = 0.25
            else:
                stair_edge_coverage = stair_edge_coverage_ratio(edges_y, amodal_mask)
                count_score = min(len(edges_y) / 5.0, 1.0)
                coverage_score = min(
                    stair_edge_coverage / max(args.minimum_auto_stair_edge_coverage, 1e-6),
                    1.0,
                )
                stair_edge_confidence = float(
                    np.clip(0.30 + 0.25 * count_score + 0.35 * coverage_score, 0.0, 0.90)
                )
                if used_gradient_edge_expansion:
                    stair_edge_confidence = min(stair_edge_confidence, 0.75)
            if not edges_y:
                edges_y = fallback
                used_regular_fallback_edges = bool(edges_y)
            stair_edge_coverage = stair_edge_coverage_ratio(edges_y, amodal_mask)
            if len(edges_y) < args.minimum_detected_step_edges:
                stair_edge_slope = 0.0

    if args.depth:
        depth = read_depth(args.depth, shape, args.depth_scale)
        depth_source = args.depth
    else:
        if geometry_mode == 'stairs':
            depth = make_step_prior_depth(shape, edges_y, args.synthetic_near, args.synthetic_far)
            depth_source = 'synthetic perspective stair prior'
        else:
            depth = make_continuous_prior_depth(shape, args.synthetic_near, args.synthetic_far)
            depth_source = 'synthetic continuous perspective prior'

    # --- Post-depth stair edge expansion using depth gradient ---
    # When a real depth map is available and RGB-based coverage is still low,
    # try to find additional stair edges from depth vertical gradient peaks.
    if (
        geometry_mode == 'stairs'
        and args.depth
        and not args.stair_edge_y
        and stair_edge_coverage < args.minimum_auto_stair_edge_coverage
    ):
        depth_edges = detect_depth_gradient_stair_edges(
            depth, visible_mask, args.max_step_edges
        )
        if depth_edges:
            expanded = merge_edge_positions(
                edges_y + depth_edges,
                max(8, int(h * 0.018)),
            )
            expanded_coverage = stair_edge_coverage_ratio(expanded, amodal_mask)
            if expanded_coverage > stair_edge_coverage:
                edges_y = expanded[: args.max_step_edges]
                stair_edge_coverage = expanded_coverage
                edge_source = f'{edge_source} + depth-gradient expansion'
                # Slightly lower confidence since depth-gradient edges are
                # a secondary signal, but higher than regular fallback.
                count_score = min(len(edges_y) / 5.0, 1.0)
                coverage_score = min(
                    stair_edge_coverage / max(args.minimum_auto_stair_edge_coverage, 1e-6),
                    1.0,
                )
                stair_edge_confidence = float(
                    np.clip(0.25 + 0.20 * count_score + 0.35 * coverage_score, 0.0, 0.85)
                )
                used_depth_gradient_edge_expansion = True

    # --- Late fallback to regular fallback edges if coverage is still low ---
    if (
        geometry_mode == 'stairs'
        and not args.stair_edge_y
        and (len(edges_y) < args.minimum_detected_step_edges or stair_edge_coverage < args.minimum_auto_stair_edge_coverage)
    ):
        fallback = fallback_edges(amodal_mask, args.fallback_step_edges)
        expanded = merge_edge_positions(edges_y + fallback, max(8, int(h * 0.012)))
        expanded_coverage = stair_edge_coverage_ratio(expanded, amodal_mask)
        if expanded_coverage > stair_edge_coverage:
            edges_y = expanded[: args.max_step_edges]
            stair_edge_coverage = expanded_coverage
            edge_source = f'{edge_source} + late fallback expansion'
            used_regular_fallback_edges = True
            stair_edge_confidence = 0.25

    fx, fy, cx, cy = camera_intrinsics(w, h, args.fx, args.fy, args.cx, args.cy)
    if geometry_mode == 'stairs':
        completed_depth, completion_log, confidence = complete_depth_by_planes(
            depth,
            completion_mask,
            amodal_mask,
            edges_y,
            fx,
            fy,
            cx,
            cy,
            args.min_fit_points,
            stair_edge_slope=stair_edge_slope,
        )
        if used_regular_fallback_edges:
            confidence[completion_mask] *= 0.45
            completion_method = 'piecewise_stair_planes_with_low_confidence_fallback_edges'
        else:
            confidence[completion_mask] *= stair_edge_confidence
            completion_method = 'piecewise_stair_planes'
    elif geometry_mode in {'walkable', 'ramp'}:
        completed_depth, completion_log, confidence = complete_depth_continuous_surface(
            depth,
            completion_mask,
            amodal_mask,
            fx,
            fy,
            cx,
            cy,
            args.min_fit_points,
        )
        completion_method = completion_log[0][3]
    else:
        completed_depth, completion_log, confidence = complete_depth_generic_inpaint(
            depth, completion_mask, amodal_mask
        )
        completion_method = 'edge_aware_inpaint'

    completed_depth = enforce_hidden_boundary_continuity(
        completed_depth,
        depth,
        visible_mask,
        completion_mask,
        args.boundary_blend_radius,
        args.boundary_anchor_power,
    )
    completed_depth, depth_correction_mask, depth_plausibility = (
        enforce_hidden_depth_plausibility(
            completed_depth,
            depth,
            visible_mask,
            completion_mask,
            geometry_mode,
        )
    )
    confidence[depth_correction_mask] = np.minimum(
        confidence[depth_correction_mask],
        0.15,
    )
    # Completion is strictly confined to the reviewed hidden region.  The
    # plausibility guard already enforces this; retain the assignment as a
    # visible invariant immediately before geometry/texture export.
    completed_depth[~completion_mask] = depth[~completion_mask]
    completed_region = completion_mask
    if args.completed_rgb:
        provided_completed_rgb = read_source_grid_rgb(
            args.completed_rgb,
            source_shape,
            shape,
        )
        clean_completed_colors = rgb.copy()
        clean_completed_colors[completed_region] = provided_completed_rgb[completed_region]
        color_completion_method = 'aligned_provided_rgb_completion'
        completed_rgb_source = str(args.completed_rgb)
    else:
        clean_completed_colors = inpaint_rgb(rgb, completed_region, radius=5.0)
        color_completion_method = 'opencv_telea_fallback'
        completed_rgb_source = None
    overlay_completed_colors = rgb.copy()
    overlay_completed_colors[completed_region] = (
        0.35 * overlay_completed_colors[completed_region] + 0.65 * np.array([255, 90, 30])
    ).astype(np.uint8)

    points = pixels_to_points(completed_depth, fx, fy, cx, cy)
    valid_3d = amodal_mask & np.isfinite(completed_depth) & (completed_depth > 0)

    save_mask(output_dir / 'obstacle_mask.png', obstacle_mask)
    save_mask(output_dir / 'amodal_stair_mask.png', amodal_mask)
    save_mask(output_dir / 'amodal_target_mask.png', amodal_mask)
    save_mask(output_dir / 'target_visible_mask.png', visible_mask)
    save_mask(output_dir / 'hidden_completion_mask.png', completion_mask)
    save_mask(output_dir / 'depth_plausibility_corrections.png', depth_correction_mask)
    save_overlay(output_dir / 'debug_overlay.png', rgb, obstacle_mask, amodal_mask, edges_y, stair_edge_slope)
    Image.fromarray(visualize_depth(depth, valid_3d)).save(output_dir / 'input_depth_vis.png')
    Image.fromarray(visualize_depth(completed_depth, valid_3d)).save(output_dir / 'completed_depth_vis.png')
    hidden_depth = np.zeros_like(completed_depth, dtype=np.float32)
    hidden_depth[completion_mask] = completed_depth[completion_mask]
    completion_delta = np.zeros_like(completed_depth, dtype=np.float32)
    completion_delta[completion_mask] = np.abs(completed_depth[completion_mask] - depth[completion_mask])
    Image.fromarray(visualize_depth(completed_depth, completion_mask)).save(
        output_dir / 'hidden_completed_depth_vis.png'
    )
    Image.fromarray(visualize_depth(completion_delta, completion_mask)).save(
        output_dir / 'completion_delta_vis.png'
    )
    Image.fromarray(np.clip(confidence * 255.0, 0, 255).astype(np.uint8)).save(
        output_dir / 'completion_confidence.png'
    )
    Image.fromarray(clean_completed_colors).save(output_dir / 'completed_rgb_inpaint.png')
    Image.fromarray(overlay_completed_colors).save(output_dir / 'completed_region_overlay.png')
    np.save(output_dir / 'completed_depth.npy', completed_depth.astype(np.float32))
    np.save(output_dir / 'hidden_completed_depth.npy', hidden_depth)
    np.save(output_dir / 'completion_delta.npy', completion_delta)
    np.save(output_dir / 'completion_confidence.npy', confidence.astype(np.float32))

    point_count = write_point_cloud_ply(output_dir / 'completed_point_cloud.ply', points, clean_completed_colors, valid_3d, args.point_stride)
    depth_values = completed_depth[valid_3d]
    auto_jump = max(0.03, 0.08 * float(np.percentile(depth_values, 95) - np.percentile(depth_values, 5))) if depth_values.size else 0.2
    mesh_jump = args.max_depth_jump if args.max_depth_jump is not None else auto_jump
    vertex_count, face_count = write_mesh_ply(
        output_dir / 'completed_mesh.ply',
        points,
        clean_completed_colors,
        valid_3d,
        args.mesh_stride,
        mesh_jump,
    )

    visible_unchanged = bool(np.array_equal(completed_depth[~completion_mask], depth[~completion_mask]))
    finite_completed = bool(np.all(np.isfinite(completed_depth[amodal_mask])))
    confidence_values = confidence[completion_mask]
    mean_confidence = float(confidence_values.mean()) if confidence_values.size else 1.0
    report_title = {
        'stairs': 'Stair Geometry Completion Report',
        'walkable': 'Walkable Surface Geometry Completion Report',
        'ramp': 'Ramp Geometry Completion Report',
        'generic': 'Generic Surface Geometry Completion Report',
    }[geometry_mode]
    report = [
        f'# {report_title}',
        '',
        f'- geometry mode: `{geometry_mode}`',
        f'- completion method: `{completion_method}`',
        f'- image: `{args.image}`',
        f'- reference image used for edges: `{edge_source}`',
        f'- obstacle mask source: `{obstacle_source}`',
        f'- amodal mask source: `{amodal_source}`',
        f'- visible mask source: `{visible_source}`',
        f'- depth source: `{depth_source}`',
        f'- resolution used: `{w}x{h}`',
        f'- intrinsics: `fx={fx:.3f}, fy={fy:.3f}, cx={cx:.3f}, cy={cy:.3f}`',
        f'- stair edge y positions: `{edges_y}`',
        f'- stair edge slope: `{stair_edge_slope:.6f}`',
        f'- stair edge angle degrees: `{np.degrees(np.arctan(stair_edge_slope)):.6f}`',
        f'- used regular fallback stair edges: `{used_regular_fallback_edges}`',
        f'- used perspective-tolerant edge expansion: `{used_perspective_edge_expansion}`',
        f'- used visible-mask row-gradient edge expansion: `{used_gradient_edge_expansion}`',
        f'- stair edge confidence: `{stair_edge_confidence:.6f}`',
        f'- stair edge vertical coverage: `{stair_edge_coverage:.6f}`',
        f'- visible depth unchanged outside hidden region: `{visible_unchanged}`',
        f'- completed target depth finite: `{finite_completed}`',
        f'- boundary blend radius: `{args.boundary_blend_radius}`',
        f'- boundary anchor power: `{args.boundary_anchor_power:.6f}`',
        f'- mean hidden completion confidence: `{mean_confidence:.6f}`',
        f'- depth plausibility status: `{depth_plausibility["status"]}`',
        f'- depth plausibility candidate eligible for review: `{depth_plausibility["candidate_eligible_for_review"]}`',
        f'- corrected implausible hidden depths: `{depth_plausibility["corrected_pixel_count"]}/{depth_plausibility["hidden_pixel_count"]}`',
        f'- accepted observed-scale depth interval: `{depth_plausibility.get("accepted_depth_lower", "unavailable")} .. {depth_plausibility.get("accepted_depth_upper", "unavailable")}`',
        f'- point cloud vertices: `{point_count}`',
        f'- mesh vertices/faces: `{vertex_count}/{face_count}`',
        '- clean RGB completion: `completed_rgb_inpaint.png`',
        '',
        '## Completion Regions',
        '',
    ]
    for top, bottom, samples, mode in completion_log:
        report.append(f'- y `{top}:{bottom}` visible samples `{samples}` -> `{mode}`')
    prior_note = (
        '- If no depth map is provided, stairs use a synthetic step prior; other modes use a continuous perspective prior. Both are debug-only.'
    )
    report.extend([
        '',
        '## Notes',
        '',
        prior_note,
        '- Structured accessibility geometry constraints are written to `accessibility_geometry_analysis.json`.',
        '- If monocular or heuristic depth is provided, the output is still relative unless calibrated metric depth is used.',
        '- Hidden-depth plausibility limits are inferred from the observed source-depth scale; they are not metric safety thresholds.',
        '- Replace the default obstacle box with a real suitcase/occluder mask for meaningful completion.',
        '- For path planning, provide calibrated intrinsics and metric depth or LiDAR points; otherwise units are arbitrary.',
    ])
    (output_dir / 'run_report.md').write_text('\n'.join(report), encoding='utf-8')
    analysis = build_accessibility_geometry_analysis(
        sample_id=args.sample_id,
        category=args.category,
        geometry_mode=geometry_mode,
        visible_mask=visible_mask,
        amodal_mask=amodal_mask,
        hidden_mask=completion_mask,
        obstacle_mask=obstacle_mask,
        depth=depth,
        completed_depth=completed_depth,
        confidence=confidence,
        completion_method=completion_method,
        edges_y=edges_y,
        stair_edge_slope=stair_edge_slope,
        edge_source=edge_source,
        stair_edge_confidence=stair_edge_confidence,
        stair_edge_coverage=stair_edge_coverage,
        used_regular_fallback_edges=used_regular_fallback_edges,
        visible_depth_unchanged=visible_unchanged,
        completed_target_depth_finite=finite_completed,
        point_cloud_vertices=point_count,
        mesh_vertices=vertex_count,
        mesh_faces=face_count,
    )
    analysis_path = output_dir / 'accessibility_geometry_analysis.json'
    analysis_path.write_text(json.dumps(analysis, indent=2, ensure_ascii=False), encoding='utf-8')
    manifest = {
        'geometry_mode': geometry_mode,
        'category': analysis['category'],
        'sample_id': args.sample_id,
        'completion_method': completion_method,
        'color_completion_method': color_completion_method,
        'completed_rgb_source': completed_rgb_source,
        'completion_log': [
            {'top': int(top), 'bottom': int(bottom), 'visible_samples': int(samples), 'mode': mode}
            for top, bottom, samples, mode in completion_log
        ],
        'depth_source': depth_source,
        'used_regular_fallback_stair_edges': used_regular_fallback_edges,
        'used_perspective_edge_expansion': used_perspective_edge_expansion,
        'used_gradient_edge_expansion': used_gradient_edge_expansion,
        'used_depth_gradient_edge_expansion': used_depth_gradient_edge_expansion,
        'stair_edge_confidence': stair_edge_confidence,
        'stair_edge_coverage': stair_edge_coverage,
        'stair_edges_y': edges_y,
        'stair_edge_slope': stair_edge_slope,
        'stair_edge_angle_degrees': float(np.degrees(np.arctan(stair_edge_slope))),
        'visible_depth_unchanged_outside_hidden': visible_unchanged,
        'completed_target_depth_finite': finite_completed,
        'visible_pixel_count': int(visible_mask.sum()),
        'amodal_pixel_count': int(amodal_mask.sum()),
        'hidden_pixel_count': int(completion_mask.sum()),
        'obstacle_pixel_count': int(obstacle_mask.sum()),
        'obstacle_target_overlap_pixel_count': int((obstacle_mask & amodal_mask).sum()),
        'obstacle_visible_overlap_ratio': float((obstacle_mask & visible_mask).sum() / max(int(visible_mask.sum()), 1)),
        'hidden_changed_ratio': float(((np.abs(completed_depth - depth) > 1e-5) & completion_mask).sum() / max(int(completion_mask.sum()), 1)),
        'boundary_blend_radius': args.boundary_blend_radius,
        'boundary_anchor_power': args.boundary_anchor_power,
        'mean_hidden_completion_confidence': mean_confidence,
        'depth_plausibility': depth_plausibility,
        'depth_plausibility_corrections': str(output_dir / 'depth_plausibility_corrections.png'),
        'geometry_candidate_eligible_for_review': depth_plausibility['candidate_eligible_for_review'],
        'point_cloud_vertices': point_count,
        'mesh_vertices': vertex_count,
        'mesh_faces': face_count,
        'completed_rgb_inpaint': str(output_dir / 'completed_rgb_inpaint.png'),
        'accessibility_geometry_analysis': str(analysis_path),
        'depth_is_metric_truth': False,
    }
    (output_dir / 'geometry_manifest.json').write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False), encoding='utf-8'
    )

    print(f'Wrote outputs to {output_dir}')
    print(f'Geometry mode: {geometry_mode}')
    print(f'Completion method: {completion_method}')
    print(f'Point cloud vertices: {point_count}')
    print(f'Mesh vertices/faces: {vertex_count}/{face_count}')


def build_parser():
    parser = argparse.ArgumentParser(description='Category-aware 3D amodal completion for accessibility surfaces.')
    parser.add_argument('--image', default=None, help='Occluded RGB image (required when running reconstruction).')
    parser.add_argument('--sample-id', default=None, help='Optional stable sample id written to JSON outputs.')
    parser.add_argument(
        '--category',
        choices=['stairs', 'ramp', 'walkway', 'walkable', 'curb_cut', 'raised_curb', 'tactile_paving', 'unknown'],
        default=None,
        help='Accessibility target category. Defaults to the category implied by --geometry-mode.',
    )
    parser.add_argument('--reference-image', default=None, help='Optional clean/similar stair image used to estimate stair edge layout.')
    parser.add_argument(
        '--completed-rgb',
        default=None,
        help=(
            'Optional display-oriented RGB completion on the exact source grid. '
            'Only hidden target pixels are used as mesh colors; when omitted, '
            'OpenCV Telea is retained as a debug fallback.'
        ),
    )
    parser.add_argument('--depth', default=None, help='Optional depth map: .npy, .npz, 16-bit png, or grayscale image.')
    parser.add_argument('--depth-scale', type=float, default=1.0, help='Multiplier applied to loaded depth values.')
    parser.add_argument('--obstacle-mask', default=None, help='Binary mask of the occluding object, white/nonzero is obstacle.')
    parser.add_argument('--amodal-mask', default=None, help='Binary mask of the full target region, including occluded area.')
    parser.add_argument('--visible-mask', default=None, help='Reviewed visible target mask. Hidden completion is amodal minus visible.')
    parser.add_argument(
        '--geometry-mode',
        choices=['stairs', 'walkable', 'ramp', 'generic'],
        default=None,
        help='Geometry prior. Defaults to the category-compatible mode; contradictory category/mode pairs are rejected.',
    )
    parser.add_argument('--occlusion-box', action='append', type=parse_box, default=[], help='Obstacle box x1,y1,x2,y2. Coordinates can be normalized or pixels. Repeatable.')
    parser.add_argument('--no-default-obstacle', action='store_true', help='Do not use the built-in approximate suitcase box when no mask/box is given.')
    parser.add_argument('--output-dir', default='./output/stair_geometry/', help='Output directory.')
    parser.add_argument('--max-size', type=int, default=1280, help='Resize longest image side before processing. Use 0 to keep original size.')
    parser.add_argument('--max-step-edges', type=int, default=9)
    parser.add_argument('--fallback-step-edges', type=int, default=5)
    parser.add_argument('--minimum-detected-step-edges', type=int, default=2)
    parser.add_argument('--minimum-auto-stair-edge-coverage', type=float, default=0.35)
    parser.add_argument('--stair-edge-y', action='append', type=int, default=[], help='Reviewed stair tread-edge y coordinate. Repeatable; disables regular fallback edges.')
    parser.add_argument('--stair-edge-angle-degrees', type=float, default=0.0, help='Reviewed tread-edge angle for --stair-edge-y; 0 means horizontal.')
    parser.add_argument('--synthetic-near', type=float, default=1.0)
    parser.add_argument('--synthetic-far', type=float, default=6.0)
    parser.add_argument('--fx', type=float, default=None)
    parser.add_argument('--fy', type=float, default=None)
    parser.add_argument('--cx', type=float, default=None)
    parser.add_argument('--cy', type=float, default=None)
    parser.add_argument('--min-fit-points', type=int, default=800)
    parser.add_argument('--boundary-blend-radius', type=int, default=12)
    parser.add_argument(
        '--boundary-anchor-power',
        type=float,
        default=0.45,
        help='Power applied to hidden-boundary anchoring weights; lower values preserve visible-depth continuity over a wider band.',
    )
    parser.add_argument('--point-stride', type=int, default=3)
    parser.add_argument('--mesh-stride', type=int, default=5)
    parser.add_argument('--max-depth-jump', type=float, default=None, help='Max depth discontinuity allowed when making mesh faces.')
    return parser


if __name__ == '__main__':
    args = build_parser().parse_args()
    if args.max_size == 0:
        args.max_size = None
    run(args)