tools/accessibility_3d_completion.py

#!/usr/bin/env python3
"""AccessibilityAmodal adapter for the licensed Amodal3R visual 3D backend.

This file contains project-owned input/output orchestration. The model package
imported as ``amodal3d`` remains third-party code under its upstream license.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib
import importlib.util
import json
import math
import os
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

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

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

os.environ["ATTN_BACKEND"] = "xformers"
os.environ["SPARSE_ATTN_BACKEND"] = "xformers"
os.environ["XFORMERS_DISABLED"] = "1"
os.environ["SPCONV_ALGO"] = "native"
os.environ["TORCH_HOME"] = os.environ.get(
    "AMODAL3D_TORCH_HOME", str(PROJECT_ROOT / "weights" / "torch")
)

VISIBLE_VALUE = 188
OCCLUDED_VALUE = 0
BACKGROUND_VALUE = 255
THREE_VALUE_MASK_VALUES = (OCCLUDED_VALUE, VISIBLE_VALUE, BACKGROUND_VALUE)
ACCEPTED_COMPLETION_STATUSES = frozenset({"candidate_selected_for_review"})


def resolve_path(value: str | Path) -> Path:
    path = Path(value).expanduser()
    return path.resolve() if path.is_absolute() else (PROJECT_ROOT / path).resolve()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def portable_file_record(path: Path) -> dict[str, Any]:
    """Record an input without embedding a workstation-specific absolute path."""

    resolved = path.resolve()
    try:
        base = "project_root"
        relative = resolved.relative_to(PROJECT_ROOT).as_posix()
    except ValueError:
        base = "external_input"
        relative = resolved.name
    return {
        "base": base,
        "path": relative,
        "sha256": sha256_file(resolved),
        "bytes": resolved.stat().st_size,
    }


def _declared_selected_path_matches(
    declared: str,
    *,
    completion_manifest: Path,
    completed_image: Path,
) -> bool:
    path = Path(declared).expanduser()
    if path.is_absolute():
        return path.resolve() == completed_image.resolve()
    candidates = (
        completion_manifest.parent / path,
        completion_manifest.parent.parent / path,
    )
    return any(candidate.resolve() == completed_image.resolve() for candidate in candidates)


def _resolve_completion_manifest_artifact(
    declared: str,
    *,
    completion_manifest: Path,
) -> Path:
    path = Path(declared).expanduser()
    candidates = (
        (path,) if path.is_absolute() else (
            completion_manifest.parent / path,
            completion_manifest.parent.parent / path,
        )
    )
    existing = [candidate.resolve() for candidate in candidates if candidate.is_file()]
    if not existing:
        raise FileNotFoundError(
            f"Completion manifest artifact does not exist: {declared}"
        )
    return existing[0]


def validate_completed_rgb_texture_preservation(
    *,
    original: Path,
    completed: Path,
    completion_manifest: Path,
    payload: dict[str, Any],
) -> dict[str, Any]:
    """Prove that a 2D completion preserves source RGB outside its edit mask."""

    files = payload.get("files")
    file_records = files if isinstance(files, dict) else {}
    declared_mask = (
        payload.get("appearance_generation_mask")
        or file_records.get("generation_mask")
    )
    if not isinstance(declared_mask, str):
        raise ValueError(
            "Accepted completed RGB manifest must declare its exact generation "
            "mask as appearance_generation_mask or files.generation_mask"
        )
    generation_mask = _resolve_completion_manifest_artifact(
        declared_mask,
        completion_manifest=completion_manifest,
    )
    with (
        Image.open(original) as original_source,
        Image.open(completed) as completed_source,
        Image.open(generation_mask) as mask_source,
    ):
        original_image = ImageOps.exif_transpose(original_source).convert("RGB")
        completed_image = ImageOps.exif_transpose(completed_source).convert("RGB")
        mask_image = ImageOps.exif_transpose(mask_source).convert("L")
    if original_image.size != completed_image.size:
        raise ValueError(
            "Accepted completed RGB must use the same display raster as the original"
        )
    if mask_image.size != original_image.size:
        raise ValueError(
            "Completion generation mask must align exactly with original/completed RGB"
        )

    original_array = np.asarray(original_image, dtype=np.uint8)
    completed_array = np.asarray(completed_image, dtype=np.uint8)
    generation = np.asarray(mask_image, dtype=np.uint8) > 127
    feather_radius = float(payload.get("feather_radius") or 0.0)
    if not math.isfinite(feather_radius) or feather_radius < 0:
        raise ValueError("Completion manifest feather_radius must be non-negative")
    allowed = generation.astype(np.uint8)
    feather_support_radius = int(math.ceil(3.0 * feather_radius))
    if feather_support_radius > 0:
        kernel = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE,
            (
                feather_support_radius * 2 + 1,
                feather_support_radius * 2 + 1,
            ),
        )
        allowed = cv2.dilate(allowed, kernel)
    allowed = allowed > 0
    changed = np.any(original_array != completed_array, axis=2)
    changed_outside = changed & ~allowed
    changed_inside = changed & allowed
    changed_outside_count = int(changed_outside.sum())
    if changed_outside_count:
        raise ValueError(
            "Completed RGB changes source texture outside its declared generation "
            f"mask/feather support at {changed_outside_count} pixels"
        )
    changed_inside_count = int(changed_inside.sum())
    if changed_inside_count == 0:
        raise ValueError(
            "Completed RGB does not change any pixel inside its declared edit region"
        )
    outside_count = int((~allowed).sum())
    return {
        "validated": True,
        "generation_mask": portable_file_record(generation_mask),
        "feather_radius": feather_radius,
        "feather_support_radius_pixels": feather_support_radius,
        "allowed_edit_pixel_count": int(allowed.sum()),
        "changed_inside_allowed_region_pixel_count": changed_inside_count,
        "changed_outside_allowed_region_pixel_count": 0,
        "source_rgb_identity_outside_allowed_region": True,
        "source_rgb_identity_outside_allowed_region_ratio": (
            1.0 if outside_count else None
        ),
    }


def select_backend_rgb(args: argparse.Namespace) -> tuple[Path, dict[str, Any]]:
    """Choose the explicitly requested RGB conditioning source.

    The original image remains the canonical geometry/provenance source.
    Completion arguments cannot replace it unless ``--conditioning-rgb
    completed`` is explicit and the completion manifest passes the quality
    gate.
    """

    original = resolve_path(args.image)
    if not original.is_file():
        raise FileNotFoundError(original)
    conditioning_rgb = str(getattr(args, "conditioning_rgb", "original"))
    completed_value = getattr(args, "completed_image", None)
    manifest_value = getattr(args, "completion_manifest", None)
    completion_inputs_supplied = bool(completed_value or manifest_value)

    if conditioning_rgb == "original":
        original_record = portable_file_record(original)
        return original, {
            "conditioning_rgb_mode": "original",
            "backend_rgb_role": "original_rgb_no_accepted_2d_completion",
            "backend_rgb": original_record,
            "original_rgb": original_record,
            "selected_completion_rgb": None,
            "completion_manifest": None,
            "completion_status": None,
            "quality_gate_accepted": False,
            "completion_inputs_supplied": completion_inputs_supplied,
            "completion_inputs_ignored": completion_inputs_supplied,
            "original_vs_completed": {
                "same_raster_size": None,
                "same_sha256": None,
            },
        }

    if conditioning_rgb != "completed":
        raise ValueError(
            "--conditioning-rgb must be either 'original' or 'completed'"
        )
    if not completed_value or not manifest_value:
        raise ValueError(
            "--conditioning-rgb completed requires both --completed-image and "
            "--completion-manifest"
        )

    completed = resolve_path(completed_value)
    completion_manifest = resolve_path(manifest_value)
    if not completed.is_file():
        raise FileNotFoundError(completed)
    if not completion_manifest.is_file():
        raise FileNotFoundError(completion_manifest)
    payload = json.loads(completion_manifest.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("Completion manifest must contain a JSON object")
    status = str(payload.get("status") or "")
    if status not in ACCEPTED_COMPLETION_STATUSES:
        raise ValueError(
            "Refusing completed RGB for learned visual 3D because completion "
            f"status is not accepted: {status or '<missing>'}"
        )
    files = payload.get("files")
    declared_from_files = (
        files.get("selected_completed_rgb") if isinstance(files, dict) else None
    )
    declared = payload.get("selected_completed_rgb") or declared_from_files
    if not isinstance(declared, str) or not _declared_selected_path_matches(
        declared,
        completion_manifest=completion_manifest,
        completed_image=completed,
    ):
        raise ValueError(
            "Completion manifest does not identify --completed-image as its "
            "selected completion"
        )

    texture_preservation = validate_completed_rgb_texture_preservation(
        original=original,
        completed=completed,
        completion_manifest=completion_manifest,
        payload=payload,
    )
    original_record = portable_file_record(original)
    completed_record = portable_file_record(completed)
    return completed, {
        "conditioning_rgb_mode": "completed",
        "backend_rgb_role": (
            "quality_gate_accepted_obstacle_removed_rgb_preserving_original_texture"
        ),
        "backend_rgb": completed_record,
        "original_rgb": original_record,
        "selected_completion_rgb": completed_record,
        "completion_manifest": portable_file_record(completion_manifest),
        "completion_status": status,
        "completion_manifest_selected_rgb": Path(declared).name,
        "quality_gate_accepted": True,
        "completion_inputs_supplied": True,
        "completion_inputs_ignored": False,
        "original_vs_completed": {
            "same_raster_size": True,
            "same_sha256": original_record["sha256"] == completed_record["sha256"],
        },
        "original_texture_preservation": texture_preservation,
    }


def load_backend_runtime():
    """Import the licensed third-party runtime only for an actual GPU run."""

    from amodal3d.pipelines import Amodal3RImageTo3DPipeline
    from amodal3d.utils import render_utils

    return Amodal3RImageTo3DPipeline, render_utils


def extract_glb(gs, mesh, mesh_simplify=0.95, texture_size=1024, export_path="output.glb"):
    from amodal3d.utils import postprocessing_utils

    glb = postprocessing_utils.to_glb(
        gs,
        mesh,
        simplify=mesh_simplify,
        texture_size=texture_size,
        verbose=False,
    )
    glb.export(export_path)
    return export_path


def save_mesh(mesh_result, filename):
    vertices = (
        mesh_result.vertices.cpu().numpy()
        if hasattr(mesh_result.vertices, "cpu")
        else mesh_result.vertices
    )
    faces = (
        mesh_result.faces.cpu().numpy()
        if hasattr(mesh_result.faces, "cpu")
        else mesh_result.faces
    )
    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
    if mesh_result.vertex_attrs is not None:
        attrs = (
            mesh_result.vertex_attrs.cpu().numpy()
            if hasattr(mesh_result.vertex_attrs, "cpu")
            else mesh_result.vertex_attrs
        )
        mesh.visual.vertex_colors = attrs
    mesh.export(filename)


def parse_box(box):
    values = [float(value) for value in box.split(",")]
    if len(values) != 4:
        raise argparse.ArgumentTypeError("box must be x1,y1,x2,y2")
    if any(value < 0 or value > 1 for value in values):
        raise argparse.ArgumentTypeError("box coordinates must be normalized to [0, 1]")
    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 make_stair_scene_mask(image, occlusion_boxes=None, save_path=None):
    width, height = image.size
    mask = np.full((height, width), VISIBLE_VALUE, dtype=np.uint8)
    mask[: int(height * 0.08), :] = BACKGROUND_VALUE
    for x1, y1, x2, y2 in occlusion_boxes or []:
        left = int(round(x1 * width))
        top = int(round(y1 * height))
        right = int(round(x2 * width))
        bottom = int(round(y2 * height))
        mask[top:bottom, left:right] = OCCLUDED_VALUE
    result = Image.fromarray(mask, mode="L")
    if save_path is not None:
        result.save(save_path)
    return result


def load_three_value_condition_mask(
    path: str | Path,
    *,
    expected_size: tuple[int, int],
) -> tuple[Image.Image, Path, dict[str, Any]]:
    """Load and strictly validate an aligned Amodal3R condition PNG."""

    resolved = resolve_path(path)
    if not resolved.is_file():
        raise FileNotFoundError(resolved)
    with Image.open(resolved) as source:
        image_format = source.format
        condition = ImageOps.exif_transpose(source)
        if image_format != "PNG":
            raise ValueError(
                f"Condition mask must be a PNG file; detected {image_format or 'unknown'}"
            )
        if condition.mode != "L":
            raise ValueError(
                "Condition mask PNG must be single-channel 8-bit grayscale "
                f"(mode L); got mode {condition.mode}"
            )
        if condition.size != expected_size:
            raise ValueError(
                "Condition mask/RGB raster mismatch: "
                f"mask={condition.size}, rgb={expected_size}. Refusing to resize."
            )
        values = np.asarray(condition, dtype=np.uint8).copy()

    observed_values, observed_counts = np.unique(values, return_counts=True)
    observed = {
        int(value): int(count)
        for value, count in zip(observed_values.tolist(), observed_counts.tolist())
    }
    invalid_values = sorted(set(observed) - set(THREE_VALUE_MASK_VALUES))
    if invalid_values:
        raise ValueError(
            "Condition mask PNG contains invalid pixel values "
            f"{invalid_values}; allowed exact values are "
            f"{list(THREE_VALUE_MASK_VALUES)}"
        )
    visible_count = observed.get(VISIBLE_VALUE, 0)
    if visible_count == 0:
        raise ValueError(
            "Condition mask PNG must contain at least one visible-target pixel "
            f"with value {VISIBLE_VALUE}"
        )

    width, height = expected_size
    statistics = {
        "strict_three_value_validation": True,
        "allowed_values": list(THREE_VALUE_MASK_VALUES),
        "observed_values": sorted(observed),
        "width": int(width),
        "height": int(height),
        "total_pixel_count": int(values.size),
        "hidden_pixel_count": observed.get(OCCLUDED_VALUE, 0),
        "visible_pixel_count": visible_count,
        "background_pixel_count": observed.get(BACKGROUND_VALUE, 0),
        "hidden_region_present": observed.get(OCCLUDED_VALUE, 0) > 0,
    }
    return Image.fromarray(values, mode="L"), resolved, statistics


def square_focus_crop(
    image: Image.Image,
    mask: Image.Image,
    *,
    padding_ratio: float,
) -> tuple[Image.Image, Image.Image, dict[str, Any]]:
    """Crop/pad aligned RGB and mask around the target without distortion."""

    if not math.isfinite(padding_ratio) or padding_ratio < 0:
        raise ValueError("--focus-crop-padding-ratio must be non-negative")
    if image.size != mask.size:
        raise ValueError("Focus crop requires aligned RGB and condition mask")

    values = np.asarray(mask, dtype=np.uint8)
    target_y, target_x = np.nonzero(values != BACKGROUND_VALUE)
    if target_x.size == 0:
        raise ValueError(
            "Focus crop requires at least one non-background target pixel"
        )

    bbox_left = int(target_x.min())
    bbox_top = int(target_y.min())
    bbox_right = int(target_x.max()) + 1
    bbox_bottom = int(target_y.max()) + 1
    bbox_width = bbox_right - bbox_left
    bbox_height = bbox_bottom - bbox_top
    side = max(
        1,
        int(
            math.ceil(
                max(bbox_width, bbox_height)
                * (1.0 + 2.0 * padding_ratio)
            )
        ),
    )
    center_x = 0.5 * (bbox_left + bbox_right)
    center_y = 0.5 * (bbox_top + bbox_bottom)
    crop_left = int(math.floor(center_x - 0.5 * side))
    crop_top = int(math.floor(center_y - 0.5 * side))
    crop_right = crop_left + side
    crop_bottom = crop_top + side

    source_width, source_height = image.size
    source_left = max(crop_left, 0)
    source_top = max(crop_top, 0)
    source_right = min(crop_right, source_width)
    source_bottom = min(crop_bottom, source_height)
    paste_left = source_left - crop_left
    paste_top = source_top - crop_top

    focused_image = Image.new("RGB", (side, side), (0, 0, 0))
    focused_mask = Image.new("L", (side, side), BACKGROUND_VALUE)
    source_box = (source_left, source_top, source_right, source_bottom)
    focused_image.paste(image.crop(source_box), (paste_left, paste_top))
    focused_mask.paste(mask.crop(source_box), (paste_left, paste_top))

    focused_values = np.asarray(focused_mask, dtype=np.uint8)
    observed_values, observed_counts = np.unique(
        focused_values, return_counts=True
    )
    observed = {
        int(value): int(count)
        for value, count in zip(
            observed_values.tolist(), observed_counts.tolist()
        )
    }
    metadata = {
        "applied": True,
        "policy": "square_target_bbox_crop_with_background_padding",
        "padding_ratio": float(padding_ratio),
        "source_size": [int(source_width), int(source_height)],
        "target_bbox_xyxy": [
            bbox_left,
            bbox_top,
            bbox_right,
            bbox_bottom,
        ],
        "crop_box_xyxy": [crop_left, crop_top, crop_right, crop_bottom],
        "source_intersection_xyxy": [
            source_left,
            source_top,
            source_right,
            source_bottom,
        ],
        "output_size": [side, side],
        "observed_values": sorted(observed),
        "hidden_pixel_count": observed.get(OCCLUDED_VALUE, 0),
        "visible_pixel_count": observed.get(VISIBLE_VALUE, 0),
        "background_pixel_count": observed.get(BACKGROUND_VALUE, 0),
    }
    return focused_image, focused_mask, metadata


def load_inputs(args, output_dir):
    backend_rgb_path, input_provenance = select_backend_rgb(args)
    with Image.open(backend_rgb_path) as source:
        image = ImageOps.exif_transpose(source).convert("RGB")
    if args.mask:
        mask_source = resolve_path(args.mask)
    else:
        mask_path = output_dir / "auto_stair_mask.png"
        make_stair_scene_mask(image, args.occlusion_box, mask_path)
        mask_source = mask_path.resolve()
    mask, mask_source, mask_statistics = load_three_value_condition_mask(
        mask_source,
        expected_size=image.size,
    )
    focus_crop_padding_ratio = getattr(
        args, "focus_crop_padding_ratio", None
    )
    if focus_crop_padding_ratio is not None:
        source_mask_statistics = dict(mask_statistics)
        image, mask, focus_crop = square_focus_crop(
            image,
            mask,
            padding_ratio=float(focus_crop_padding_ratio),
        )
        mask_statistics = {
            **mask_statistics,
            "width": int(mask.width),
            "height": int(mask.height),
            "total_pixel_count": int(mask.width * mask.height),
            "observed_values": focus_crop["observed_values"],
            "hidden_pixel_count": focus_crop["hidden_pixel_count"],
            "visible_pixel_count": focus_crop["visible_pixel_count"],
            "background_pixel_count": focus_crop["background_pixel_count"],
            "hidden_region_present": focus_crop["hidden_pixel_count"] > 0,
            "focus_crop": focus_crop,
            "pre_focus_crop_statistics": source_mask_statistics,
        }
    return (
        image,
        mask,
        mask_source,
        mask_statistics,
        backend_rgb_path,
        input_provenance,
    )


def require_cuda(torch_module: Any | None = None) -> dict[str, Any]:
    """Fail before model loading unless the process has a usable CUDA allocation."""

    if torch_module is None:
        try:
            import torch as torch_module
        except ImportError as exc:
            raise RuntimeError(
                "Amodal3R visual 3D requires PyTorch with CUDA support"
            ) from exc
    cuda = getattr(torch_module, "cuda", None)
    if cuda is None or not cuda.is_available():
        raise RuntimeError(
            "Amodal3R visual 3D requires an available CUDA GPU for learned "
            "generation and diff_gaussian_rasterization. Run this command inside "
            "a Slurm GPU allocation."
        )
    device_count = getattr(cuda, "device_count", lambda: 1)()
    torch_version = getattr(torch_module, "__version__", None)
    version_namespace = getattr(torch_module, "version", None)
    return {
        "required": True,
        "available": True,
        "device_type": "cuda",
        "device_count": int(device_count),
        "torch_version": str(torch_version) if torch_version is not None else None,
        "torch_cuda_build": (
            str(getattr(version_namespace, "cuda"))
            if version_namespace is not None
            and getattr(version_namespace, "cuda", None) is not None
            else None
        ),
    }


def require_gpu_renderers(
    *,
    allow_gaussian_only: bool,
    find_spec=importlib.util.find_spec,
    import_module=importlib.import_module,
) -> dict[str, Any]:
    """Preflight the CUDA rasterizers before loading the Amodal3R weights."""

    gaussian_available = find_spec("diff_gaussian_rasterization") is not None
    if not gaussian_available:
        raise RuntimeError(
            "Amodal3R requires diff_gaussian_rasterization for the primary "
            "CUDA rotating render."
        )
    try:
        import_module("diff_gaussian_rasterization")
    except Exception as exc:
        raise RuntimeError(
            "diff_gaussian_rasterization is installed but its CUDA extension "
            "could not be imported."
        ) from exc
    mesh_available = find_spec("nvdiffrast") is not None
    if not mesh_available and not allow_gaussian_only:
        raise RuntimeError(
            "Full GPU rendering requires nvdiffrast for the dense FlexiCubes "
            "mesh rotation. Install the project --nvdiffrast dependency, or use "
            "--allow-gaussian-only only for an explicit debug run."
        )
    mesh_context_ready = False
    if mesh_available:
        try:
            dr = import_module("nvdiffrast.torch")
            context = dr.RasterizeCudaContext(device="cuda")
            del context
            mesh_context_ready = True
        except Exception as exc:
            raise RuntimeError(
                "nvdiffrast is installed but its CUDA raster context could not "
                "be created on the allocated GPU."
            ) from exc
    return {
        "gaussian": {
            "package": "diff_gaussian_rasterization",
            "available": gaussian_available,
            "required": True,
            "device": "cuda",
            "runtime_import_succeeded": True,
        },
        "dense_mesh": {
            "package": "nvdiffrast",
            "available": mesh_available,
            "required": not allow_gaussian_only,
            "device": "cuda" if mesh_available else None,
            "runtime_import_succeeded": mesh_available,
            "cuda_context_preflight_succeeded": mesh_context_ready,
        },
        "cpu_render_fallback_allowed": False,
        "gaussian_only_debug_mode": bool(allow_gaussian_only),
    }


def write_gif_atomic(path: Path, frames: list[np.ndarray], *, fps: int) -> None:
    temporary = path.with_name(f".{path.stem}.{os.getpid()}.tmp.gif")
    try:
        imageio.mimsave(
            temporary,
            frames,
            duration=1000.0 / float(fps),
            loop=0,
        )
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def copy_file_atomic(source: Path, destination: Path) -> None:
    """Atomically materialize a byte-identical compatibility artifact."""

    temporary = destination.with_name(
        f".{destination.stem}.{os.getpid()}.tmp{destination.suffix}"
    )
    try:
        shutil.copyfile(source, temporary)
        os.replace(temporary, destination)
    finally:
        temporary.unlink(missing_ok=True)


def inspect_gif(path: Path) -> dict[str, Any]:
    with Image.open(path) as image:
        return {
            "path": path.name,
            "frames": int(getattr(image, "n_frames", 1)),
            "width": int(image.width),
            "height": int(image.height),
        }


def validate_render_artifacts(
    *,
    output_dir: Path,
    video_frames: int,
    nviews: int,
    dense_mesh_rendered: bool,
    mesh_face_count: int,
) -> dict[str, Any]:
    """Validate the full-GPU publication contract before writing manifest.json."""

    gaussian = inspect_gif(output_dir / "sample_gaussian.gif")
    combined = inspect_gif(output_dir / "sample_multi.gif")
    if gaussian["frames"] != video_frames or combined["frames"] != video_frames:
        raise RuntimeError(
            "GPU GIF frame-count mismatch: "
            f"gaussian={gaussian['frames']}, combined={combined['frames']}, "
            f"expected={video_frames}"
        )
    # The 9527 reference contract uses sample_multi.gif as a compatibility
    # alias for the source-colored Gaussian render.  Dense-mesh normal maps are
    # diagnostics and must never widen or recolor the primary presentation.
    expected_combined_width = gaussian["width"]
    if (
        combined["width"] != expected_combined_width
        or combined["height"] != gaussian["height"]
    ):
        raise RuntimeError(
            "Reference-compatible GIF dimensions do not match the Gaussian render"
        )
    if sha256_file(output_dir / "sample_gaussian.gif") != sha256_file(
        output_dir / "sample_multi.gif"
    ):
        raise RuntimeError(
            "sample_multi.gif must be a byte-identical compatibility alias of "
            "the source-colored sample_gaussian.gif"
        )
    gaussian_views = [
        output_dir / f"{index:03d}_gs.png" for index in range(nviews)
    ]
    if not all(path.is_file() and path.stat().st_size > 0 for path in gaussian_views):
        raise RuntimeError("GPU Gaussian multiview output is incomplete")

    mesh_gif = None
    mesh_views: list[Path] = []
    if dense_mesh_rendered:
        mesh_gif = inspect_gif(output_dir / "sample_mesh.gif")
        if (
            mesh_gif["frames"] != video_frames
            or mesh_gif["width"] != gaussian["width"]
            or mesh_gif["height"] != gaussian["height"]
        ):
            raise RuntimeError("GPU dense-mesh GIF does not align with Gaussian GIF")
        mesh_views = [
            output_dir / f"{index:03d}_mesh.png" for index in range(nviews)
        ]
        if not all(path.is_file() and path.stat().st_size > 0 for path in mesh_views):
            raise RuntimeError("GPU dense-mesh multiview output is incomplete")
    if mesh_face_count <= 0:
        raise RuntimeError("FlexiCubes output contains no triangle faces")
    contact_sheet = output_dir / "multiview_contact_sheet.jpg"
    if not contact_sheet.is_file() or contact_sheet.stat().st_size <= 0:
        raise RuntimeError("GPU multiview contact sheet is missing")
    with Image.open(contact_sheet) as sheet:
        expected_sheet_size = (
            gaussian["width"] * 2,
            gaussian["height"] * ((nviews + 1) // 2),
        )
        if sheet.size != expected_sheet_size:
            raise RuntimeError(
                "Primary contact sheet must use only source-colored Gaussian "
                f"views in a 2-column layout: got={sheet.size}, "
                f"expected={expected_sheet_size}"
            )
    return {
        "validated": True,
        "gaussian_gif": gaussian,
        "mesh_gif": mesh_gif,
        "combined_gif": combined,
        "combined_gif_is_byte_identical_gaussian_alias": True,
        "contact_sheet_layout": "gaussian_color_only_2_columns",
        "gaussian_view_count": len(gaussian_views),
        "mesh_view_count": len(mesh_views),
        "dense_mesh_rendered_on_gpu": dense_mesh_rendered,
        "mesh_face_count": int(mesh_face_count),
        "cpu_render_fallback_used": False,
    }


def build_output_manifest(
    *,
    args: argparse.Namespace,
    input_provenance: dict[str, Any],
    mask_source: str | Path,
    mask_statistics: dict[str, Any],
    cuda_runtime: dict[str, Any],
    gpu_renderer_runtime: dict[str, Any],
    render_validation: dict[str, Any],
    mesh_gif: str | None,
    nvdiffrast_available: bool,
    mesh_renderer_mode: str,
    glb_export_mode: str | None,
) -> dict[str, Any]:
    """Build the portable visual-candidate manifest and evidence policy."""

    return {
        "schema_version": "accessibilityamodal_visual_3d_candidate_v1",
        "created_at_utc": datetime.now(timezone.utc).isoformat(),
        "input_provenance": input_provenance,
        "three_value_mask": portable_file_record(resolve_path(mask_source)),
        "three_value_mask_statistics": dict(mask_statistics),
        "third_party_backend_provenance": {
            "name": "Amodal3R",
            "python_package": "amodal3d",
            "pipeline_class": "Amodal3RImageTo3DPipeline",
            "model_identifier_or_path": args.model,
            "upstream_identity_preserved": True,
        },
        "representation_contract": {
            "primary_representation": "Amodal3R Gaussian",
            "primary_render_backend": "CUDA diff_gaussian_rasterization",
            "primary_render_device": "cuda",
            "primary_output": "sample_gaussian.gif",
            "reference_compatible_output": "sample_multi.gif",
            "reference_compatible_output_matches_primary": True,
            "primary_appearance": (
                "quality_gated_obstacle_removed_color_preserving_original_texture"
                if input_provenance["conditioning_rgb_mode"] == "completed"
                else "source_rgb_conditioned_color"
            ),
            "appearance_conditioning_mode": input_provenance[
                "conditioning_rgb_mode"
            ],
            "primary_contact_sheet": "multiview_contact_sheet.jpg",
            "primary_contact_sheet_content": "gaussian_color_only",
            "paired_dense_geometry": "FlexiCubes triangle mesh",
            "paired_dense_geometry_output": "mesh.ply",
            "paired_dense_geometry_render": (
                "sample_mesh.gif" if nvdiffrast_available else None
            ),
            "paired_dense_geometry_render_role": (
                "diagnostic_normal_map_only_not_source_rgb_texture"
                if nvdiffrast_available
                else None
            ),
            "interactive_textured_surface": (
                "mesh.glb" if args.export_glb else None
            ),
            "interactive_texture_source": (
                "GPU Gaussian appearance baked to UV/PBR base-color texture"
                if args.export_glb
                else None
            ),
            "vggt_is_primary": False,
            "discrete_point_cloud_is_primary": False,
        },
        "cuda_runtime": dict(cuda_runtime),
        "gpu_renderer_runtime": dict(gpu_renderer_runtime),
        "render_validation": dict(render_validation),
        "seed": args.seed,
        "nviews": args.nviews,
        "video_frames": args.video_frames,
        "output_kind": (
            "AccessibilityAmodal learned visual candidate via the licensed "
            "third-party Amodal3R backend"
        ),
        "outputs": {
            "combined_gif": "sample_multi.gif",
            "gaussian_gif": "sample_gaussian.gif",
            "mesh_gif": mesh_gif,
            "mesh_normal_diagnostic_gif": mesh_gif,
            "contact_sheet": "multiview_contact_sheet.jpg",
            "mesh": "mesh.ply",
            "glb": "mesh.glb" if args.export_glb else None,
            "conditioning_rgb_model_input": "conditioning_rgb_model_input.png",
            "condition_mask_model_input": "condition_mask_model_input.png",
        },
        "nvdiffrast_available": nvdiffrast_available,
        "mesh_renderer_mode": mesh_renderer_mode,
        "glb_export_mode": glb_export_mode,
        "metric_geometry": False,
        "passability_evidence": False,
        "automatic_passability_claim": False,
        "human_review_required": True,
        "evidence_policy": {
            "role": "learned_visual_candidate_only",
            "metric_geometry": False,
            "passability_evidence": False,
            "automatic_passability_claim": False,
            "warning": (
                "This learned visual 3D candidate is not calibrated geometry and "
                "must not be used to decide whether a person can pass."
            ),
        },
    }


def run(args):
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    # manifest.json is the success marker. Remove stale primary markers before a
    # rerun so a failed GPU mesh pass cannot be mistaken for a complete result.
    for filename in (
        "manifest.json",
        "sample_multi.gif",
        "sample_mesh.gif",
        "mesh.glb",
    ):
        (output_dir / filename).unlink(missing_ok=True)
    (
        image,
        mask,
        mask_source,
        mask_statistics,
        backend_rgb_path,
        input_provenance,
    ) = load_inputs(args, output_dir)
    image.save(output_dir / "conditioning_rgb_model_input.png")
    mask.save(output_dir / "condition_mask_model_input.png")
    cuda_runtime = require_cuda()
    gpu_renderer_runtime = require_gpu_renderers(
        allow_gaussian_only=args.allow_gaussian_only,
    )
    nvdiffrast_available = bool(
        gpu_renderer_runtime["dense_mesh"]["available"]
    )
    print(f"Original image: {args.image}")
    print(f"Learned visual backend RGB: {backend_rgb_path}")
    print(f"Backend RGB role: {input_provenance['backend_rgb_role']}")
    if input_provenance["completion_inputs_ignored"]:
        print(
            "Completion RGB arguments were supplied but ignored because "
            "--conditioning-rgb defaults to original."
        )
    print(f"Mask: {mask_source}")
    print(f"Output dir: {output_dir}")

    pipeline_class, render_utils = load_backend_runtime()
    pipeline = pipeline_class.from_pretrained(args.model)
    pipeline.cuda()
    outputs = pipeline.run_multi_image(
        [image],
        [mask],
        seed=args.seed,
        sparse_structure_sampler_params={"steps": args.ss_steps, "cfg_strength": args.ss_cfg},
        slat_sampler_params={"steps": args.slat_steps, "cfg_strength": args.slat_cfg},
        erode_kernel_size=args.erode_kernel_size,
    )

    video_gs = render_utils.render_video(
        outputs["gaussian"][0],
        bg_color=(1, 1, 1),
        num_frames=args.video_frames,
    )["color"]
    write_gif_atomic(output_dir / "sample_gaussian.gif", video_gs, fps=24)
    gaussian = outputs["gaussian"][0]
    multi_view_gs, _, _ = render_utils.render_multiview(
        gaussian, nviews=args.nviews, bg_color=(1, 1, 1)
    )
    mesh = outputs["mesh"][0]
    for index, output in enumerate(multi_view_gs["color"]):
        output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
        cv2.imwrite(str(output_dir / f"{index:03d}_gs.png"), output)
    mesh_path = output_dir / "mesh.ply"
    save_mesh(mesh, mesh_path)
    mesh_face_count = int(mesh.faces.shape[0])

    mesh_gif: str | None = None
    mesh_renderer_mode = "skipped: explicit Gaussian-only debug mode"
    if nvdiffrast_available:
        video_mesh = render_utils.render_video(
            mesh,
            bg_color=(1, 1, 1),
            num_frames=args.video_frames,
        )["normal"]
        write_gif_atomic(output_dir / "sample_mesh.gif", video_mesh, fps=24)
        copy_file_atomic(
            output_dir / "sample_gaussian.gif",
            output_dir / "sample_multi.gif",
        )
        multi_view_mesh, _, _ = render_utils.render_multiview(
            mesh, nviews=args.nviews, bg_color=(1, 1, 1)
        )
        for index, output in enumerate(multi_view_mesh["normal"]):
            output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
            cv2.imwrite(str(output_dir / f"{index:03d}_mesh.png"), output)
        previews = list(multi_view_gs["color"])
        mesh_gif = "sample_mesh.gif"
        mesh_renderer_mode = "nvdiffrast normal-map diagnostic rendering"
    else:
        previews = list(multi_view_gs["color"])
        copy_file_atomic(
            output_dir / "sample_gaussian.gif",
            output_dir / "sample_multi.gif",
        )

    rows = []
    for start in range(0, len(previews), 2):
        row = previews[start : start + 2]
        if len(row) == 1:
            row.append(np.full_like(row[0], 255))
        rows.append(np.concatenate(row, axis=1))
    contact_sheet = np.concatenate(rows, axis=0)
    contact_sheet_path = output_dir / "multiview_contact_sheet.jpg"
    temporary_contact_sheet = contact_sheet_path.with_name(
        f".{contact_sheet_path.stem}.{os.getpid()}.tmp.jpg"
    )
    try:
        Image.fromarray(contact_sheet).save(temporary_contact_sheet, quality=92)
        os.replace(temporary_contact_sheet, contact_sheet_path)
    finally:
        temporary_contact_sheet.unlink(missing_ok=True)
    glb_export_mode = None
    if args.export_glb:
        if nvdiffrast_available:
            extract_glb(
                outputs["gaussian"][0],
                outputs["mesh"][0],
                mesh_simplify=args.mesh_simplify,
                texture_size=args.texture_size,
                export_path=str(output_dir / "mesh.glb"),
            )
            glb_export_mode = "Amodal3R textured GLB"
        else:
            raise RuntimeError(
                "--export-glb requires nvdiffrast; CPU GLB fallback is disabled "
                "by the full-GPU rendering contract."
            )
    render_validation = validate_render_artifacts(
        output_dir=output_dir,
        video_frames=args.video_frames,
        nviews=args.nviews,
        dense_mesh_rendered=nvdiffrast_available,
        mesh_face_count=mesh_face_count,
    )
    manifest = build_output_manifest(
        args=args,
        input_provenance=input_provenance,
        mask_source=mask_source,
        mask_statistics=mask_statistics,
        cuda_runtime=cuda_runtime,
        gpu_renderer_runtime=gpu_renderer_runtime,
        render_validation=render_validation,
        mesh_gif=mesh_gif,
        nvdiffrast_available=nvdiffrast_available,
        mesh_renderer_mode=mesh_renderer_mode,
        glb_export_mode=glb_export_mode,
    )
    (output_dir / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print("Done.")


def build_parser():
    parser = argparse.ArgumentParser(
        description="Run the AccessibilityAmodal visual-3D backend adapter."
    )
    parser.add_argument(
        "--image",
        required=True,
        help=(
            "Canonical original RGB. It remains provenance/geometry source even "
            "when an accepted selected 2D completion drives this visual backend."
        ),
    )
    parser.add_argument(
        "--conditioning-rgb",
        choices=("original", "completed"),
        default="original",
        help=(
            "RGB used to condition Amodal3R. Defaults to the canonical original. "
            "Use 'completed' explicitly to opt in to a quality-gated completion."
        ),
    )
    parser.add_argument(
        "--completed-image",
        default=None,
        help=(
            "Selected 2D completion used only with --conditioning-rgb completed. "
            "Requires an accepted --completion-manifest."
        ),
    )
    parser.add_argument(
        "--completion-manifest",
        default=None,
        help=(
            "2D manifest proving --completed-image is the selected candidate. "
            "Required only with --conditioning-rgb completed."
        ),
    )
    parser.add_argument(
        "--mask",
        default=None,
        help="Optional three-value mask: white background, gray visible, black occluded.",
    )
    parser.add_argument(
        "--focus-crop-padding-ratio",
        type=float,
        default=None,
        help=(
            "Optionally crop/pad RGB and mask to a square around the modeled "
            "target before Amodal3R's 518x518 resize. This preserves source "
            "aspect ratio and makes small target surfaces more prominent."
        ),
    )
    parser.add_argument("--output-dir", default="./output/accessibilityamodal/visual_candidate")
    parser.add_argument("--model", default="Sm0kyWu/Amodal3R")
    parser.add_argument("--occlusion-box", action="append", type=parse_box, default=[])
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--ss-steps", type=int, default=12)
    parser.add_argument("--ss-cfg", type=float, default=7.5)
    parser.add_argument("--slat-steps", type=int, default=12)
    parser.add_argument("--slat-cfg", type=float, default=3.0)
    parser.add_argument("--erode-kernel-size", type=int, default=3)
    parser.add_argument("--nviews", type=int, default=8)
    parser.add_argument("--video-frames", type=int, default=120)
    parser.add_argument(
        "--allow-gaussian-only",
        action="store_true",
        help=(
            "Explicit debug escape hatch when nvdiffrast is unavailable. "
            "The default requires both CUDA Gaussian and CUDA dense-mesh renders."
        ),
    )
    parser.add_argument("--export-glb", action="store_true")
    parser.add_argument("--mesh-simplify", type=float, default=0.5)
    parser.add_argument("--texture-size", type=int, default=1024)
    return parser


if __name__ == "__main__":
    run(build_parser().parse_args())