#!/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())