#!/usr/bin/env python3 """AccessibilityAmodal one-command mask, 2D completion, and 3D inference. This is a lightweight login-node launcher. GPU work is submitted to Slurm via ``slurm/run_accesspath_demo.sbatch``; this process can wait for the job and report the persistent result directory. The third-party Accessibility3D model is accessed only through the project adapter ``tools/accessibility_3d_completion.py``. """ from __future__ import annotations import argparse import json import os import re import shutil import subprocess import sys import time from datetime import datetime from pathlib import Path from accesspath3r.privacy import public_path, sanitize_text PROJECT_ROOT = Path(__file__).resolve().parents[1] SBATCH_SCRIPT = PROJECT_ROOT / "slurm" / "run_accesspath_demo.sbatch" PRIMARY_3D_DIRNAME = "accessibility3d" DEFAULT_PROMPT_CONFIG = ( PROJECT_ROOT / "configs" / "accessibility_mask_prompts_e5_v2.json" ) STRUCTURAL_STAIRS_PROMPT_CONFIG = ( PROJECT_ROOT / "configs" / "accessibility_stairs_structural_occlusion_prompts.json" ) OUTDOOR_STAIRS_NO_OCCLUSION_PROMPT_CONFIG = ( PROJECT_ROOT / "configs" / "accessibility_stairs_outdoor_no_occlusion_prompts.json" ) CATEGORIES = ("curb_cut", "ramp", "stairs", "tactile_paving", "walkway") GEOMETRY_ALLOWED_2D_STATUSES = { "candidate_selected_for_review", "skipped_empty_removal_mask", } TERMINAL_STATES = { "BOOT_FAIL", "CANCELLED", "COMPLETED", "DEADLINE", "FAILED", "NODE_FAIL", "OUT_OF_MEMORY", "PREEMPTED", "REVOKED", "TIMEOUT", } def safe_sample_id(value: str) -> str: result = re.sub(r"[^A-Za-z0-9_-]+", "_", value).strip("_") return result or "single_image" def normalized_state(value: str) -> str: return value.strip().split()[0].split("+")[0] if value.strip() else "" def run_command(command: list[str]) -> subprocess.CompletedProcess[str]: return subprocess.run( command, cwd=PROJECT_ROOT, check=False, capture_output=True, text=True, ) def query_job_state(job_id: str) -> str | None: queued = run_command(["squeue", "-h", "-j", job_id, "-o", "%T"]) if queued.returncode == 0 and queued.stdout.strip(): return normalized_state(queued.stdout.splitlines()[0]) accounting = run_command( [ "sacct", "-n", "-P", "-j", job_id, "--format=JobIDRaw,State,ExitCode", ] ) if accounting.returncode != 0: return None for line in accounting.stdout.splitlines(): fields = line.split("|") if len(fields) >= 2 and fields[0] == job_id: return normalized_state(fields[1]) return None def format_clock(seconds: float) -> str: seconds = max(0, int(round(seconds))) hours, remainder = divmod(seconds, 3600) minutes, seconds = divmod(remainder, 60) if hours: return f"{hours:d}:{minutes:02d}:{seconds:02d}" return f"{minutes:02d}:{seconds:02d}" def parse_timestamp(value: object) -> float | None: if not isinstance(value, str): return None try: return datetime.fromisoformat(value).timestamp() except ValueError: return None def read_progress(output_dir: Path) -> dict | None: try: value = json.loads((output_dir / "progress.json").read_text(encoding="utf-8")) except (FileNotFoundError, json.JSONDecodeError, OSError): return None return value if isinstance(value, dict) else None def infer_progress_from_artifacts( output_dir: Path, sample_id: str, include_generative_3d: bool, ) -> dict | None: """Provide useful progress for jobs launched before progress.json existed.""" presentation_manifest = ( output_dir / "04_3d_completion" / "geometry" / "presentation_models" / "manifest.json" ) legacy_variants_manifest = ( output_dir / "04_3d_completion" / "geometry" / "3d_variants" / "manifest.json" ) if not presentation_manifest.is_file() and legacy_variants_manifest.is_file(): presentation_manifest = legacy_variants_manifest learned_manifest = output_dir / PRIMARY_3D_DIRNAME / "manifest.json" current_compatibility_manifest = ( output_dir / "04_3d_completion" / "visual_candidate" / "manifest.json" ) if not learned_manifest.is_file() and current_compatibility_manifest.is_file(): learned_manifest = current_compatibility_manifest final_marker = learned_manifest if include_generative_3d else presentation_manifest if final_marker.is_file(): return { "status": "completed", "percent": 100.0, "stage_label": "All outputs completed", } milestones = ( ( learned_manifest, 96.0, "Accessibility3D CUDA Gaussian rotation and dense triangle mesh", ), ( presentation_manifest, 74.0, "Diagnostic category-constrained geometry", ), ( output_dir / "04_3d_completion" / "geometry" / "completed_mesh_turntable_slow.gif", 44.0, "Continuous-surface and solid 3D views", ), ( output_dir / "04_3d_completion" / "geometry" / "geometry_manifest.json", 43.0, "Basic 3D turntable preview", ), ( output_dir / "03_2d_completion" / "completed_rgb_selected.png", 34.0, "Depth and 3D geometry reconstruction", ), ( output_dir / "02_masks" / "metadata.json", 20.0, "GPU generative 2D completion and candidate selection", ), ( output_dir / "01_sam3" / "summary.json", 16.0, "Visible, hidden, amodal, and obstacle masks", ), ) for marker, percent, label in milestones: if marker.is_file(): return {"status": "running", "percent": percent, "stage_label": label} return None def interpolated_percent(progress: dict, now: float) -> float: try: percent = float(progress.get("percent", 0.0)) except (TypeError, ValueError): percent = 0.0 if progress.get("status") != "running": return max(0.0, min(100.0, percent)) try: end_percent = float(progress.get("stage_end_percent", percent)) expected = float(progress.get("stage_expected_seconds", 0.0)) except (TypeError, ValueError): return max(0.0, min(99.0, percent)) stage_started = parse_timestamp(progress.get("stage_started_at")) if stage_started is None or expected <= 0 or end_percent <= percent: return max(0.0, min(99.0, percent)) fraction = max(0.0, min(0.98, (now - stage_started) / expected)) interpolated = percent + (end_percent - percent) * fraction return max(0.0, min(99.0, interpolated)) def estimated_remaining_seconds( progress: dict | None, now: float, elapsed: float, base_total: float, ) -> float: """Keep ETA conservative when a stage or an earlier stage runs long.""" adjusted_total = max(1.0, base_total) if not progress or progress.get("status") != "running": return max(0.0, adjusted_total - elapsed) overall_started = parse_timestamp(progress.get("started_at")) stage_started = parse_timestamp(progress.get("stage_started_at")) try: start_percent = float(progress.get("percent", 0.0)) stage_expected = float(progress.get("stage_expected_seconds", 0.0)) except (TypeError, ValueError): return max(0.0, adjusted_total - elapsed) if overall_started is not None and stage_started is not None: planned_stage_start = max(0.0, start_percent) / 100.0 * base_total actual_stage_start = max(0.0, stage_started - overall_started) adjusted_total += max(0.0, actual_stage_start - planned_stage_start) adjusted_total += max(0.0, now - stage_started - stage_expected) return max(0.0, adjusted_total - elapsed) class ProgressDisplay: def __init__(self, job_id: str) -> None: self.job_id = job_id self.is_tty = sys.stdout.isatty() self.last_key: tuple[str, str] | None = None self.last_width = 0 def update( self, percent: float, label: str, elapsed: float, eta: float | None, state: str, final: bool = False, ) -> None: width = 28 filled = min(width, max(0, int(round(width * percent / 100.0)))) bar = "#" * filled + "-" * (width - filled) if state == "PENDING": timing = f"queue {format_clock(elapsed)} | GPU ETA ~{format_clock(eta or 0)}" elif state == "FAILED": timing = f"elapsed {format_clock(elapsed)} | failed" elif eta is None: timing = f"elapsed {format_clock(elapsed)} | ETA calculating" elif eta <= 0 and not final: timing = f"elapsed {format_clock(elapsed)} | finishing" else: timing = f"elapsed {format_clock(elapsed)} | ETA ~{format_clock(eta)}" line = ( f"[Slurm {self.job_id}] [{bar}] {percent:5.1f}% | " f"{label} | {timing}" ) key = (state, label) if self.is_tty: padding = " " * max(0, self.last_width - len(line)) print(f"\r{line}{padding}", end="\n" if final else "", flush=True) self.last_width = len(line) elif key != self.last_key or final: print(line, flush=True) self.last_key = key def wait_for_job( job_id: str, poll_seconds: float, output_dir: Path, sample_id: str, include_generative_3d: bool, estimated_seconds: float, ) -> str: missing_polls = 0 submitted_at = time.monotonic() running_at: float | None = None display = ProgressDisplay(job_id) while True: state = query_job_state(job_id) now_wall = time.time() now_monotonic = time.monotonic() progress = read_progress(output_dir) or infer_progress_from_artifacts( output_dir, sample_id, include_generative_3d, ) if state: missing_polls = 0 if state == "RUNNING" and running_at is None: running_at = now_monotonic else: # sacct can lag briefly after a job leaves squeue. missing_polls += 1 if missing_polls >= 12 and not progress: raise RuntimeError( "Unable to read the job state from squeue or sacct. " f"Check manually with: squeue -j {job_id}" ) progress_status = progress.get("status") if progress else None if progress_status == "completed": runtime_started = parse_timestamp(progress.get("started_at")) elapsed = now_wall - runtime_started if runtime_started else 0.0 display.update(100.0, "All outputs completed", elapsed, 0.0, "COMPLETED", final=True) return "COMPLETED" if progress_status == "failed": runtime_started = parse_timestamp(progress.get("started_at")) elapsed = now_wall - runtime_started if runtime_started else 0.0 display.update( interpolated_percent(progress, now_wall), str(progress.get("message") or progress.get("stage_label") or "GPU stage failed"), elapsed, None, "FAILED", final=True, ) return "FAILED" if state == "PENDING": display.update( 0.0, "Waiting for a GPU", now_monotonic - submitted_at, estimated_seconds, state, ) else: if running_at is None: running_at = now_monotonic runtime_started = parse_timestamp(progress.get("started_at")) if progress else None elapsed = ( max(0.0, now_wall - runtime_started) if runtime_started is not None else now_monotonic - running_at ) total_estimate = estimated_seconds if progress: try: total_estimate = float(progress.get("total_estimated_seconds", estimated_seconds)) except (TypeError, ValueError): pass percent = interpolated_percent(progress, now_wall) if progress else 0.0 label = ( str(progress.get("stage_label") or "Initializing GPU job") if progress else "Initializing GPU job" ) display.update( percent, label, elapsed, estimated_remaining_seconds(progress, now_wall, elapsed, total_estimate), state or "UNKNOWN", final=bool(state in TERMINAL_STATES), ) if state in TERMINAL_STATES: return state time.sleep(poll_seconds) def estimate_runtime_seconds(args: argparse.Namespace) -> int: # Broad SAM3 prompt evaluation and the presentation-model renderer scale # with image size. These values are calibrated from the 1368x1824 9528 # run; GLB texture baking is accounted for separately because it dominated # that job's final stage instead of behaving like the normal mesh export. seconds = (5 if args.reviewed_visible_workspace else 150) + 13 + (24 if not args.no_depth else 6) + 3 + 420 if not args.no_2d: seconds += 40 if args.fast_2d_baseline: seconds += 12 if not args.no_generative_3d: seconds += 75 if args.export_glb: seconds += 540 return seconds def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description=( "Submit one accessibility image to SAM3 mask inference, amodal/hidden " "completion, 2D completion, Depth Anything geometry reconstruction, " "and Accessibility3D CUDA Gaussian rendering with a dense triangle mesh." ) ) parser.add_argument("image_pos", nargs="?", help="Input RGB image (positional form).") parser.add_argument("--image", dest="image_opt", help="Input RGB image.") parser.add_argument("--category", required=True, choices=CATEGORIES) parser.add_argument("--sample-id", default=None) parser.add_argument( "--output-dir", default=None, help=( "Persistent output directory. " "Default: output/inference/_