| |
| """Run prompt-consensus accessibility masks on a small reviewed pilot set. |
| |
| Supported backends: |
| sam3 Meta SAM 3 text-prompt concept segmentation (preferred) |
| grounded_sam GroundingDINO boxes refined by SAM ViT-H (local control) |
| |
| The tool saves per-prompt masks, a high-recall union, a majority-vote core, an |
| uncertainty mask, both raw/high-recall and filtered obstacle masks, overlays, |
| and machine-readable quality scores. Automatic output is a review proposal, |
| never ground truth. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import html |
| import json |
| import math |
| import os |
| import re |
| import sys |
| import traceback |
| from collections import Counter |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| import cv2 |
| import numpy as np |
| from PIL import Image, ImageOps |
|
|
|
|
| def read_jsonl(path: Path) -> list[dict[str, Any]]: |
| rows: list[dict[str, Any]] = [] |
| with path.open("r", encoding="utf-8") as handle: |
| for line in handle: |
| if line.strip(): |
| rows.append(json.loads(line)) |
| return rows |
|
|
|
|
| def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp_path = path.with_suffix(path.suffix + ".tmp") |
| with tmp_path.open("w", encoding="utf-8") as handle: |
| for row in rows: |
| handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n") |
| tmp_path.replace(path) |
|
|
|
|
| def json_dump(path: Path, value: Any) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp_path = path.with_suffix(path.suffix + ".tmp") |
| with tmp_path.open("w", encoding="utf-8") as handle: |
| json.dump(value, handle, ensure_ascii=False, indent=2, sort_keys=True) |
| tmp_path.replace(path) |
|
|
|
|
| STANDARD_SAMPLE_OUTPUTS = ( |
| "target_visible_candidate.png", |
| "target_visible_core.png", |
| "target_uncertain.png", |
| "obstacle_unfiltered_evidence.png", |
| "obstacle_all_detected.png", |
| "obstacle.png", |
| "obstacle_red.png", |
| "overlay.png", |
| "quality.json", |
| ) |
| OBSTACLE_EVIDENCE_POLICY_VERSION = 1 |
|
|
|
|
| def load_resume_quality(sample_dir: Path, sample_id: str, backend_name: str) -> dict[str, Any] | None: |
| quality_path = sample_dir / "quality.json" |
| if not quality_path.is_file(): |
| return None |
| if any(not (sample_dir / filename).is_file() for filename in STANDARD_SAMPLE_OUTPUTS): |
| return None |
| with quality_path.open("r", encoding="utf-8") as handle: |
| quality = json.load(handle) |
| if quality.get("sample_id") != sample_id: |
| return None |
| if quality.get("backend") != backend_name: |
| return None |
| |
| |
| |
| |
| if quality.get("obstacle_evidence_policy_version") != OBSTACLE_EVIDENCE_POLICY_VERSION: |
| return None |
| if quality.get("review_status") == "error" or quality.get("error"): |
| return None |
| return quality |
|
|
|
|
| def safe_name(text: str) -> str: |
| return re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_") |
|
|
|
|
| def prompt_list(value: Any) -> list[str]: |
| """Normalize a prompt list while accepting the legacy JSON shape.""" |
| if isinstance(value, str): |
| raw = [value] |
| elif isinstance(value, (list, tuple)): |
| raw = value |
| else: |
| raw = [] |
| return list(dict.fromkeys(str(prompt).strip() for prompt in raw if str(prompt).strip())) |
|
|
|
|
| def obstacle_prompt_resolution( |
| config: dict[str, Any], category: str, backend_name: str | None = None |
| ) -> dict[str, Any]: |
| """Resolve backend-aware obstacle prompts without breaking legacy configs. |
| |
| When a backend-specific list is present it replaces the corresponding |
| legacy global/category list. The other level still falls back separately, |
| so a config may specialize only global prompts or only one category. With |
| ``backend_name=None`` the output is exactly the historical legacy union. |
| """ |
| global_prompts = prompt_list(config.get("obstacle_prompts")) |
| category_map = config.get("category_obstacle_prompts") |
| category_prompts = prompt_list( |
| category_map.get(category) if isinstance(category_map, dict) else None |
| ) |
| global_source = "legacy.obstacle_prompts" |
| category_source = "legacy.category_obstacle_prompts" |
| if backend_name: |
| by_backend = config.get("obstacle_prompts_by_backend") |
| if isinstance(by_backend, dict) and backend_name in by_backend: |
| global_prompts = prompt_list(by_backend.get(backend_name)) |
| global_source = f"obstacle_prompts_by_backend.{backend_name}" |
| category_by_backend = config.get("category_obstacle_prompts_by_backend") |
| backend_categories = ( |
| category_by_backend.get(backend_name) |
| if isinstance(category_by_backend, dict) |
| else None |
| ) |
| if isinstance(backend_categories, dict) and category in backend_categories: |
| category_prompts = prompt_list(backend_categories.get(category)) |
| category_source = ( |
| f"category_obstacle_prompts_by_backend.{backend_name}.{category}" |
| ) |
| prompts = list(dict.fromkeys([*global_prompts, *category_prompts])) |
| return { |
| "backend": backend_name, |
| "category": category, |
| "global_source": global_source, |
| "category_source": category_source, |
| "prompts": prompts, |
| } |
|
|
|
|
| def obstacle_prompts_for_category( |
| config: dict[str, Any], category: str, backend_name: str | None = None |
| ) -> list[str]: |
| """Return backend-aware obstacle prompts, with legacy fallback support.""" |
| return list(obstacle_prompt_resolution(config, category, backend_name)["prompts"]) |
|
|
|
|
| def as_mapping(value: Any) -> dict[str, Any]: |
| return dict(value) if isinstance(value, dict) else {} |
|
|
|
|
| def merge_mappings(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: |
| """Small recursive merge for backend-specific evidence policy overrides.""" |
| result = dict(base) |
| for key, value in override.items(): |
| if isinstance(value, dict) and isinstance(result.get(key), dict): |
| result[key] = merge_mappings(dict(result[key]), value) |
| else: |
| result[key] = value |
| return result |
|
|
|
|
| def merge_prompt_config(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: |
| """Merge a small, output-run-specific prompt overlay safely. |
| |
| Nested objects are merged, while prompt lists are appended in stable |
| de-duplicated order. This lets a finite remediation run add a narrowly |
| justified prompt without copying or silently editing the frozen base |
| configuration used by a control run. Scalar policy values remain an |
| explicit override, so their provenance can be locked by the Slurm wrapper. |
| """ |
| result = dict(base) |
| for key, value in override.items(): |
| current = result.get(key) |
| if isinstance(current, dict) and isinstance(value, dict): |
| result[key] = merge_prompt_config(dict(current), value) |
| elif isinstance(current, list) and isinstance(value, list): |
| result[key] = list(dict.fromkeys([*current, *value])) |
| else: |
| result[key] = value |
| return result |
|
|
|
|
| def float_or_default(value: Any, default: float) -> float: |
| try: |
| return float(value) |
| except (TypeError, ValueError): |
| return default |
|
|
|
|
| def normalized_float_mapping(value: Any) -> dict[str, float]: |
| mapping = as_mapping(value) |
| return { |
| str(key).strip().lower(): float_or_default(item, 0.0) |
| for key, item in mapping.items() |
| if str(key).strip() |
| } |
|
|
|
|
| def normalized_override_mapping(value: Any) -> dict[str, dict[str, Any]]: |
| mapping = as_mapping(value) |
| return { |
| str(key).strip().lower(): as_mapping(item) |
| for key, item in mapping.items() |
| if str(key).strip() and isinstance(item, dict) |
| } |
|
|
|
|
| def resolve_obstacle_evidence_policy( |
| config: dict[str, Any], backend_name: str |
| ) -> dict[str, Any]: |
| """Compile a serializable trusted-obstacle-evidence policy. |
| |
| New configs may use ``obstacle_evidence_policy`` and optional |
| ``obstacle_evidence_policy_by_backend.<backend>``. The aliases |
| ``obstacle_all_detected_policy`` / ``..._by_backend`` are accepted for |
| early experiment configs. In a legacy config, the existing |
| ``obstacle_maximum_area_ratio`` map becomes the per-prompt area cap, |
| confidence filtering remains disabled, and a conservative route-like |
| component guard protects downstream hidden completion from broad scene |
| surfaces. |
| """ |
| base = as_mapping( |
| config.get("obstacle_evidence_policy", config.get("obstacle_all_detected_policy", {})) |
| ) |
| by_backend = as_mapping( |
| config.get( |
| "obstacle_evidence_policy_by_backend", |
| config.get("obstacle_all_detected_policy_by_backend", {}), |
| ) |
| ) |
| backend_override = as_mapping(by_backend.get(backend_name)) |
| raw = merge_mappings(base, backend_override) |
| legacy_maximums = normalized_float_mapping(config.get("obstacle_maximum_area_ratio")) |
| maximums = dict(legacy_maximums) |
| maximums.update( |
| normalized_float_mapping( |
| raw.get("maximum_area_ratio", raw.get("maximum_area_ratio_by_prompt", {})) |
| ) |
| ) |
| prompt_overrides = normalized_override_mapping( |
| raw.get("prompt_overrides", raw.get("prompts", {})) |
| ) |
| confidence_by_prompt = normalized_float_mapping( |
| raw.get("minimum_confidence_by_prompt", raw.get("minimum_prompt_confidence_by_prompt", {})) |
| ) |
| route_like_raw = as_mapping( |
| raw.get("route_like", raw.get("route_like_component_policy", {})) |
| ) |
| |
| |
| |
| route_like = { |
| "enabled": bool(route_like_raw.get("enabled", True)), |
| "minimum_bbox_width_ratio": float_or_default( |
| route_like_raw.get("minimum_bbox_width_ratio"), 0.80 |
| ), |
| "minimum_bbox_height_ratio": float_or_default( |
| route_like_raw.get("minimum_bbox_height_ratio"), 0.12 |
| ), |
| "minimum_component_area_ratio": float_or_default( |
| route_like_raw.get("minimum_component_area_ratio"), 0.03 |
| ), |
| } |
| default_maximum = float_or_default( |
| raw.get("default_maximum_area_ratio", maximums.get("default", 0.15)), |
| 0.15, |
| ) |
| return { |
| "version": OBSTACLE_EVIDENCE_POLICY_VERSION, |
| "backend": backend_name, |
| "source": { |
| "global": "obstacle_evidence_policy" if base else "legacy_defaults", |
| "backend_override_present": bool(backend_override), |
| "legacy_obstacle_maximum_area_ratio_used": bool(legacy_maximums), |
| }, |
| "minimum_component_area_pixels": max( |
| 1, int(float_or_default(raw.get("minimum_component_area_pixels"), 16.0)) |
| ), |
| "minimum_component_area_ratio": max( |
| 0.0, float_or_default(raw.get("minimum_component_area_ratio"), 0.00002) |
| ), |
| "minimum_prompt_confidence": min( |
| 1.0, |
| max( |
| 0.0, |
| float_or_default( |
| raw.get( |
| "minimum_prompt_confidence", |
| raw.get("minimum_confidence", 0.0), |
| ), |
| 0.0, |
| ), |
| ), |
| ), |
| "minimum_confidence_by_prompt": confidence_by_prompt, |
| "default_maximum_area_ratio": min(1.0, max(0.0, default_maximum)), |
| "maximum_area_ratio_by_prompt": maximums, |
| "full_frame_area_ratio": min( |
| 1.0, max(0.0, float_or_default(raw.get("full_frame_area_ratio"), 0.90)) |
| ), |
| "prompt_overrides": prompt_overrides, |
| "route_like": route_like, |
| } |
|
|
|
|
| def evidence_minimum_component_area(policy: dict[str, Any], image_area: int) -> int: |
| return max( |
| int(policy["minimum_component_area_pixels"]), |
| int(round(image_area * float(policy["minimum_component_area_ratio"]))), |
| ) |
|
|
|
|
| def evidence_prompt_limits(prompt: str, policy: dict[str, Any]) -> tuple[float, float]: |
| key = prompt.strip().lower() |
| override = as_mapping(policy["prompt_overrides"].get(key)) |
| confidence = float_or_default( |
| override.get( |
| "minimum_confidence", |
| policy["minimum_confidence_by_prompt"].get(key, policy["minimum_prompt_confidence"]), |
| ), |
| float(policy["minimum_prompt_confidence"]), |
| ) |
| maximum_area = float_or_default( |
| override.get( |
| "maximum_area_ratio", |
| policy["maximum_area_ratio_by_prompt"].get( |
| key, policy["default_maximum_area_ratio"] |
| ), |
| ), |
| float(policy["default_maximum_area_ratio"]), |
| ) |
| return min(1.0, max(0.0, confidence)), min(1.0, max(0.0, maximum_area)) |
|
|
|
|
| def save_mask(path: Path, mask: np.ndarray) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(path) |
|
|
|
|
| def save_red_mask(path: Path, mask: np.ndarray) -> None: |
| """Save an RGB review visualization while preserving obstacle.png as binary.""" |
| path.parent.mkdir(parents=True, exist_ok=True) |
| color = np.zeros((*mask.shape, 3), dtype=np.uint8) |
| color[mask] = (230, 45, 45) |
| Image.fromarray(color, mode="RGB").save(path) |
|
|
|
|
| def remove_small_components(mask: np.ndarray, minimum_area: int) -> np.ndarray: |
| count, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8) |
| output = np.zeros_like(mask, dtype=bool) |
| for index in range(1, count): |
| if int(stats[index, cv2.CC_STAT_AREA]) >= minimum_area: |
| output |= labels == index |
| return output |
|
|
|
|
| def retain_core_connected_union(union: np.ndarray, core: np.ndarray, minimum_area: int) -> np.ndarray: |
| union = remove_small_components(union, minimum_area) |
| count, labels, stats, _ = cv2.connectedComponentsWithStats(union.astype(np.uint8), 8) |
| output = np.zeros_like(union, dtype=bool) |
| for index in range(1, count): |
| component = labels == index |
| if int(stats[index, cv2.CC_STAT_AREA]) >= minimum_area and np.any(component & core): |
| output |= component |
| return output |
|
|
|
|
| def retain_nearby_obstacles( |
| obstacle: np.ndarray, |
| target: np.ndarray, |
| minimum_area: int, |
| maximum_area_ratio: float = 0.15, |
| ) -> np.ndarray: |
| """Keep compact obstacle components spatially adjacent to the target surface.""" |
| distance = max(9, int(round(min(target.shape) * 0.035)) | 1) |
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (distance, distance)) |
| nearby = cv2.dilate(target.astype(np.uint8), kernel).astype(bool) |
| count, labels, stats, _ = cv2.connectedComponentsWithStats(obstacle.astype(np.uint8), 8) |
| output = np.zeros_like(obstacle, dtype=bool) |
| image_area = obstacle.size |
| for index in range(1, count): |
| area = int(stats[index, cv2.CC_STAT_AREA]) |
| component = labels == index |
| if ( |
| area >= minimum_area |
| and area / image_area <= maximum_area_ratio |
| and np.any(component & nearby) |
| ): |
| output |= component |
| return output |
|
|
|
|
| def retain_compact_obstacles( |
| obstacle: np.ndarray, |
| minimum_area: int, |
| maximum_area_ratio: float = 0.12, |
| ) -> np.ndarray: |
| """Keep compact scene obstacles without requiring target adjacency.""" |
| count, labels, stats, _ = cv2.connectedComponentsWithStats(obstacle.astype(np.uint8), 8) |
| output = np.zeros_like(obstacle, dtype=bool) |
| image_area = obstacle.size |
| for index in range(1, count): |
| area = int(stats[index, cv2.CC_STAT_AREA]) |
| if minimum_area <= area <= image_area * maximum_area_ratio: |
| output |= labels == index |
| return output |
|
|
|
|
| def route_like_component_count(mask: np.ndarray, route_policy: dict[str, Any]) -> int: |
| """Count broad ground/route-like components that must not seed hidden masks.""" |
| if not bool(route_policy.get("enabled", False)) or not np.any(mask): |
| return 0 |
| count, _, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8) |
| height, width = mask.shape |
| route_like = 0 |
| for index in range(1, count): |
| area = int(stats[index, cv2.CC_STAT_AREA]) |
| width_ratio = float(stats[index, cv2.CC_STAT_WIDTH] / max(width, 1)) |
| height_ratio = float(stats[index, cv2.CC_STAT_HEIGHT] / max(height, 1)) |
| area_ratio = float(area / max(mask.size, 1)) |
| if ( |
| width_ratio >= float(route_policy["minimum_bbox_width_ratio"]) |
| and height_ratio >= float(route_policy["minimum_bbox_height_ratio"]) |
| and area_ratio >= float(route_policy["minimum_component_area_ratio"]) |
| ): |
| route_like += 1 |
| return route_like |
|
|
|
|
| def trusted_obstacle_prompt_evidence( |
| *, |
| prompt: str, |
| group: str, |
| prediction: PromptPrediction, |
| minimum_component_area: int, |
| policy: dict[str, Any], |
| ) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]: |
| """Return all cleaned evidence and the subset safe to seed hidden completion. |
| |
| The first mask is intentionally high recall: it only removes microscopic |
| components. The second is accepted only if its prompt confidence and |
| image-area sanity checks pass and it is not route-like. Rejecting at the |
| prompt level makes the audit rationale easy to understand and guarantees |
| that a broad false positive cannot leak into ``obstacle_all_detected``. |
| """ |
| cleaned = remove_small_components(prediction.mask, minimum_component_area) |
| minimum_confidence, maximum_area_ratio = evidence_prompt_limits(prompt, policy) |
| area_ratio = float(cleaned.mean()) |
| reasons: list[str] = [] |
| if not np.any(cleaned): |
| reasons.append("empty_after_small_component_removal") |
| if float(prediction.confidence) < minimum_confidence: |
| reasons.append("confidence_below_minimum") |
| if area_ratio > maximum_area_ratio: |
| reasons.append("area_ratio_exceeds_prompt_maximum") |
| if area_ratio > float(policy["full_frame_area_ratio"]): |
| reasons.append("area_ratio_exceeds_full_frame_guard") |
| route_like_count = route_like_component_count(cleaned, policy["route_like"]) |
| if route_like_count: |
| reasons.append("route_like_component_detected") |
| accepted = not reasons |
| trusted = cleaned if accepted else np.zeros_like(cleaned) |
| decision = { |
| "prompt": prompt, |
| "group": group, |
| "confidence": round(float(prediction.confidence), 6), |
| "instance_count": int(prediction.instance_count), |
| "minimum_confidence": round(minimum_confidence, 6), |
| "maximum_area_ratio": round(maximum_area_ratio, 6), |
| "cleaned_pixels": int(cleaned.sum()), |
| "cleaned_area_ratio": round(area_ratio, 6), |
| "route_like_component_count": route_like_count, |
| "decision": "accepted" if accepted else "rejected", |
| "reasons": reasons, |
| "trusted_pixels": int(trusted.sum()), |
| } |
| return cleaned, trusted, decision |
|
|
|
|
| def area_plausibility(area_ratio: float, minimum: float, maximum: float) -> float: |
| if area_ratio <= 0 or area_ratio < minimum / 2 or area_ratio > min(1.0, maximum * 1.2): |
| return 0.0 |
| if minimum <= area_ratio <= maximum: |
| return 1.0 |
| if area_ratio < minimum: |
| return max(0.0, area_ratio / minimum) |
| return max(0.0, 1.0 - (area_ratio - maximum) / max(1.0 - maximum, 1e-6)) |
|
|
|
|
| @dataclass |
| class PromptPrediction: |
| mask: np.ndarray |
| confidence: float |
| instance_count: int |
| boxes: list[list[float]] |
|
|
|
|
| class Sam3Backend: |
| name = "sam3" |
|
|
| def __init__(self, args: argparse.Namespace): |
| repo = Path(args.sam3_repo).resolve() |
| checkpoint = Path(args.sam3_checkpoint).resolve() |
| if not repo.is_dir(): |
| raise FileNotFoundError(f"SAM 3 repository not found: {repo}") |
| if not checkpoint.is_file(): |
| raise FileNotFoundError( |
| f"SAM 3 checkpoint not found: {checkpoint}. Accept the model terms and download sam3.pt first." |
| ) |
| sys.path.insert(0, str(repo)) |
| import torch |
| from sam3.model.sam3_image_processor import Sam3Processor |
| from sam3.model_builder import build_sam3_image_model |
|
|
| self.torch = torch |
| self.device = args.device |
| self.autocast_device = str(args.device).split(":", 1)[0] |
| self.autocast_dtype = torch.bfloat16 |
| model = build_sam3_image_model( |
| device=args.device, |
| checkpoint_path=str(checkpoint), |
| load_from_HF=False, |
| compile=args.compile, |
| ) |
| self.processor = Sam3Processor( |
| model, |
| device=args.device, |
| confidence_threshold=args.confidence_threshold, |
| ) |
| self.state: dict[str, Any] | None = None |
|
|
| def begin_image(self, image: Image.Image) -> None: |
| with self.torch.autocast( |
| device_type=self.autocast_device, |
| dtype=self.autocast_dtype, |
| enabled=self.autocast_device == "cuda", |
| ): |
| self.state = self.processor.set_image(image) |
|
|
| def predict(self, prompt: str) -> PromptPrediction: |
| if self.state is None: |
| raise RuntimeError("begin_image must be called before predict") |
| with self.torch.autocast( |
| device_type=self.autocast_device, |
| dtype=self.autocast_dtype, |
| enabled=self.autocast_device == "cuda", |
| ): |
| output = self.processor.set_text_prompt(prompt, self.state) |
| masks_tensor = output["masks"] |
| scores_tensor = output["scores"] |
| boxes_tensor = output["boxes"] |
| if masks_tensor.numel() == 0: |
| shape = (int(output["original_height"]), int(output["original_width"])) |
| return PromptPrediction(np.zeros(shape, dtype=bool), 0.0, 0, []) |
| masks = masks_tensor.detach().cpu().numpy().astype(bool) |
| while masks.ndim > 3 and masks.shape[1] == 1: |
| masks = masks[:, 0] |
| union = np.any(masks, axis=0) |
| scores = scores_tensor.detach().float().cpu().numpy() |
| boxes = boxes_tensor.detach().float().cpu().numpy().tolist() |
| return PromptPrediction(union, float(scores.max()), int(len(scores)), boxes) |
|
|
|
|
| class GroundedSamBackend: |
| name = "grounded_sam" |
|
|
| def __init__(self, args: argparse.Namespace): |
| repo = Path(args.grounded_sam_repo).resolve() |
| config = Path(args.grounding_dino_config).resolve() |
| dino_checkpoint = Path(args.grounding_dino_checkpoint).resolve() |
| sam_checkpoint = Path(args.sam_checkpoint).resolve() |
| for path in (repo, config, dino_checkpoint, sam_checkpoint): |
| if not path.exists(): |
| raise FileNotFoundError(f"Required Grounded-SAM path not found: {path}") |
| sys.path[:0] = [str(repo), str(repo / "GroundingDINO"), str(repo / "segment_anything")] |
| import torch |
| import GroundingDINO.groundingdino.datasets.transforms as T |
| from GroundingDINO.groundingdino.models import build_model |
| from GroundingDINO.groundingdino.util.slconfig import SLConfig |
| from GroundingDINO.groundingdino.util.utils import clean_state_dict |
| from segment_anything import SamPredictor, sam_model_registry |
|
|
| model_args = SLConfig.fromfile(str(config)) |
| model_args.device = args.device |
| if args.bert_path: |
| model_args.bert_base_uncased_path = str(Path(args.bert_path).resolve()) |
| self.torch = torch |
| self.transforms = T.Compose( |
| [ |
| T.RandomResize([800], max_size=1333), |
| T.ToTensor(), |
| T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
| ] |
| ) |
| self.device = args.device |
| self.box_threshold = args.box_threshold |
| self.model = build_model(model_args) |
| checkpoint = torch.load(str(dino_checkpoint), map_location="cpu", weights_only=False) |
| self.model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) |
| self.model.to(args.device).eval() |
| sam = sam_model_registry[args.sam_model_type](checkpoint=str(sam_checkpoint)).to(args.device) |
| self.predictor = SamPredictor(sam) |
| self.image_pil: Image.Image | None = None |
| self.image_tensor = None |
| self.image_rgb: np.ndarray | None = None |
|
|
| def begin_image(self, image: Image.Image) -> None: |
| self.image_pil = image |
| self.image_rgb = np.asarray(image) |
| self.image_tensor, _ = self.transforms(image, None) |
| self.predictor.set_image(self.image_rgb) |
|
|
| def predict(self, prompt: str) -> PromptPrediction: |
| if self.image_pil is None or self.image_tensor is None or self.image_rgb is None: |
| raise RuntimeError("begin_image must be called before predict") |
| caption = prompt.lower().strip() |
| if not caption.endswith("."): |
| caption += "." |
| with self.torch.inference_mode(): |
| outputs = self.model(self.image_tensor[None].to(self.device), captions=[caption]) |
| logits = outputs["pred_logits"].detach().cpu().sigmoid()[0] |
| boxes = outputs["pred_boxes"].detach().cpu()[0] |
| confidence = logits.max(dim=1).values |
| keep = confidence > self.box_threshold |
| boxes = boxes[keep] |
| confidence = confidence[keep] |
| height, width = self.image_rgb.shape[:2] |
| if boxes.shape[0] == 0: |
| return PromptPrediction(np.zeros((height, width), dtype=bool), 0.0, 0, []) |
| scale = self.torch.tensor([width, height, width, height], dtype=boxes.dtype) |
| boxes = boxes * scale |
| boxes[:, :2] -= boxes[:, 2:] / 2 |
| boxes[:, 2:] += boxes[:, :2] |
| transformed = self.predictor.transform.apply_boxes_torch(boxes, (height, width)).to(self.device) |
| with self.torch.inference_mode(): |
| masks, sam_scores, _ = self.predictor.predict_torch( |
| point_coords=None, |
| point_labels=None, |
| boxes=transformed, |
| multimask_output=False, |
| ) |
| masks_np = masks[:, 0].detach().cpu().numpy().astype(bool) |
| combined_confidence = float( |
| math.sqrt(float(confidence.max()) * float(sam_scores.detach().cpu().max())) |
| ) |
| return PromptPrediction( |
| np.any(masks_np, axis=0), |
| combined_confidence, |
| int(len(masks_np)), |
| boxes.tolist(), |
| ) |
|
|
|
|
| def build_backend(args: argparse.Namespace): |
| if args.backend == "sam3": |
| return Sam3Backend(args) |
| if args.backend == "grounded_sam": |
| return GroundedSamBackend(args) |
| raise ValueError(f"Unsupported backend: {args.backend}") |
|
|
|
|
| def prompt_consensus( |
| predictions: list[PromptPrediction], minimum_area: int |
| ) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[str, Any]]: |
| if not predictions: |
| raise ValueError("At least one prompt prediction is required") |
| masks = np.stack([prediction.mask for prediction in predictions]) |
| votes = masks.sum(axis=0) |
| union = votes > 0 |
| majority = max(2, math.ceil(len(predictions) / 2)) if len(predictions) > 1 else 1 |
| core = votes >= majority |
| fallback = False |
| if not np.any(core) and np.any(union): |
| fallback = True |
| best_index = int(np.argmax([prediction.confidence for prediction in predictions])) |
| core = predictions[best_index].mask.copy() |
| candidate = retain_core_connected_union(union, core, minimum_area) |
| core = remove_small_components(core & candidate, minimum_area) |
| uncertain = candidate & ~core |
| union_area = int(candidate.sum()) |
| agreement = float(core.sum() / union_area) if union_area else 0.0 |
| confidences = [prediction.confidence for prediction in predictions] |
| metadata = { |
| "majority_vote": majority, |
| "fallback_to_best_prompt": fallback, |
| "agreement": round(agreement, 6), |
| "mean_prompt_confidence": round(float(np.mean(confidences)), 6), |
| "maximum_prompt_confidence": round(float(np.max(confidences)), 6), |
| "detected_prompt_count": sum(value > 0 for value in confidences), |
| } |
| return candidate, core, uncertain, metadata |
|
|
|
|
| def overlay_masks( |
| image: np.ndarray, target: np.ndarray, obstacle: np.ndarray, uncertain: np.ndarray |
| ) -> np.ndarray: |
| overlay = image.astype(np.float32).copy() |
| colors = [ |
| (target, np.asarray([30, 210, 70], dtype=np.float32), 0.38), |
| (uncertain, np.asarray([250, 210, 30], dtype=np.float32), 0.55), |
| |
| (obstacle, np.asarray([230, 45, 45], dtype=np.float32), 0.58), |
| ] |
| for mask, color, alpha in colors: |
| overlay[mask] = overlay[mask] * (1.0 - alpha) + color * alpha |
| return np.clip(overlay, 0, 255).astype(np.uint8) |
|
|
|
|
| def score_category( |
| mask: np.ndarray, |
| consensus: dict[str, Any], |
| thresholds: dict[str, float], |
| category: str, |
| ) -> tuple[float, dict[str, float]]: |
| area_ratio = float(mask.mean()) |
| plausibility = area_plausibility( |
| area_ratio, |
| thresholds["minimum_target_area_ratio"], |
| thresholds["maximum_target_area_ratio"], |
| ) |
| bottom_contact = float(mask[int(mask.shape[0] * 0.8) :, :].mean() > 0.005) |
| confidence = float(consensus["mean_prompt_confidence"]) |
| agreement = float(consensus["agreement"]) |
| evidence = 0.55 * confidence + 0.35 * agreement + 0.10 * bottom_contact |
| score = evidence * (0.20 + 0.80 * plausibility) |
| if category == "walkway": |
| score *= 0.94 |
| signals = { |
| "area_ratio": round(area_ratio, 6), |
| "area_plausibility": round(plausibility, 6), |
| "bottom_contact": bottom_contact, |
| "score": round(score, 6), |
| } |
| return score, signals |
|
|
|
|
| def predict_prompt_set( |
| backend, |
| prompts: list[str], |
| output_dir: Path, |
| minimum_component_area: int, |
| save_prompt_masks: bool = True, |
| ) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[str, Any]]: |
| predictions = [] |
| prompt_metadata = [] |
| for prompt in prompts: |
| prediction = backend.predict(prompt) |
| predictions.append(prediction) |
| if save_prompt_masks: |
| save_mask(output_dir / f"{safe_name(prompt)}.png", prediction.mask) |
| prompt_metadata.append( |
| { |
| "prompt": prompt, |
| "confidence": round(prediction.confidence, 6), |
| "instance_count": prediction.instance_count, |
| "boxes_xyxy": prediction.boxes, |
| } |
| ) |
| candidate, core, uncertain, consensus = prompt_consensus( |
| predictions, minimum_component_area |
| ) |
| consensus["prompts"] = prompt_metadata |
| return candidate, core, uncertain, consensus |
|
|
|
|
| def process_sample( |
| row: dict[str, Any], |
| backend, |
| config: dict[str, Any], |
| output_root: Path, |
| args: argparse.Namespace, |
| ) -> dict[str, Any]: |
| sample_id = row["sample_id"] |
| sample_dir = output_root / "samples" / sample_id |
| quality_path = sample_dir / "quality.json" |
| if args.resume: |
| quality = load_resume_quality(sample_dir, sample_id, backend.name) |
| if quality is not None: |
| return quality |
|
|
| |
| |
| |
| image = ImageOps.exif_transpose(Image.open(row["image_path"])).convert("RGB") |
| rgb = np.asarray(image) |
| image_area = rgb.shape[0] * rgb.shape[1] |
| minimum_component_area = max(args.minimum_component_area, int(image_area * 0.0002)) |
| backend.begin_image(image) |
|
|
| requested_categories = args.categories or list(config["categories"]) |
| fixed_category = row.get("category") |
| if fixed_category: |
| requested_categories = [fixed_category] |
| if args.verify_manifest_category and not row.get("category_reviewed", False): |
| requested_categories.extend( |
| config.get("category_confusions", {}).get(fixed_category, []) |
| ) |
| requested_categories = list(dict.fromkeys(requested_categories)) |
|
|
| category_results: dict[str, dict[str, Any]] = {} |
| masks: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]] = {} |
| thresholds = config["quality_thresholds"] |
| save_full_artifacts = args.artifact_level == "full" |
| for category in requested_categories: |
| if category not in config["categories"]: |
| raise ValueError(f"Unknown category in pilot manifest: {category}") |
| prompt_dir = sample_dir / "prompts" / category |
| candidate, core, uncertain, consensus = predict_prompt_set( |
| backend, |
| config["categories"][category]["prompts"], |
| prompt_dir, |
| minimum_component_area, |
| save_prompt_masks=save_full_artifacts, |
| ) |
| |
| cat_max_area = config["categories"][category].get( |
| "max_area_ratio", |
| config.get("tactile_paving_max_area_ratio", None), |
| ) |
| candidate_area_ratio = float(candidate.mean()) if candidate.any() else 0.0 |
| if cat_max_area is not None and candidate_area_ratio > float(cat_max_area): |
| consensus["area_cap_rejected"] = True |
| consensus["area_cap_ratio"] = round(candidate_area_ratio, 6) |
| consensus["area_cap_limit"] = float(cat_max_area) |
| candidate = np.zeros_like(candidate) |
| core = np.zeros_like(core) |
| uncertain = np.zeros_like(uncertain) |
| else: |
| consensus["area_cap_rejected"] = False |
| |
| score, signals = score_category(candidate, consensus, thresholds, category) |
| if save_full_artifacts: |
| save_mask(sample_dir / "categories" / f"{category}_candidate.png", candidate) |
| save_mask(sample_dir / "categories" / f"{category}_core.png", core) |
| save_mask(sample_dir / "categories" / f"{category}_uncertain.png", uncertain) |
| masks[category] = (candidate, core, uncertain) |
| category_results[category] = {"consensus": consensus, **signals} |
|
|
| best_scored_category = max( |
| category_results, key=lambda key: category_results[key]["score"] |
| ) |
| |
| |
| |
| selected_category = fixed_category or best_scored_category |
| category_disagreement = bool( |
| fixed_category and best_scored_category != fixed_category |
| ) |
| target, target_core, target_uncertain = masks[selected_category] |
|
|
| prompt_resolution = obstacle_prompt_resolution( |
| config, selected_category, backend.name |
| ) |
| obstacle_predictions: list[tuple[str, PromptPrediction]] = [] |
| obstacle_metadata = [] |
| for prompt in prompt_resolution["prompts"]: |
| prediction = backend.predict(prompt) |
| obstacle_predictions.append((prompt, prediction)) |
| if save_full_artifacts: |
| save_mask( |
| sample_dir / "prompts" / "obstacles" / f"{safe_name(prompt)}.png", |
| prediction.mask, |
| ) |
| obstacle_metadata.append( |
| { |
| "prompt": prompt, |
| "group": "dynamic_or_compact", |
| "confidence": round(prediction.confidence, 6), |
| "instance_count": prediction.instance_count, |
| "boxes_xyxy": prediction.boxes, |
| } |
| ) |
| maximum_ratios = config.get("obstacle_maximum_area_ratio", {}) |
| filtered_dynamic_obstacles = [] |
| for prompt, prediction in obstacle_predictions: |
| cleaned = remove_small_components(prediction.mask, minimum_component_area) |
| maximum_area_ratio = float( |
| maximum_ratios.get(prompt, maximum_ratios.get("default", 0.15)) |
| ) |
| filtered_dynamic_obstacles.append( |
| retain_nearby_obstacles( |
| cleaned, |
| target, |
| minimum_component_area, |
| maximum_area_ratio=maximum_area_ratio, |
| ) |
| ) |
| dynamic_obstacle = ( |
| np.any(np.stack(filtered_dynamic_obstacles), axis=0) |
| if filtered_dynamic_obstacles |
| else np.zeros_like(target) |
| ) |
|
|
| barrier_predictions: list[tuple[str, PromptPrediction]] = [] |
| barrier_backends = config.get("barrier_obstacle_backends", ["sam3"]) |
| barrier_prompts = ( |
| config.get("barrier_obstacle_prompts", []) |
| if backend.name in barrier_backends |
| else [] |
| ) |
| for prompt in barrier_prompts: |
| prediction = backend.predict(prompt) |
| barrier_predictions.append((prompt, prediction)) |
| if save_full_artifacts: |
| save_mask( |
| sample_dir / "prompts" / "obstacles" / f"{safe_name(prompt)}.png", |
| prediction.mask, |
| ) |
| obstacle_metadata.append( |
| { |
| "prompt": prompt, |
| "group": "barrier_consensus", |
| "confidence": round(prediction.confidence, 6), |
| "instance_count": prediction.instance_count, |
| "boxes_xyxy": prediction.boxes, |
| } |
| ) |
| if barrier_predictions: |
| barrier_votes = np.stack( |
| [prediction.mask for _, prediction in barrier_predictions] |
| ).sum(axis=0) |
| required_votes = max(2, math.ceil(len(barrier_predictions) / 2)) |
| barrier_obstacle = barrier_votes >= required_votes |
| barrier_obstacle = retain_compact_obstacles( |
| barrier_obstacle, minimum_component_area |
| ) |
| else: |
| barrier_obstacle = np.zeros_like(target) |
|
|
| |
| |
| |
| |
| |
| evidence_policy = resolve_obstacle_evidence_policy(config, backend.name) |
| evidence_minimum_area = evidence_minimum_component_area(evidence_policy, image_area) |
| unfiltered_evidence_masks: list[np.ndarray] = [] |
| trusted_evidence_masks: list[np.ndarray] = [] |
| obstacle_evidence_decisions: list[dict[str, Any]] = [] |
| for group, predictions in ( |
| ("dynamic_or_compact", obstacle_predictions), |
| ("barrier_consensus", barrier_predictions), |
| ): |
| for prompt, prediction in predictions: |
| unfiltered, trusted, decision = trusted_obstacle_prompt_evidence( |
| prompt=prompt, |
| group=group, |
| prediction=prediction, |
| minimum_component_area=evidence_minimum_area, |
| policy=evidence_policy, |
| ) |
| unfiltered_evidence_masks.append(unfiltered) |
| trusted_evidence_masks.append(trusted) |
| obstacle_evidence_decisions.append(decision) |
| obstacle_unfiltered_evidence = ( |
| np.any(np.stack(unfiltered_evidence_masks), axis=0) |
| if unfiltered_evidence_masks |
| else np.zeros_like(target) |
| ) |
| obstacle_all_detected = ( |
| np.any(np.stack(trusted_evidence_masks), axis=0) |
| if trusted_evidence_masks |
| else np.zeros_like(target) |
| ) |
| obstacle = dynamic_obstacle | barrier_obstacle |
| target = target & ~obstacle |
| target_core = target_core & target |
| target_uncertain = target_uncertain & target |
|
|
| kernel_size = max(5, int(round(min(rgb.shape[:2]) * 0.012)) | 1) |
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) |
| target_dilated = cv2.dilate(target.astype(np.uint8), kernel).astype(bool) |
| obstacle_dilated = cv2.dilate(obstacle.astype(np.uint8), kernel).astype(bool) |
| adjacency = target_dilated & obstacle_dilated |
| adjacency_ratio = float(adjacency.sum() / max(target.sum(), 1)) |
|
|
| save_mask(sample_dir / "target_visible_candidate.png", target) |
| save_mask(sample_dir / "target_visible_core.png", target_core) |
| save_mask(sample_dir / "target_uncertain.png", target_uncertain) |
| save_mask(sample_dir / "obstacle_unfiltered_evidence.png", obstacle_unfiltered_evidence) |
| save_mask(sample_dir / "obstacle_all_detected.png", obstacle_all_detected) |
| save_mask(sample_dir / "obstacle.png", obstacle) |
| save_red_mask(sample_dir / "obstacle_red.png", obstacle) |
| Image.fromarray(overlay_masks(rgb, target_core, obstacle, target_uncertain)).save( |
| sample_dir / "overlay.png", quality=92 |
| ) |
|
|
| selected = category_results[selected_category] |
| score = float(selected["score"]) |
| target_area_ratio = float(target.mean()) |
| uncertainty_ratio = float(target_uncertain.sum() / max(target.sum(), 1)) |
| obstacle_area_ratio = float(obstacle.mean()) |
| area_valid = ( |
| thresholds["minimum_target_area_ratio"] |
| <= target_area_ratio |
| <= thresholds["maximum_target_area_ratio"] |
| ) |
| if not area_valid: |
| review_status = "reject_or_reprompt" |
| elif ( |
| score >= thresholds["auto_accept_score"] |
| and not selected["consensus"]["fallback_to_best_prompt"] |
| and uncertainty_ratio <= 0.35 |
| and obstacle_area_ratio <= 0.20 |
| ): |
| review_status = "candidate_accept_after_visual_review" |
| elif score >= thresholds["manual_review_score"]: |
| review_status = "manual_review" |
| else: |
| review_status = "reject_or_reprompt" |
| if ( |
| category_disagreement |
| and not row.get("category_reviewed", False) |
| and review_status == "candidate_accept_after_visual_review" |
| ): |
| review_status = "manual_review" |
| quality = { |
| "sample_id": sample_id, |
| "image_path": row["image_path"], |
| "backend": backend.name, |
| "selected_category": selected_category, |
| "category_verification": { |
| "manifest_category": fixed_category, |
| "best_scored_category": best_scored_category, |
| "category_disagreement": category_disagreement, |
| "category_reviewed": bool(row.get("category_reviewed", False)), |
| }, |
| "quality_score": round(score, 6), |
| "review_status": review_status, |
| "target_area_ratio": round(target_area_ratio, 6), |
| "target_core_ratio": round(float(target_core.mean()), 6), |
| "uncertainty_ratio_within_target": round( |
| uncertainty_ratio, 6 |
| ), |
| "obstacle_area_ratio": round(obstacle_area_ratio, 6), |
| "obstacle_unfiltered_evidence_pixels": int(obstacle_unfiltered_evidence.sum()), |
| "obstacle_unfiltered_evidence_area_ratio": round( |
| float(obstacle_unfiltered_evidence.mean()), 6 |
| ), |
| "obstacle_all_detected_pixels": int(obstacle_all_detected.sum()), |
| "obstacle_all_detected_area_ratio": round(float(obstacle_all_detected.mean()), 6), |
| "obstacle_all_detected_policy": "sanitized_high_recall_prompt_union_v2", |
| "obstacle_evidence_policy_version": OBSTACLE_EVIDENCE_POLICY_VERSION, |
| "obstacle_evidence_minimum_component_area_pixels": evidence_minimum_area, |
| "obstacle_evidence_artifact_semantics": { |
| "obstacle_unfiltered_evidence.png": ( |
| "union of every obstacle/barrier prompt after only small-component removal; " |
| "audit evidence only, never direct hidden-completion support" |
| ), |
| "obstacle_all_detected.png": ( |
| "high-recall prompt union after confidence, per-prompt area, full-frame, " |
| "and route-like sanity filters; candidate evidence for later hidden completion" |
| ), |
| "obstacle.png": ( |
| "unchanged nearby/compact obstacle proposal used for visible-mask cleanup; " |
| "not the high-recall evidence layer" |
| ), |
| }, |
| "obstacle_prompt_resolution": prompt_resolution, |
| "obstacle_evidence_policy": evidence_policy, |
| "obstacle_evidence_prompt_decisions": obstacle_evidence_decisions, |
| "obstacle_evidence_decision_counts": dict( |
| sorted(Counter(item["decision"] for item in obstacle_evidence_decisions).items()) |
| ), |
| "barrier_obstacle_area_ratio": round(float(barrier_obstacle.mean()), 6), |
| "target_obstacle_adjacency_ratio": round(adjacency_ratio, 6), |
| "category_results": category_results, |
| "obstacle_prompts": obstacle_metadata, |
| "automatic_mask_is_ground_truth": False, |
| "artifact_level": args.artifact_level, |
| } |
| json_dump(quality_path, quality) |
| return quality |
|
|
|
|
| def build_html_report(output_dir: Path, results: list[dict[str, Any]]) -> None: |
| rows = [] |
| for result in sorted(results, key=lambda item: item.get("quality_score", -1), reverse=True): |
| sample_id = result["sample_id"] |
| overlay = f"samples/{sample_id}/overlay.png" |
| error = result.get("error") |
| details = html.escape(error) if error else ( |
| f'{html.escape(result["selected_category"])} | ' |
| f'score={result["quality_score"]:.3f} | ' |
| f'{html.escape(result["review_status"])}' |
| ) |
| rows.append( |
| f'<article><img src="{html.escape(overlay)}" loading="lazy">' |
| f'<div><strong>{html.escape(sample_id)}</strong><br>{details}</div></article>' |
| ) |
| document = """<!doctype html><meta charset="utf-8"><title>Accessibility mask pilot</title> |
| <style>body{font-family:sans-serif;margin:20px}main{display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:16px}article{border:1px solid #bbb;padding:8px}img{width:100%;height:auto}div{margin-top:6px;font-size:14px}</style> |
| <h1>Accessibility mask pilot</h1><p>Green: consensus core; yellow: uncertain target; red: obstacle. All masks require visual review.</p><main>""" |
| document += "\n".join(rows) + "</main>" |
| (output_dir / "report.html").write_text(document, encoding="utf-8") |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--pilot-manifest", default="output/accessibility_clip_screen/pilot_manifest.jsonl") |
| parser.add_argument("--image", default=None, help="Single-image mode input image; bypasses --pilot-manifest.") |
| parser.add_argument("--single-sample-id", default=None, help="Stable output ID used with --image.") |
| parser.add_argument( |
| "--single-category", |
| choices=["curb_cut", "ramp", "stairs", "tactile_paving", "walkway"], |
| default=None, |
| help="Known category for --image. Omit to score all configured categories.", |
| ) |
| parser.add_argument("--prompt-config", default="configs/accessibility_mask_prompts.json") |
| parser.add_argument( |
| "--prompt-config-override", |
| default=None, |
| help=( |
| "Optional JSON overlay merged into --prompt-config for this output-only run. " |
| "Nested prompt lists are appended and de-duplicated; the base config is never edited." |
| ), |
| ) |
| parser.add_argument("--output-dir", default="output/accessibility_mask_proposals/sam3") |
| parser.add_argument("--backend", choices=["sam3", "grounded_sam"], default="sam3") |
| parser.add_argument("--device", default="cuda") |
| parser.add_argument("--categories", nargs="*", default=None) |
| parser.add_argument("--limit", type=int, default=0) |
| parser.add_argument("--minimum-component-area", type=int, default=128) |
| parser.add_argument("--confidence-threshold", type=float, default=0.38) |
| parser.add_argument("--compile", action="store_true") |
| parser.add_argument("--resume", action="store_true") |
| parser.add_argument( |
| "--artifact-level", |
| choices=["full", "standard"], |
| default="full", |
| help="standard omits regenerable per-prompt/category PNGs for large batch runs.", |
| ) |
| parser.add_argument( |
| "--verify-manifest-category", |
| action="store_true", |
| help="Compare configured confusable categories unless category_reviewed=true.", |
| ) |
| parser.add_argument( |
| "--allow-test-split", |
| action="store_true", |
| help=( |
| "Explicitly allow frozen test RGB rows for a pre-registered front-end " |
| "inference run. Default behavior remains validation-only." |
| ), |
| ) |
| parser.add_argument("--sam3-repo", default="repos/sam3") |
| parser.add_argument("--sam3-checkpoint", default="weights/sam3/sam3.pt") |
| parser.add_argument( |
| "--grounded-sam-repo", |
| default="repos/grounded-sam-compat", |
| ) |
| parser.add_argument( |
| "--grounding-dino-config", |
| default="repos/grounded-sam-compat/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py", |
| ) |
| parser.add_argument( |
| "--grounding-dino-checkpoint", |
| default="weights/grounded-sam/groundingdino_swint_ogc.pth", |
| ) |
| parser.add_argument( |
| "--sam-checkpoint", |
| default="weights/grounded-sam/sam_vit_h_4b8939.pth", |
| ) |
| parser.add_argument("--sam-model-type", choices=["vit_h", "vit_l", "vit_b"], default="vit_h") |
| parser.add_argument("--bert-path", default=None) |
| parser.add_argument("--box-threshold", type=float, default=0.28) |
| return parser |
|
|
|
|
| def main() -> int: |
| args = build_parser().parse_args() |
| config_path = Path(args.prompt_config).resolve() |
| config_override_path = ( |
| Path(args.prompt_config_override).resolve() |
| if args.prompt_config_override |
| else None |
| ) |
| output_dir = Path(args.output_dir).resolve() |
| output_dir.mkdir(parents=True, exist_ok=True) |
| with config_path.open("r", encoding="utf-8") as handle: |
| config = json.load(handle) |
| if not isinstance(config, dict): |
| raise ValueError(f"prompt config must be a JSON object: {config_path}") |
| if config_override_path is not None: |
| with config_override_path.open("r", encoding="utf-8") as handle: |
| override = json.load(handle) |
| if not isinstance(override, dict): |
| raise ValueError(f"prompt config override must be a JSON object: {config_override_path}") |
| config = merge_prompt_config(config, override) |
| if args.image: |
| image_path = Path(args.image).expanduser().resolve() |
| if not image_path.is_file(): |
| raise FileNotFoundError(image_path) |
| sample_id = args.single_sample_id or safe_name(image_path.stem) or "single_image" |
| rows = [ |
| { |
| "sample_id": sample_id, |
| "image_path": str(image_path), |
| "category": args.single_category, |
| "category_reviewed": bool(args.single_category), |
| } |
| ] |
| else: |
| pilot_path = Path(args.pilot_manifest).resolve() |
| rows = read_jsonl(pilot_path) |
| if args.limit > 0: |
| rows = rows[: args.limit] |
| allowed_splits = {None, "validation"} |
| if args.allow_test_split: |
| allowed_splits.add("test") |
| invalid_splits = sorted( |
| {str(row.get("benchmark_split")) for row in rows if row.get("benchmark_split") not in allowed_splits} |
| ) |
| if invalid_splits: |
| raise ValueError( |
| "Proposal manifest contains unsupported benchmark split(s): " |
| f"{invalid_splits}. Pass --allow-test-split only for an explicitly " |
| "registered test front-end inference run." |
| ) |
| test_accessed = any(row.get("benchmark_split") == "test" for row in rows) |
| backend = None |
| initialization_error = None |
| try: |
| backend = build_backend(args) |
| if backend.torch.cuda.is_available(): |
| backend.torch.cuda.reset_peak_memory_stats() |
| except Exception as exc: |
| traceback.print_exc() |
| initialization_error = f"{type(exc).__name__}: {exc}" |
| results = [] |
| if initialization_error is not None: |
| for row in rows: |
| results.append( |
| { |
| "sample_id": row["sample_id"], |
| "image_path": row["image_path"], |
| "backend": args.backend, |
| "quality_score": -1.0, |
| "review_status": "error", |
| "error": f"model_initialization_failed: {initialization_error}", |
| } |
| ) |
| write_jsonl(output_dir / "results.jsonl", results) |
| else: |
| assert backend is not None |
| for index, row in enumerate(rows, start=1): |
| print(f"[{index}/{len(rows)}] {row['sample_id']}", flush=True) |
| try: |
| result = process_sample(row, backend, config, output_dir, args) |
| except Exception as exc: |
| traceback.print_exc() |
| result = { |
| "sample_id": row["sample_id"], |
| "image_path": row["image_path"], |
| "backend": args.backend, |
| "quality_score": -1.0, |
| "review_status": "error", |
| "error": f"{type(exc).__name__}: {exc}", |
| } |
| results.append(result) |
| write_jsonl(output_dir / "results.jsonl", results) |
| build_html_report(output_dir, results) |
| failed = sum(item.get("review_status") == "error" for item in results) |
| peak_bytes = ( |
| int(backend.torch.cuda.max_memory_allocated()) |
| if backend is not None and backend.torch.cuda.is_available() |
| else 0 |
| ) |
| summary = { |
| "backend": args.backend, |
| "sample_count": len(results), |
| "completed": len(results) - failed, |
| "failed": failed, |
| "failure_rate": failed / len(results) if results else 1.0, |
| "status_counts": dict( |
| sorted( |
| { |
| status: sum(item.get("review_status") == status for item in results) |
| for status in {item.get("review_status") for item in results} |
| }.items() |
| ) |
| ), |
| "mean_quality_score": round( |
| float(np.mean([item["quality_score"] for item in results if item["quality_score"] >= 0])), |
| 6, |
| ) |
| if any(item["quality_score"] >= 0 for item in results) |
| else None, |
| "automatic_masks_are_ground_truth": False, |
| "prompt_config": str(config_path), |
| "prompt_config_override": str(config_override_path) if config_override_path else None, |
| "model_initialization_error": initialization_error, |
| "peak_gpu_memory_bytes": peak_bytes, |
| "peak_gpu_memory_gib": peak_bytes / (1024**3), |
| "environment": { |
| "python": sys.version, |
| "executable": sys.executable, |
| "torch": getattr(getattr(backend, "torch", None), "__version__", None), |
| "cuda_runtime": getattr(getattr(backend, "torch", None), "version", None).cuda |
| if backend is not None |
| else None, |
| "gpu_name": backend.torch.cuda.get_device_name(0) |
| if backend is not None and backend.torch.cuda.is_available() |
| else None, |
| "slurm_job_id": os.environ.get("SLURM_JOB_ID"), |
| }, |
| "resolved_arguments": vars(args), |
| "test_accessed": test_accessed, |
| "test_access_explicitly_allowed": bool(args.allow_test_split), |
| } |
| json_dump(output_dir / "summary.json", summary) |
| print(json.dumps(summary, ensure_ascii=False, indent=2)) |
| return 0 if all(item.get("review_status") != "error" for item in results) else 2 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|