tools/accessibility_2d_completion.py

#!/usr/bin/env python3
"""GPU generative 2D completion for accessibility-scene occlusions.

The module is a project-owned orchestration layer.  It uses a locally supplied
Diffusers inpainting checkpoint as a backend, preserves all pixels outside the
reviewed completion mask, ranks several candidates using staircase-oriented
image evidence, and publishes one deterministic selected result.
"""

from __future__ import annotations

import argparse
import json
import math
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFilter, ImageOps


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

from accessibilityamodal.visual_completion import (  # noqa: E402
    NEGATIVE_PROMPT,
    PROMPTS,
    apply_selection_clutter_penalty,
    build_completion_envelope,
    build_visual_removal_mask,
    candidate_clutter_metrics,
    candidate_quality,
    derive_hidden_mask,
    mask_statistics,
    select_candidate,
)


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


def load_rgb(path: Path) -> Image.Image:
    """Read RGB in the display orientation used to generate the proposal masks."""
    return ImageOps.exif_transpose(Image.open(path)).convert("RGB")


def load_binary_mask(path: Path, size: tuple[int, int]) -> Image.Image:
    mask = ImageOps.exif_transpose(Image.open(path)).convert("L")
    if mask.size != size:
        raise ValueError(
            f"Mask/RGB raster mismatch for {path}: mask={mask.size}, rgb={size}. "
            "Refusing to resize because this can hide EXIF-orientation misalignment."
        )
    return Image.fromarray((np.asarray(mask) > 127).astype(np.uint8) * 255, mode="L")


def portable_input_reference(path: Path) -> dict[str, str]:
    """Describe an input without embedding a machine-specific absolute path."""

    try:
        return {
            "base": "project_root",
            "path": path.resolve().relative_to(PROJECT_ROOT).as_posix(),
        }
    except ValueError:
        return {"base": "external_input", "path": path.name}


def load_completion_masks(
    *,
    size: tuple[int, int],
    hidden_path: Path | None,
    obstacle_path: Path | None,
    target_amodal_path: Path | None,
    target_visible_path: Path | None,
) -> tuple[Image.Image, Image.Image, dict[str, Any]]:
    """Load geometry hidden and build the separate visual removal mask."""

    if bool(target_amodal_path) != bool(target_visible_path):
        raise ValueError(
            "--target-amodal-mask and --target-visible-mask must be provided together"
        )
    if hidden_path is None and target_amodal_path is None:
        raise ValueError(
            "Provide --mask, or provide both --target-amodal-mask and --target-visible-mask"
        )

    hidden_from_file = (
        np.asarray(load_binary_mask(hidden_path, size)) > 127
        if hidden_path is not None
        else None
    )
    target_amodal = (
        np.asarray(load_binary_mask(target_amodal_path, size)) > 127
        if target_amodal_path is not None
        else None
    )
    target_visible = (
        np.asarray(load_binary_mask(target_visible_path, size)) > 127
        if target_visible_path is not None
        else None
    )
    hidden_from_targets = (
        derive_hidden_mask(target_amodal, target_visible)
        if target_amodal is not None and target_visible is not None
        else None
    )
    if (
        hidden_from_file is not None
        and hidden_from_targets is not None
        and not np.array_equal(hidden_from_file, hidden_from_targets)
    ):
        mismatch = int(np.count_nonzero(hidden_from_file ^ hidden_from_targets))
        raise ValueError(
            "--mask disagrees with target_amodal AND NOT target_visible "
            f"at {mismatch} pixels"
        )

    geometry_hidden = (
        hidden_from_targets if hidden_from_targets is not None else hidden_from_file
    )
    assert geometry_hidden is not None
    obstacle = (
        np.asarray(load_binary_mask(obstacle_path, size)) > 127
        if obstacle_path is not None
        else None
    )
    visual_removal, policy_stats = build_visual_removal_mask(
        geometry_hidden,
        obstacle,
    )
    report: dict[str, Any] = {
        "hidden_source": (
            "target_amodal_minus_target_visible"
            if hidden_from_targets is not None
            else "legacy_hidden_mask"
        ),
        "geometry_hidden": policy_stats["geometry_hidden"],
        "obstacle_input": {
            "provided": obstacle is not None,
            **policy_stats["obstacle_input"],
        },
        "obstacle_retained": policy_stats["obstacle_retained"],
        "visual_removal": policy_stats["visual_removal"],
        "visual_removal_policy": policy_stats,
        "target_amodal": (
            {"provided": True, **mask_statistics(target_amodal)}
            if target_amodal is not None
            else {"provided": False}
        ),
        "target_visible": (
            {"provided": True, **mask_statistics(target_visible)}
            if target_visible is not None
            else {"provided": False}
        ),
    }
    return (
        Image.fromarray(geometry_hidden.astype(np.uint8) * 255, mode="L"),
        Image.fromarray(visual_removal.astype(np.uint8) * 255, mode="L"),
        report,
    )


def dilate(mask: Image.Image, radius: int) -> Image.Image:
    if radius <= 0:
        return mask
    array = np.asarray(mask) > 127
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1))
    return Image.fromarray(cv2.dilate(array.astype(np.uint8), kernel) * 255, mode="L")


def crop_around_mask(mask: Image.Image, padding: int) -> tuple[int, int, int, int]:
    binary = np.asarray(mask) > 127
    ys, xs = np.where(binary)
    if not len(xs):
        raise ValueError("The 2D completion mask is empty")
    width, height = mask.size
    x1, x2 = int(xs.min()), int(xs.max()) + 1
    y1, y2 = int(ys.min()), int(ys.max()) + 1
    x1, y1 = max(0, x1 - padding), max(0, y1 - padding)
    x2, y2 = min(width, x2 + padding), min(height, y2 + padding)

    # A near-square crop gives the diffusion model enough visible context on
    # both sides of a narrow structural occluder.
    side = max(x2 - x1, y2 - y1)
    cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0
    x1, y1 = int(round(cx - side / 2.0)), int(round(cy - side / 2.0))
    x2, y2 = x1 + side, y1 + side
    if x1 < 0:
        x2, x1 = x2 - x1, 0
    if y1 < 0:
        y2, y1 = y2 - y1, 0
    if x2 > width:
        x1, x2 = max(0, x1 - (x2 - width)), width
    if y2 > height:
        y1, y2 = max(0, y1 - (y2 - height)), height
    return x1, y1, x2, y2


def model_size(size: tuple[int, int], maximum: int) -> tuple[int, int]:
    width, height = size
    scale = min(maximum / max(width, height), 1.0)
    return (
        max(64, int(round(width * scale / 8.0)) * 8),
        max(64, int(round(height * scale / 8.0)) * 8),
    )


def composite_generated_crop(
    source: Image.Image,
    generated: Image.Image,
    mask_crop: Image.Image,
    box: tuple[int, int, int, int],
    feather_radius: float,
    preserve_crop: Image.Image | None = None,
) -> Image.Image:
    x1, y1, x2, y2 = box
    crop_size = (x2 - x1, y2 - y1)
    generated = generated.convert("RGB").resize(crop_size, Image.Resampling.LANCZOS)
    alpha = mask_crop.resize(crop_size, Image.Resampling.NEAREST)
    if feather_radius > 0:
        alpha = alpha.filter(ImageFilter.GaussianBlur(feather_radius))
    if preserve_crop is not None:
        preserve = (
            np.asarray(
                preserve_crop.resize(crop_size, Image.Resampling.NEAREST).convert("L")
            )
            > 127
        )
        alpha_array = np.asarray(alpha).copy()
        alpha_array[preserve] = 0
        alpha = Image.fromarray(alpha_array, mode="L")
    result = source.copy()
    result.paste(generated, (x1, y1), alpha)
    return result


def labelled_panel(image: Image.Image, label: str, size: tuple[int, int]) -> Image.Image:
    body = ImageOps.contain(image.convert("RGB"), (size[0], size[1] - 34))
    panel = Image.new("RGB", size, "white")
    ImageDraw.Draw(panel).text((9, 10), label, fill=(20, 20, 20))
    panel.paste(body, ((size[0] - body.width) // 2, 34 + (size[1] - 34 - body.height) // 2))
    return panel


def write_candidate_grid(
    path: Path,
    source: Image.Image,
    mask: Image.Image,
    candidates: list[dict[str, Any]],
) -> None:
    overlay = source.copy()
    overlay.paste(Image.new("RGB", source.size, (255, 45, 30)), mask=mask)
    overlay = Image.blend(source, overlay, 0.52)
    panel_size = (400, 560)
    items = [labelled_panel(source, "source", panel_size), labelled_panel(overlay, "completion mask", panel_size)]
    for row in candidates:
        quality = row["quality"]
        label = (
            f"candidate {row['index']} | score {quality['score']:.3f} "
            f"| gate {quality.get('gate_score', quality['score']):.3f}"
        )
        if "selection_score" in quality:
            label += f" | select {quality['selection_score']:.3f}"
        if quality.get("quality_flags"):
            label += " | " + ",".join(quality["quality_flags"])
        if row.get("selected"):
            label += " | REVIEW PICK"
        items.append(labelled_panel(Image.open(row["completed_rgb"]), label, panel_size))
    columns = 3
    rows = math.ceil(len(items) / columns)
    sheet = Image.new("RGB", (columns * panel_size[0], rows * panel_size[1]), (240, 240, 240))
    for index, item in enumerate(items):
        sheet.paste(item, ((index % columns) * panel_size[0], (index // columns) * panel_size[1]))
    sheet.save(path, quality=94, subsampling=0)


def load_backend(model: Path, device: str):
    from diffusers import StableDiffusionInpaintPipeline

    dtype = torch.float16 if device.startswith("cuda") else torch.float32
    pipeline = StableDiffusionInpaintPipeline.from_pretrained(
        str(model),
        torch_dtype=dtype,
        variant="fp16" if dtype == torch.float16 else None,
        safety_checker=None,
        requires_safety_checker=False,
        local_files_only=True,
    )
    pipeline = pipeline.to(device)
    pipeline.set_progress_bar_config(desc="Accessibility GPU 2D", leave=False)
    return pipeline


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", required=True)
    parser.add_argument(
        "--mask",
        default=None,
        help=(
            "Legacy reviewed hidden mask. Optional when both target masks are supplied; "
            "this geometry mask is never expanded to whole obstacle instances."
        ),
    )
    parser.add_argument(
        "--obstacle-mask",
        default=None,
        help=(
            "Detected obstacle mask. Visual removal keeps only complete connected "
            "components intersecting hidden; nearby non-occluding people remain."
        ),
    )
    parser.add_argument("--target-amodal-mask", default=None)
    parser.add_argument("--target-visible-mask", default=None)
    parser.add_argument("--category", choices=tuple(PROMPTS), required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--model", default="weights/stable-diffusion-inpainting")
    parser.add_argument("--prompt", default=None)
    parser.add_argument("--negative-prompt", default=NEGATIVE_PROMPT)
    parser.add_argument("--seed", type=int, default=9527)
    parser.add_argument("--steps", type=int, default=45)
    parser.add_argument("--guidance-scale", type=float, default=6.5)
    parser.add_argument("--num-candidates", type=int, default=4)
    parser.add_argument("--crop-padding", type=int, default=220)
    parser.add_argument("--resolution", type=int, default=768)
    parser.add_argument("--mask-dilate", type=int, default=8)
    parser.add_argument("--feather-radius", type=float, default=3.0)
    parser.add_argument("--device", default="cuda")
    return parser


def main() -> int:
    args = build_parser().parse_args()
    if args.num_candidates < 1 or args.steps < 1:
        raise ValueError("num-candidates and steps must be positive")
    image_path = resolve_path(args.image)
    mask_path = resolve_path(args.mask) if args.mask else None
    obstacle_path = resolve_path(args.obstacle_mask) if args.obstacle_mask else None
    target_amodal_path = (
        resolve_path(args.target_amodal_mask) if args.target_amodal_mask else None
    )
    target_visible_path = (
        resolve_path(args.target_visible_mask) if args.target_visible_mask else None
    )
    model_path = resolve_path(args.model)
    output_dir = resolve_path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    source = load_rgb(image_path)
    geometry_hidden_mask, visual_removal_mask, mask_report = load_completion_masks(
        size=source.size,
        hidden_path=mask_path,
        obstacle_path=obstacle_path,
        target_amodal_path=target_amodal_path,
        target_visible_path=target_visible_path,
    )
    if obstacle_path is not None:
        obstacle_array = np.asarray(load_binary_mask(obstacle_path, source.size)) > 127
        protected_non_occluding = obstacle_array & ~(
            np.asarray(visual_removal_mask) > 127
        )
    else:
        protected_non_occluding = np.zeros(
            (source.height, source.width),
            dtype=bool,
        )
        obstacle_array = None
    completion_envelope, envelope_report = build_completion_envelope(
        np.asarray(visual_removal_mask) > 127,
        obstacle_array,
    )
    completion_envelope_mask = Image.fromarray(
        completion_envelope.astype(np.uint8) * 255,
        mode="L",
    )
    target_amodal_array = (
        np.asarray(load_binary_mask(target_amodal_path, source.size)) > 127
        if target_amodal_path is not None
        else None
    )
    generation_mask = dilate(completion_envelope_mask, args.mask_dilate)
    # Dilation and feathering must not nibble into a nearby person/object that
    # the component policy intentionally excluded.
    generation_array = (np.asarray(generation_mask) > 127) & ~protected_non_occluding
    generation_mask = Image.fromarray(
        generation_array.astype(np.uint8) * 255,
        mode="L",
    )
    protected_non_occluding_mask = Image.fromarray(
        protected_non_occluding.astype(np.uint8) * 255,
        mode="L",
    )
    mask_report["generation_after_dilation"] = mask_statistics(
        generation_array
    )
    mask_report["completion_envelope"] = envelope_report
    mask_report["protected_non_occluding_obstacle"] = mask_statistics(
        protected_non_occluding
    )
    mask_report["generation_mask_dilate_radius"] = args.mask_dilate
    geometry_hidden_mask.save(output_dir / "geometry_hidden_mask.png")
    # Preserve the legacy review filename for downstream readers.
    geometry_hidden_mask.save(output_dir / "reviewed_hidden_mask.png")
    visual_removal_mask.save(output_dir / "visual_removal_mask.png")
    completion_envelope_mask.save(output_dir / "completion_envelope_mask.png")
    generation_mask.save(output_dir / "generation_mask.png")
    protected_non_occluding_mask.save(
        output_dir / "protected_non_occluding_obstacle_mask.png"
    )

    input_references = {
        "image": portable_input_reference(image_path),
        "legacy_hidden": (
            portable_input_reference(mask_path) if mask_path is not None else None
        ),
        "obstacle": (
            portable_input_reference(obstacle_path) if obstacle_path is not None else None
        ),
        "target_amodal": (
            portable_input_reference(target_amodal_path)
            if target_amodal_path is not None
            else None
        ),
        "target_visible": (
            portable_input_reference(target_visible_path)
            if target_visible_path is not None
            else None
        ),
    }
    if not np.any(np.asarray(generation_mask) > 127):
        selected = output_dir / "completed_rgb_selected.png"
        source.save(selected)
        manifest = {
            "schema_version": "accessibilityamodal_visual_completion_v1",
            "status": "skipped_empty_removal_mask",
            "human_review_required": True,
            "automatic_passability_claim": False,
            "inputs": input_references,
            "mask_statistics": mask_report,
            "files": {
                "geometry_hidden_mask": "geometry_hidden_mask.png",
                "visual_removal_mask": "visual_removal_mask.png",
                "completion_envelope_mask": "completion_envelope_mask.png",
                "generation_mask": "generation_mask.png",
                "protected_non_occluding_obstacle_mask": (
                    "protected_non_occluding_obstacle_mask.png"
                ),
                "selected_completed_rgb": selected.name,
            },
            "selected_completed_rgb": selected.name,
            "warning": (
                "Visual completion outputs are review candidates, not ground truth "
                "or evidence that a route is passable."
            ),
        }
        (output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
        return 0

    box = crop_around_mask(generation_mask, args.crop_padding)
    source_crop = source.crop(box)
    mask_crop = generation_mask.crop(box)
    preserve_crop = protected_non_occluding_mask.crop(box)
    inference_size = model_size(source_crop.size, args.resolution)
    model_image = source_crop.resize(inference_size, Image.Resampling.LANCZOS)
    model_mask = mask_crop.resize(inference_size, Image.Resampling.NEAREST)
    device = args.device if not args.device.startswith("cuda") or torch.cuda.is_available() else "cpu"
    prompt = args.prompt or PROMPTS[args.category]
    pipeline = load_backend(model_path, device)

    candidates: list[dict[str, Any]] = []
    for index in range(args.num_candidates):
        seed = args.seed + index * 1009
        generator = torch.Generator(device=device).manual_seed(seed)
        generated = pipeline(
            prompt=prompt,
            negative_prompt=args.negative_prompt,
            image=model_image,
            mask_image=model_mask,
            num_inference_steps=args.steps,
            guidance_scale=args.guidance_scale,
            generator=generator,
        ).images[0]
        completed = composite_generated_crop(
            source,
            generated,
            mask_crop,
            box,
            args.feather_radius,
            preserve_crop=preserve_crop,
        )
        completed_path = output_dir / f"completed_rgb_candidate_{index:02d}.png"
        crop_path = output_dir / f"generated_crop_candidate_{index:02d}.png"
        completed.save(completed_path)
        generated.save(crop_path)
        quality = candidate_quality(
            completed,
            geometry_hidden_mask,
            args.category,
        )
        quality["clutter_metrics"] = candidate_clutter_metrics(
            completed,
            source,
            completion_envelope,
            target_amodal_array,
        )
        candidates.append(
            {
                "index": index,
                "seed": seed,
                "completed_rgb": str(completed_path),
                "generated_crop": str(crop_path),
                "quality": quality,
            }
        )

    apply_selection_clutter_penalty(candidates)
    selected_row, status = select_candidate(candidates)
    selected_row["selected"] = True
    selected_path = output_dir / "completed_rgb_selected.png"
    shutil.copy2(selected_row["completed_rgb"], selected_path)
    source_crop.save(output_dir / "source_context_crop.png")
    write_candidate_grid(output_dir / "candidate_comparison.jpg", source, generation_mask, candidates)

    manifest_candidates = []
    for row in candidates:
        portable_row = dict(row)
        portable_row["completed_rgb"] = Path(row["completed_rgb"]).name
        portable_row["generated_crop"] = Path(row["generated_crop"]).name
        manifest_candidates.append(portable_row)
    manifest = {
        "schema_version": "accessibilityamodal_visual_completion_v1",
        "created_at_utc": datetime.now(timezone.utc).isoformat(),
        "status": status,
        "human_review_required": True,
        "automatic_passability_claim": False,
        "pipeline": "AccessibilityAmodal GPU 2D Visual Completion",
        "backend": "Diffusers StableDiffusionInpaintPipeline",
        "inputs": input_references,
        "mask_statistics": mask_report,
        "quality_evaluation_mask": "geometry_hidden_mask.png",
        "category": args.category,
        "model": portable_input_reference(model_path),
        "prompt": prompt,
        "negative_prompt": args.negative_prompt,
        "device": device,
        "steps": args.steps,
        "guidance_scale": args.guidance_scale,
        "crop_box": box,
        "inference_size": inference_size,
        "mask_dilate": args.mask_dilate,
        "feather_radius": args.feather_radius,
        "selected_index": selected_row["index"],
        "selection_score_policy": (
            "absolute_hidden_surface_quality_minus_up_to_0.15_"
            "cohort_relative_outside_target_clutter_rank"
        ),
        "selected_completed_rgb": selected_path.name,
        "candidate_comparison": "candidate_comparison.jpg",
        "files": {
            "geometry_hidden_mask": "geometry_hidden_mask.png",
            "visual_removal_mask": "visual_removal_mask.png",
            "completion_envelope_mask": "completion_envelope_mask.png",
            "generation_mask": "generation_mask.png",
            "protected_non_occluding_obstacle_mask": (
                "protected_non_occluding_obstacle_mask.png"
            ),
            "source_context_crop": "source_context_crop.png",
            "selected_completed_rgb": selected_path.name,
            "candidate_comparison": "candidate_comparison.jpg",
        },
        "candidates": manifest_candidates,
        "warning": (
            "Generative 2D outputs are visual review candidates, not ground truth "
            "or evidence that a route is passable."
        ),
    }
    (output_dir / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    if status == "withheld_needs_review":
        print(
            f"All candidates carry quality risks; retained candidate "
            f"{selected_row['index']} for review and withheld automatic acceptance."
        )
    else:
        print(
            f"Selected visual candidate {selected_row['index']} by gate score "
            f"{selected_row['quality']['gate_score']:.4f}; human review remains required."
        )
    print(f"Wrote GPU 2D review candidate to {selected_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())