accessibilityamodal/pipeline.py

"""AccessibilityAmodal mask construction and geometry completion.

The script is an orchestration layer. It does not vendor the external projects;
instead it consumes their outputs or calls a running LISA Gradio server when
available. This keeps the project-owned AccessibilityAmodal workflow separate
from the incompatible dependency trees of LISA, Grounded-SAM, Amodal-Wild,
Diffusion-VAS, pix2gestalt, and the optional licensed visual 3D backend.
"""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any
from urllib.request import urlopen

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

from accessibilityamodal.category_overrides import (
    DEFAULT_CATEGORY_OVERRIDES_PATH,
    load_reviewed_category_overrides,
    resolve_reviewed_category,
)
from accessibilityamodal.visual_completion import (
    build_completion_envelope,
    build_visual_removal_mask,
)


PROJECT_ROOT = Path(__file__).resolve().parents[1]


ACCESSIBILITY_PRESETS = {
    'stairs': {
        'target_prompt': 'Can you segment all visible stair steps and stair walking surfaces in this image? Please output segmentation mask.',
        'obstacle_prompt': 'Can you segment the suitcase, luggage, person, vehicle, or any object occluding the stairs in this image? Please output segmentation mask.',
    },
    'tactile_paving': {
        'target_prompt': 'Can you segment the visible tactile paving or blind sidewalk guiding tiles in this image? Please output segmentation mask.',
        'obstacle_prompt': 'Can you segment the object blocking the tactile paving path in this image? Please output segmentation mask.',
    },
    'ramp': {
        'target_prompt': 'Can you segment the visible wheelchair ramp or accessible sloped walking surface in this image? Please output segmentation mask.',
        'obstacle_prompt': 'Can you segment the object blocking the wheelchair ramp or accessible route in this image? Please output segmentation mask.',
    },
    'curb_cut': {
        'target_prompt': 'Can you segment the visible curb cut ramp, sidewalk transition, and accessible sloped crossing surface in this image? Please output segmentation mask.',
        'obstacle_prompt': 'Can you segment objects blocking the curb cut, sidewalk transition, or accessible crossing in this image? Please output segmentation mask.',
    },
    'raised_curb': {
        'target_prompt': 'Can you segment the visible raised road curb, including its continuous top edge and vertical face, in this image? Exclude stairs and curb-cut ramps. Please output segmentation mask.',
        'obstacle_prompt': 'Can you segment people, vehicles, vegetation, or other objects occluding the raised curb in this image? Please output segmentation mask.',
    },
    'walkway': {
        'target_prompt': 'Can you segment the visible accessible pedestrian walkway or sidewalk support surface in this image? Please output segmentation mask.',
        'obstacle_prompt': 'Can you segment obstacles blocking the accessible walkway or sidewalk path in this image? Please output segmentation mask.',
    },
    'walkable': {
        'target_prompt': 'Can you segment the visible accessible walkable path for a wheelchair or delivery robot in this image? Please output segmentation mask.',
        'obstacle_prompt': 'Can you segment obstacles blocking the accessible walking path in this image? Please output segmentation mask.',
    },
}


def read_rgb(path: str | Path) -> np.ndarray:
    """Read RGB in the same display orientation used by the mask proposal stage."""
    return np.array(ImageOps.exif_transpose(Image.open(path)).convert('RGB'))


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


def read_mask(path: str | Path, shape: tuple[int, int]) -> np.ndarray:
    mask = np.array(ImageOps.exif_transpose(Image.open(path)).convert('L')) > 127
    h, w = shape
    if mask.shape != (h, w):
        raise ValueError(
            f'Mask/RGB raster mismatch for {path}: mask={mask.shape}, rgb={(h, w)}. '
            'Refusing to resize because it can hide EXIF-orientation misalignment.'
        )
    return mask


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


def mask_from_lisa_json(path: str | Path, shape: tuple[int, int]) -> np.ndarray:
    with open(path, 'r', encoding='utf-8') as f:
        data: Any = json.load(f)
    if isinstance(data, dict) and 'data' in data:
        data = data['data']
    arr = np.array(data)
    if arr.ndim == 3:
        arr = arr[..., 0]
    mask = arr.astype(np.float32) > 0
    h, w = shape
    if mask.shape != (h, w):
        raise ValueError(
            f'LISA mask/RGB raster mismatch for {path}: mask={mask.shape}, rgb={(h, w)}. '
            'Refusing to resize because the LISA input must share the normalized display-oriented grid.'
        )
    return mask


def lisa_server_reachable(server_url: str, timeout: float = 2.0) -> bool:
    try:
        with urlopen(server_url, timeout=timeout):
            return True
    except Exception:
        return False


def call_lisa(server_url: str, image_path: str, prompt: str, out_prefix: Path, shape: tuple[int, int]) -> np.ndarray:
    try:
        from gradio_client import Client
    except ImportError as exc:
        raise RuntimeError('gradio_client is not installed in this environment') from exc

    client = Client(server_url)
    result = client.predict(prompt, image_path, api_name='/predict')
    if not isinstance(result, (tuple, list)) or len(result) < 2:
        raise RuntimeError(f'Unexpected LISA response: {result!r}')

    rendered_path = Path(result[0])
    json_path = Path(result[1])
    copied_rendered = out_prefix.with_suffix('.png')
    copied_json = out_prefix.with_suffix('.json')
    if rendered_path.exists():
        shutil.copyfile(rendered_path, copied_rendered)
    shutil.copyfile(json_path, copied_json)
    return mask_from_lisa_json(copied_json, shape)


def largest_components(mask: np.ndarray, keep: int = 3, min_area: int = 64) -> np.ndarray:
    num, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), connectivity=8)
    if num <= 1:
        return mask.astype(bool)
    areas = [(idx, stats[idx, cv2.CC_STAT_AREA]) for idx in range(1, num)]
    areas = [(idx, area) for idx, area in areas if area >= min_area]
    areas.sort(key=lambda item: item[1], reverse=True)
    kept = np.zeros_like(mask, dtype=bool)
    for idx, _ in areas[:keep]:
        kept |= labels == idx
    return kept


def heuristic_target_mask(rgb: np.ndarray, preset: str) -> np.ndarray:
    h, w = rgb.shape[:2]
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(cv2.GaussianBlur(gray, (5, 5), 0), 60, 160)
    lower = np.zeros((h, w), dtype=np.uint8)
    lower[int(0.12 * h):, :] = 1
    if preset == 'stairs':
        horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (max(25, w // 18), 3))
        horizontal = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, horizontal_kernel)
        band = cv2.dilate(horizontal, cv2.getStructuringElement(cv2.MORPH_RECT, (max(35, w // 14), max(9, h // 80))))
        mask = (band > 0) & (lower > 0)
        if mask.sum() < 0.03 * h * w:
            mask = lower.astype(bool)
        return mask
    return lower.astype(bool)


def heuristic_obstacle_mask(rgb: np.ndarray, boxes: list[tuple[float, float, float, float]], target_mask: np.ndarray) -> np.ndarray:
    h, w = target_mask.shape
    if boxes:
        return boxes_to_mask(boxes, (h, w))
    # Conservative central occluder prior for one-off experiments. It is not a
    # replacement for a real LISA/Grounded-SAM obstacle mask.
    return boxes_to_mask([(0.36, 0.32, 0.68, 0.88)], (h, w)) & cv2.dilate(
        target_mask.astype(np.uint8),
        cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (max(31, w // 24), max(31, h // 24))),
    ).astype(bool)


def row_span_amodal_mask(visible_mask: np.ndarray, obstacle_mask: np.ndarray, fallback_mask: np.ndarray) -> np.ndarray:
    h, w = visible_mask.shape
    rows = np.where(visible_mask.any(axis=1))[0]
    if rows.size < 2:
        return fallback_mask

    lefts = np.full(h, np.nan, dtype=np.float32)
    rights = np.full(h, np.nan, dtype=np.float32)
    for y in rows:
        xs = np.where(visible_mask[y])[0]
        if xs.size:
            lefts[y] = np.percentile(xs, 3)
            rights[y] = np.percentile(xs, 97)

    known = np.where(np.isfinite(lefts) & np.isfinite(rights))[0]
    y_all = np.arange(h)
    left_interp = np.interp(y_all, known, lefts[known])
    right_interp = np.interp(y_all, known, rights[known])

    margin = max(20, int(0.04 * w))
    top = max(0, int(known.min() - 0.04 * h))
    bottom = min(h, int(known.max() + 0.08 * h))
    amodal = np.zeros((h, w), dtype=bool)
    for y in range(top, bottom):
        left = int(np.clip(left_interp[y] - margin, 0, w - 1))
        right = int(np.clip(right_interp[y] + margin, left + 1, w))
        amodal[y, left:right] = True

    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (max(21, w // 48), max(21, h // 48)))
    amodal = cv2.morphologyEx(amodal.astype(np.uint8), cv2.MORPH_CLOSE, kernel).astype(bool)
    amodal |= visible_mask
    return amodal


def reviewed_geometry_support_prior(
    reviewed_amodal: np.ndarray,
    target_visible: np.ndarray,
    obstacle_mask: np.ndarray,
    *,
    envelope_margin_fraction: float = 0.012,
    maximum_hull_expansion_ratio: float = 2.0,
) -> tuple[np.ndarray, dict[str, Any]]:
    """Bridge reviewed occluder gaps without replacing the reviewed mask.

    Accessibility annotations can be semantically correct yet retain a person-
    shaped notch in the support surface.  Raw triangulation then preserves that
    notch as a hole.  This helper builds a *geometry-only hypothesis*:

    1. retain complete obstacle components that intersect the reviewed hidden
       target;
    2. form small component bounding envelopes, matching the 2D removal policy;
    3. add only envelope pixels that lie inside a robust convex support hull.

    The hull is derived from non-trivial reviewed components, so isolated mask
    speckles cannot stretch it across the image.  Highly concave/ambiguous
    targets fail closed and keep the reviewed mask unchanged.
    """

    amodal = np.asarray(reviewed_amodal, dtype=bool)
    visible = np.asarray(target_visible, dtype=bool)
    obstacle = np.asarray(obstacle_mask, dtype=bool)
    if amodal.ndim != 2 or visible.ndim != 2 or obstacle.ndim != 2:
        raise ValueError("reviewed amodal, visible, and obstacle masks must be 2D")
    if amodal.shape != visible.shape or amodal.shape != obstacle.shape:
        raise ValueError("reviewed amodal, visible, and obstacle masks must align")
    if envelope_margin_fraction < 0:
        raise ValueError("envelope_margin_fraction must be non-negative")
    if maximum_hull_expansion_ratio < 1:
        raise ValueError("maximum_hull_expansion_ratio must be at least 1")

    reviewed = amodal | visible
    hidden = reviewed & ~visible
    visual_removal, removal_stats = build_visual_removal_mask(hidden, obstacle)
    completion_envelope, envelope_stats = build_completion_envelope(
        visual_removal,
        obstacle,
        margin_fraction=envelope_margin_fraction,
    )

    reviewed_pixels = int(reviewed.sum())
    base_metadata: dict[str, Any] = {
        "policy": (
            "reviewed_mask_plus_occluder_envelope_inside_robust_convex_support_hull"
        ),
        "reviewed_pixel_count": reviewed_pixels,
        "reviewed_hidden_pixel_count": int(hidden.sum()),
        "visual_removal_pixel_count": int(visual_removal.sum()),
        "completion_envelope_pixel_count": int(completion_envelope.sum()),
        "envelope_margin_fraction": float(envelope_margin_fraction),
        "maximum_hull_expansion_ratio": float(maximum_hull_expansion_ratio),
        "obstacle_components_retained": int(
            removal_stats["obstacle_components_retained"]
        ),
        "envelope_component_count": int(envelope_stats["component_count"]),
    }
    if reviewed_pixels == 0:
        return reviewed.copy(), {
            **base_metadata,
            "status": "unchanged_empty_reviewed_mask",
            "hull_source_pixel_count": 0,
            "hull_pixel_count": 0,
            "hull_expansion_ratio": None,
            "added_support_pixel_count": 0,
        }

    count, labels, component_stats, _ = cv2.connectedComponentsWithStats(
        reviewed.astype(np.uint8),
        connectivity=8,
    )
    minimum_component_area = max(32, int(round(reviewed_pixels * 0.002)))
    hull_source = np.zeros_like(reviewed)
    retained_component_count = 0
    for label in range(1, count):
        area = int(component_stats[label, cv2.CC_STAT_AREA])
        if area >= minimum_component_area:
            hull_source |= labels == label
            retained_component_count += 1
    if not hull_source.any():
        hull_source = reviewed.copy()
        retained_component_count = max(0, count - 1)

    ys, xs = np.where(hull_source)
    points = np.column_stack([xs, ys]).astype(np.int32)
    hull = np.zeros_like(reviewed, dtype=np.uint8)
    if points.shape[0] >= 3:
        cv2.fillConvexPoly(hull, cv2.convexHull(points), 1)
    else:
        hull[hull_source] = 1
    hull_mask = hull.astype(bool)
    hull_source_pixels = int(hull_source.sum())
    hull_pixels = int(hull_mask.sum())
    hull_ratio = hull_pixels / max(hull_source_pixels, 1)

    metadata = {
        **base_metadata,
        "minimum_hull_component_area": minimum_component_area,
        "hull_source_component_count": retained_component_count,
        "hull_source_pixel_count": hull_source_pixels,
        "hull_source_fraction_of_reviewed": round(
            hull_source_pixels / reviewed_pixels,
            8,
        ),
        "hull_pixel_count": hull_pixels,
        "hull_expansion_ratio": round(hull_ratio, 8),
    }
    if hull_ratio > maximum_hull_expansion_ratio:
        return reviewed.copy(), {
            **metadata,
            "status": "unchanged_ambiguous_hull_expansion",
            "added_support_pixel_count": 0,
        }

    added_support = completion_envelope & hull_mask & ~reviewed
    support_prior = reviewed | added_support
    return support_prior, {
        **metadata,
        "status": (
            "augmented_occluded_support_hypothesis"
            if added_support.any()
            else "unchanged_no_supported_gap"
        ),
        "added_support_pixel_count": int(added_support.sum()),
        "support_prior_pixel_count": int(support_prior.sum()),
        "support_prior_is_metric_truth": False,
        "human_review_required": True,
    }


def maybe_refine_masks_with_sam(
    args,
    rgb: np.ndarray,
    target_visible: np.ndarray,
    obstacle: np.ndarray,
    output_dir: Path,
) -> tuple[np.ndarray, np.ndarray, str | None]:
    if not args.sam_refine:
        return target_visible, obstacle, None
    if not args.sam_checkpoint:
        message = 'SAM refinement requested but --sam-checkpoint was not provided'
        if args.sam_fallback == 'error':
            raise ValueError(message)
        print(f'{message}; keeping coarse masks', file=sys.stderr)
        return target_visible, obstacle, message

    pre_target = output_dir / 'pre_sam_target_visible_mask.png'
    pre_obstacle = output_dir / 'pre_sam_obstacle_mask.png'
    save_mask(pre_target, target_visible)
    save_mask(pre_obstacle, obstacle)

    sam_output_dir = output_dir / 'sam_refine'
    cmd = [
        sys.executable,
        '-m',
        'accessibilityamodal.sam_refinement',
        '--image', args.image,
        '--target-mask', str(pre_target),
        '--obstacle-mask', str(pre_obstacle),
        '--output-dir', str(sam_output_dir),
        '--sam-repo', args.sam_repo,
        '--sam-checkpoint', args.sam_checkpoint,
        '--sam-model-type', args.sam_model_type,
        '--device', args.sam_device,
        '--keep-target-components', str(args.keep_target_components),
        '--keep-obstacle-components', str(args.keep_obstacle_components),
    ]
    print('Running:', ' '.join(cmd))
    try:
        subprocess.run(cmd, check=True)
    except Exception as exc:
        if args.sam_fallback == 'error':
            raise
        message = f'SAM refinement failed; keeping coarse masks: {exc}'
        print(message, file=sys.stderr)
        return target_visible, obstacle, message

    refined_target = read_mask(sam_output_dir / 'target_visible_mask_sam.png', rgb.shape[:2])
    refined_obstacle = read_mask(sam_output_dir / 'obstacle_mask_sam.png', rgb.shape[:2])
    return refined_target, refined_obstacle, f'SAM: {args.sam_model_type}'


def write_three_value_mask(
    path: Path, target_visible: np.ndarray, amodal_mask: np.ndarray
) -> None:
    mask = np.full(amodal_mask.shape, 255, dtype=np.uint8)
    hidden = amodal_mask & ~target_visible
    mask[target_visible] = 188
    mask[hidden] = 0
    Image.fromarray(mask).save(path)


def overlay(rgb: np.ndarray, masks: list[tuple[np.ndarray, tuple[int, int, int], float]]) -> np.ndarray:
    out = rgb.astype(np.float32).copy()
    for mask, color, alpha in masks:
        if mask.any():
            out[mask] = out[mask] * (1.0 - alpha) + np.array(color, dtype=np.float32) * alpha
    return np.clip(out, 0, 255).astype(np.uint8)


def build_masks(args) -> dict[str, Path]:
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    rgb = read_rgb(args.image)
    shape = rgb.shape[:2]
    preset = ACCESSIBILITY_PRESETS[args.preset]
    target_prompt = args.target_prompt or preset['target_prompt']
    obstacle_prompt = args.obstacle_prompt or preset['obstacle_prompt']
    lisa_image_path = output_dir / 'lisa_input_display_oriented.png'
    if args.use_lisa:
        Image.fromarray(rgb).save(lisa_image_path)

    if args.target_visible_mask:
        target_visible = read_mask(args.target_visible_mask, shape)
        target_source = args.target_visible_mask
    elif args.use_lisa and lisa_server_reachable(args.lisa_server_url):
        target_visible = call_lisa(args.lisa_server_url, str(lisa_image_path), target_prompt, output_dir / 'lisa_target_visible', shape)
        target_source = f'LISA: {target_prompt}'
    else:
        if args.use_lisa and args.lisa_fallback == 'error':
            raise RuntimeError(f'LISA server is not reachable: {args.lisa_server_url}')
        target_visible = heuristic_target_mask(rgb, args.preset)
        target_source = 'heuristic target mask'

    if not args.target_visible_mask:
        target_visible = largest_components(
            target_visible,
            keep=args.keep_target_components,
            min_area=max(64, shape[0] * shape[1] // 20000),
        )

    if args.obstacle_mask:
        obstacle = read_mask(args.obstacle_mask, shape)
        obstacle_source = args.obstacle_mask
    elif args.use_lisa and lisa_server_reachable(args.lisa_server_url):
        try:
            obstacle = call_lisa(args.lisa_server_url, str(lisa_image_path), obstacle_prompt, output_dir / 'lisa_obstacle', shape)
            obstacle_source = f'LISA: {obstacle_prompt}'
        except Exception as exc:
            if args.lisa_fallback == 'error':
                raise
            print(f'LISA obstacle segmentation failed, falling back to heuristic obstacle mask: {exc}', file=sys.stderr)
            obstacle = heuristic_obstacle_mask(rgb, args.occlusion_box, target_visible)
            obstacle_source = 'heuristic obstacle mask after LISA failure'
    else:
        if args.use_lisa and args.lisa_fallback == 'error':
            raise RuntimeError(f'LISA server is not reachable: {args.lisa_server_url}')
        obstacle = heuristic_obstacle_mask(rgb, args.occlusion_box, target_visible)
        obstacle_source = 'heuristic obstacle mask'

    if not args.obstacle_mask:
        obstacle = largest_components(
            obstacle,
            keep=args.keep_obstacle_components,
            min_area=max(64, shape[0] * shape[1] // 30000),
        )
    target_visible, obstacle, sam_source = maybe_refine_masks_with_sam(args, rgb, target_visible, obstacle, output_dir)
    if sam_source:
        target_source = f'{target_source} + {sam_source}'
        obstacle_source = f'{obstacle_source} + {sam_source}'
    # Occluders are never part of the support surface. This guard also fixes
    # broad heuristic/LISA target masks that accidentally include a person,
    # vehicle, railing, or clutter region.
    target_visible &= ~obstacle

    if args.target_amodal_mask:
        provided_amodal = read_mask(args.target_amodal_mask, shape)
        reviewed_amodal = provided_amodal | target_visible
        amodal_support_prior, support_prior_metadata = (
            reviewed_geometry_support_prior(
                reviewed_amodal,
                target_visible,
                obstacle,
            )
        )
        amodal = amodal_support_prior.copy()
        hidden = amodal_support_prior & ~target_visible
        amodal_source = args.target_amodal_mask
        completion_policy = (
            "reviewed target mask plus auditable occluder-constrained "
            "geometry support hypothesis"
        )
    else:
        fallback_target = heuristic_target_mask(rgb, args.preset)
        # Row-span completion is only a structural support prior. Restrict hidden
        # completion to observed obstacle pixels so disjoint stairs/walkways are not
        # connected across unoccluded image regions.
        amodal_support_prior = row_span_amodal_mask(
            target_visible, obstacle, fallback_target
        )
        hidden = obstacle & amodal_support_prior & ~target_visible
        amodal = target_visible | hidden
        reviewed_amodal = amodal.copy()
        amodal_source = "obstacle-constrained perspective baseline"
        completion_policy = "obstacle-constrained row-span structural support prior"
        support_prior_metadata = {
            "policy": completion_policy,
            "status": "heuristic_obstacle_constrained_support_prior",
            "reviewed_pixel_count": int(amodal.sum()),
            "added_support_pixel_count": 0,
            "support_prior_pixel_count": int(amodal_support_prior.sum()),
            "support_prior_is_metric_truth": False,
            "human_review_required": True,
        }

    visual_3d_hidden = reviewed_amodal & ~target_visible
    paths = {
        'target_visible': output_dir / 'target_visible_mask.png',
        'obstacle': output_dir / 'obstacle_mask.png',
        'amodal': output_dir / 'amodal_accessibility_mask.png',
        'reviewed_amodal': output_dir / 'reviewed_amodal_source_mask.png',
        'hidden': output_dir / 'hidden_completion_mask.png',
        'support_prior': output_dir / 'amodal_support_prior.png',
        'visual_3d_condition': output_dir / 'amodal3d_three_value_mask.png',
        'visual_3d_condition_compatibility': (
            output_dir / 'accessibilityamodal_condition_mask.png'
        ),
        'overlay': output_dir / 'mask_debug_overlay.png',
        'manifest': output_dir / 'mask_manifest.json',
    }
    save_mask(paths['target_visible'], target_visible)
    save_mask(paths['obstacle'], obstacle)
    save_mask(paths['amodal'], amodal)
    save_mask(paths['reviewed_amodal'], reviewed_amodal)
    save_mask(paths['hidden'], hidden)
    save_mask(paths['support_prior'], amodal_support_prior)
    # Visual Accessibility3D conditioning must reflect only the reviewed support.
    # The expanded support prior remains a geometry-only hypothesis consumed by
    # reconstruct_steps and must not leak into the visual three-value mask.
    write_three_value_mask(
        paths['visual_3d_condition'],
        target_visible,
        reviewed_amodal,
    )
    shutil.copyfile(
        paths['visual_3d_condition'],
        paths['visual_3d_condition_compatibility'],
    )
    debug = overlay(rgb, [
        (amodal, (0, 150, 255), 0.30),
        (target_visible, (0, 220, 80), 0.45),
        (hidden, (255, 40, 40), 0.65),
    ])
    Image.fromarray(debug).save(paths['overlay'])

    manifest = {
        'sample_id': args.sample_id,
        'image': args.image,
        'preset': args.preset,
        'source_preset': getattr(args, 'source_preset', args.preset),
        'category_override_applied': getattr(args, 'category_override_applied', False),
        'category': 'walkway' if args.preset == 'walkable' else args.preset,
        'target_prompt': target_prompt,
        'obstacle_prompt': obstacle_prompt,
        'target_source': target_source,
        'obstacle_source': obstacle_source,
        'amodal_source': amodal_source,
        'completion_policy': completion_policy,
        'target_visible_pixel_count': int(target_visible.sum()),
        'obstacle_pixel_count': int(obstacle.sum()),
        'target_obstacle_overlap_pixel_count': int((target_visible & obstacle).sum()),
        'reviewed_amodal_pixel_count': int(reviewed_amodal.sum()),
        'reviewed_hidden_pixel_count': int(visual_3d_hidden.sum()),
        'geometry_support_amodal_pixel_count': int(amodal.sum()),
        'geometry_support_hidden_pixel_count': int(hidden.sum()),
        # Compatibility counts retained for existing manifest consumers.
        'hidden_pixel_count': int(hidden.sum()),
        'amodal_pixel_count': int(amodal.sum()),
        'geometry_support_prior': support_prior_metadata,
        'mask_roles': {
            'visual_3d_condition': {
                'path': str(paths['visual_3d_condition']),
                'compatibility_path': str(
                    paths['visual_3d_condition_compatibility']
                ),
                'role': 'strict_reviewed_amodal_condition_for_amodal3d',
                'source_mask': 'reviewed_amodal',
                'uses_geometry_support_prior': False,
                'pixel_values': {
                    'background': 255,
                    'visible_target': 188,
                    'hidden_target': 0,
                },
                'amodal_pixel_count': int(reviewed_amodal.sum()),
                'hidden_pixel_count': int(visual_3d_hidden.sum()),
            },
            'geometry_support_prior': {
                'path': str(paths['support_prior']),
                'role': 'geometry_only_support_hypothesis',
                'source_mask': 'reviewed_amodal',
                'added_pixel_count': int(
                    (amodal_support_prior & ~reviewed_amodal).sum()
                ),
                'amodal_pixel_count': int(amodal_support_prior.sum()),
                'hidden_pixel_count': int(
                    (amodal_support_prior & ~target_visible).sum()
                ),
            },
            'reconstruct_steps_amodal_input': {
                'path': str(paths['amodal']),
                'role': 'geometry_support_input_for_reconstruct_steps',
                'source_mask': 'amodal_support_prior',
                'uses_geometry_support_prior': True,
            },
        },
        'outputs': {key: str(value) for key, value in paths.items() if key != 'manifest'},
        'external_repos_reviewed': {
            'amodal': 'external backend; path supplied by user',
            'Amodal-Wild': 'external backend; path supplied by user',
            'diffusion-vas': 'external backend; path supplied by user',
            'pix2gestalt': 'external backend; path supplied by user',
        },
    }
    paths['manifest'].write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding='utf-8')
    return paths


def prepare_depth(args) -> tuple[str | None, float, str]:
    if args.depth:
        return args.depth, args.depth_scale, f'user depth: {args.depth}'
    if not args.estimate_depth:
        synthetic_kind = (
            'synthetic perspective stair prior'
            if args.preset == 'stairs'
            else 'synthetic continuous perspective prior'
        )
        return None, args.depth_scale, synthetic_kind

    output_dir = Path(args.output_dir)
    depth_path = output_dir / 'estimated_depth.npy'
    depth_vis_path = output_dir / 'estimated_depth_vis.png'
    depth_manifest_path = output_dir / 'depth_manifest.json'
    cmd = [
        sys.executable,
        '-m',
        'accessibilityamodal.depth',
        '--image', args.image,
        '--output-depth', str(depth_path),
        '--output-vis', str(depth_vis_path),
        '--manifest', str(depth_manifest_path),
        '--engine', args.depth_engine,
        '--near', str(args.depth_near),
        '--far', str(args.depth_far),
        '--max-size', str(args.depth_max_size),
    ]
    if args.depth_engine == 'transformers':
        cmd.extend(['--model', args.depth_model, '--device', args.depth_device])
    elif args.depth_engine == 'depth_anything_v2':
        cmd.extend([
            '--device', args.depth_device,
            '--depth-anything-repo', args.depth_anything_repo,
            '--encoder', args.depth_encoder,
            '--input-size', str(args.depth_input_size),
        ])
        if args.depth_checkpoint:
            cmd.extend(['--checkpoint', args.depth_checkpoint])
        if args.metric_depth:
            cmd.extend(['--metric-depth', '--max-depth', str(args.max_depth)])
    elif args.depth_engine == 'vggt':
        cmd.extend([
            '--device', args.depth_device,
            '--vggt-repo', args.vggt_repo,
            '--vggt-model', args.vggt_model,
            '--vggt-input-size', str(args.vggt_input_size),
            '--output-conf', str(output_dir / 'vggt_depth_conf.npy'),
            '--output-conf-vis', str(output_dir / 'vggt_depth_conf.png'),
            '--output-raw-depth', str(output_dir / 'vggt_raw_depth.npy'),
            '--output-world-points', str(output_dir / 'vggt_world_points.npy'),
            '--output-point-conf', str(output_dir / 'vggt_world_points_conf.npy'),
            '--output-point-cloud', str(output_dir / 'vggt_point_cloud.ply'),
            '--point-conf-threshold', str(args.vggt_point_conf_threshold),
            '--max-point-cloud-points', str(args.vggt_max_point_cloud_points),
        ])
        if args.vggt_checkpoint:
            cmd.extend(['--vggt-checkpoint', args.vggt_checkpoint])
        if args.vggt_keep_raw_depth:
            cmd.append('--vggt-keep-raw-depth')
    if args.invert_depth:
        cmd.append('--invert-depth')

    print('Running:', ' '.join(cmd))
    try:
        subprocess.run(cmd, check=True)
    except Exception as exc:
        if args.depth_fallback == 'error':
            raise
        print(f'Depth estimation failed, falling back to reconstruct_steps synthetic depth: {exc}', file=sys.stderr)
        return None, args.depth_scale, f'depth estimation failed; synthetic fallback: {exc}'
    depth_kind = f'estimated depth: {args.depth_engine}'
    if args.metric_depth:
        depth_kind += ' metric'
    return str(depth_path), 1.0, depth_kind


def run_reconstruct(args, masks: dict[str, Path], depth_path: str | None, depth_scale: float) -> None:
    if args.skip_reconstruct:
        return
    category = 'walkway' if args.preset == 'walkable' else args.preset
    geometry_mode = (
        'stairs' if args.preset == 'stairs'
        else 'ramp' if args.preset == 'ramp'
        else 'walkable'
    )
    cmd = [
        sys.executable,
        '-m',
        'accessibilityamodal.reconstruct',
        '--image', args.image,
        '--obstacle-mask', str(masks['obstacle']),
        '--amodal-mask', str(masks['amodal']),
        '--visible-mask', str(masks['target_visible']),
        '--geometry-mode', geometry_mode,
        '--category', category,
        '--output-dir', str(Path(args.output_dir) / 'geometry'),
    ]
    if args.sample_id:
        cmd.extend(['--sample-id', args.sample_id])
    if args.reference_image:
        cmd.extend(['--reference-image', args.reference_image])
    completed_rgb = getattr(args, 'completed_rgb', None)
    if completed_rgb:
        cmd.extend(['--completed-rgb', completed_rgb])
    if depth_path:
        cmd.extend(['--depth', depth_path, '--depth-scale', str(depth_scale)])
    if args.fx is not None:
        cmd.extend(['--fx', str(args.fx)])
    if args.fy is not None:
        cmd.extend(['--fy', str(args.fy)])
    if args.cx is not None:
        cmd.extend(['--cx', str(args.cx)])
    if args.cy is not None:
        cmd.extend(['--cy', str(args.cy)])
    print('Running:', ' '.join(cmd))
    subprocess.run(cmd, check=True)


def write_pipeline_manifest(args, masks: dict[str, Path], depth_source: str, depth_path: str | None) -> None:
    output_dir = Path(args.output_dir)
    manifest = {
        'sample_id': args.sample_id,
        'image': args.image,
        'reference_image': args.reference_image,
        'completed_rgb': args.completed_rgb,
        'preset': args.preset,
        'source_preset': getattr(args, 'source_preset', args.preset),
        'category_override_applied': getattr(args, 'category_override_applied', False),
        'category': 'walkway' if args.preset == 'walkable' else args.preset,
        'use_lisa': args.use_lisa,
        'lisa_server_url': args.lisa_server_url if args.use_lisa else None,
        'depth_source': depth_source,
        'depth_path': depth_path,
        'reconstruct_output_dir': None if args.skip_reconstruct else str(output_dir / 'geometry'),
        'mask_outputs': {key: str(value) for key, value in masks.items()},
        'note': 'Depth from monocular/heuristic sources is relative unless calibrated metric depth is provided.',
    }
    (output_dir / 'pipeline_manifest.json').write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding='utf-8')


def build_parser():
    parser = argparse.ArgumentParser(description='Build accessibility masks from text prompts and run 3D geometry completion.')
    parser.add_argument('--image', default='./3R/2.jpg')
    parser.add_argument('--sample-id', default=None)
    parser.add_argument('--reference-image', default=None)
    parser.add_argument(
        '--completed-rgb',
        default=None,
        help='Aligned 2D completion used only to texture the reviewed hidden target region.',
    )
    parser.add_argument('--preset', choices=sorted(ACCESSIBILITY_PRESETS), default='stairs')
    parser.add_argument(
        '--category-overrides',
        default=str(DEFAULT_CATEGORY_OVERRIDES_PATH),
        help='Human-reviewed sample category overrides.',
    )
    parser.add_argument('--target-prompt', default=None)
    parser.add_argument('--obstacle-prompt', default=None)
    parser.add_argument('--target-visible-mask', default=None, help='Use an existing visible target mask instead of LISA/heuristic.')
    parser.add_argument('--target-amodal-mask', default=None, help='Use a reviewed amodal target mask and bypass automatic mask completion.')
    parser.add_argument('--obstacle-mask', default=None, help='Use an existing obstacle mask instead of LISA/heuristic.')
    parser.add_argument('--occlusion-box', action='append', type=parse_box, default=[], help='Fallback obstacle box, normalized or pixel x1,y1,x2,y2.')
    parser.add_argument('--use-lisa', action='store_true', help='Call a running LISA Gradio server for text-guided masks.')
    parser.add_argument('--lisa-server-url', default='http://127.0.0.1:7860/')
    parser.add_argument('--lisa-fallback', choices=['heuristic', 'error'], default='heuristic', help='What to do when a requested LISA mask cannot be obtained.')
    parser.add_argument('--output-dir', default='./output/accessibility_pipeline/')
    parser.add_argument('--keep-target-components', type=int, default=6)
    parser.add_argument('--keep-obstacle-components', type=int, default=4)
    parser.add_argument('--depth', default=None)
    parser.add_argument('--depth-scale', type=float, default=1.0)
    parser.add_argument('--estimate-depth', action='store_true', help='Generate a local relative depth map before reconstruction.')
    parser.add_argument('--depth-engine', choices=['heuristic', 'transformers', 'depth_anything_v2', 'vggt'], default='heuristic', help='Local depth generator used with --estimate-depth.')
    parser.add_argument('--depth-model', default='depth-anything/Depth-Anything-V2-Small-hf', help='Transformers depth model id or local path.')
    parser.add_argument('--depth-device', default='auto', help='auto, cpu, cuda, cuda:0, or a transformers device string.')
    parser.add_argument('--depth-anything-repo', default='../diffusion-vas/models/Depth_Anything_V2')
    parser.add_argument('--depth-checkpoint', default=None, help='Local Depth Anything V2 .pth checkpoint.')
    parser.add_argument('--depth-encoder', choices=['vits', 'vitb', 'vitl', 'vitg'], default='vitl')
    parser.add_argument('--metric-depth', action='store_true', help='Use a metric Depth Anything V2 checkpoint and keep metric units.')
    parser.add_argument('--max-depth', type=float, default=20.0, help='Metric depth max range in meters.')
    parser.add_argument('--depth-input-size', type=int, default=518, help='Depth Anything V2 inference input size.')
    parser.add_argument('--vggt-repo', default='vggt', help='Local facebookresearch/VGGT checkout.')
    parser.add_argument('--vggt-model', default='facebook/VGGT-1B', help='Hugging Face model id or local VGGT model directory.')
    parser.add_argument('--vggt-checkpoint', default=None, help='Optional local VGGT model.pt checkpoint.')
    parser.add_argument('--vggt-input-size', type=int, default=518)
    parser.add_argument('--vggt-keep-raw-depth', action='store_true', help='Do not near/far normalize VGGT depth before reconstruction.')
    parser.add_argument('--vggt-point-conf-threshold', type=float, default=1.0)
    parser.add_argument('--vggt-max-point-cloud-points', type=int, default=120000)
    parser.add_argument('--depth-near', type=float, default=1.0)
    parser.add_argument('--depth-far', type=float, default=6.0)
    parser.add_argument('--depth-max-size', type=int, default=1280, help='Resize longest side for local depth generation. Use 0 to keep original size.')
    parser.add_argument('--invert-depth', action='store_true', help='Invert estimated relative depth before normalization.')
    parser.add_argument('--depth-fallback', choices=['synthetic', 'error'], default='synthetic', help='What to do if --estimate-depth fails.')
    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('--sam-refine', action='store_true', help='Refine target/obstacle masks with Segment Anything box prompts.')
    parser.add_argument('--sam-repo', default='../amodal/segment-anything')
    parser.add_argument('--sam-checkpoint', default=None)
    parser.add_argument('--sam-model-type', choices=['vit_h', 'vit_l', 'vit_b', 'default'], default='vit_h')
    parser.add_argument('--sam-device', default='auto')
    parser.add_argument('--sam-fallback', choices=['coarse', 'error'], default='coarse')
    parser.add_argument('--skip-reconstruct', action='store_true')
    return parser


def main():
    args = build_parser().parse_args()
    args.source_preset = args.preset
    args.preset = resolve_reviewed_category(
        args.sample_id or '',
        args.preset,
        load_reviewed_category_overrides(args.category_overrides),
    )
    if args.preset not in ACCESSIBILITY_PRESETS:
        raise ValueError(f'Unsupported reviewed category override: {args.preset}')
    args.category_override_applied = args.preset != args.source_preset
    masks = build_masks(args)
    depth_path, depth_scale, depth_source = prepare_depth(args)
    write_pipeline_manifest(args, masks, depth_source, depth_path)
    run_reconstruct(args, masks, depth_path, depth_scale)
    print(f'Wrote mask outputs to {args.output_dir}')


if __name__ == '__main__':
    main()