#!/usr/bin/env python3 """SAM3 zero-shot text-prompt segmentation on BUS-CoT test split only.""" from __future__ import annotations import argparse import json import os import re import sys from pathlib import Path import numpy as np import pandas as pd from PIL import Image, ImageDraw, ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True SCRIPT_DIR = Path(__file__).resolve().parent BUNDLE_ROOT = SCRIPT_DIR.parent DEFAULT_SAM3 = str(Path(os.environ.get("SAM3_CHECKPOINT", BUNDLE_ROOT / "model" / "sam3_base.pt"))) SEG_SCRIPTS = SCRIPT_DIR if str(SEG_SCRIPTS) not in sys.path: sys.path.insert(0, str(SEG_SCRIPTS)) def parse_args(): p = argparse.ArgumentParser() p.add_argument("--index_csv", default="data/test_only_grounding_segmentation_baselines/index/test_index.csv") p.add_argument("--sam3_model_path", default=DEFAULT_SAM3) p.add_argument("--checkpoint_path", default=None) p.add_argument("--output_dir", default="outputs/test_only_grounding_segmentation_baselines/sam3_text_prompt") p.add_argument("--prompts", nargs="+", default=["breast tumor", "breast lesion", "breast mass", "ultrasound breast lesion", "hypoechoic breast mass"]) p.add_argument("--max_samples", type=int, default=None) p.add_argument("--overwrite", action="store_true") p.add_argument("--resume", action="store_true") p.add_argument("--image_size", type=int, default=512) p.add_argument("--threshold", type=float, default=0.5) p.add_argument("--device", default="cuda") p.add_argument("--encoder_trainable", choices=["frozen", "lora", "last_block"], default="frozen") p.add_argument("--lora_rank", type=int, default=8) p.add_argument("--lora_alpha", type=float, default=16) return p.parse_args() def slug(text: str) -> str: return re.sub(r"[^A-Za-z0-9]+", "_", text).strip("_").lower() def read_mask(path: str) -> np.ndarray: return np.asarray(Image.open(path).convert("L")) > 0 def mask_bbox(mask: np.ndarray) -> list[int] | None: ys, xs = np.where(mask) if xs.size == 0: return None return [int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1] def overlap(pred: np.ndarray, gt: np.ndarray) -> dict: if pred.shape != gt.shape: pred = np.asarray(Image.fromarray(pred.astype(np.uint8) * 255).resize((gt.shape[1], gt.shape[0]), Image.NEAREST)) > 0 inter = int(np.logical_and(pred, gt).sum()) p, g = int(pred.sum()), int(gt.sum()) union = int(np.logical_or(pred, gt).sum()) return { "dice": 2 * inter / (p + g) if p + g else 1.0, "mask_iou": inter / union if union else 1.0, "pred_area": p, "gt_area": g, "area_ratio": p / g if g else np.nan, } def bbox_iou(a: list[int] | None, b: list[int] | None) -> float: if not a or not b: return np.nan ax1, ay1, ax2, ay2 = a bx1, by1, bx2, by2 = b ix1, iy1, ix2, iy2 = max(ax1, bx1), max(ay1, by1), min(ax2, bx2), min(ay2, by2) inter = max(0, ix2 - ix1) * max(0, iy2 - iy1) aa = max(0, ax2 - ax1) * max(0, ay2 - ay1) ba = max(0, bx2 - bx1) * max(0, by2 - by1) return inter / (aa + ba - inter) if aa + ba - inter else 0.0 def load_predictor(args, prompt: str): from sam3_buscot_runner import SAM3BuscotPredictor return SAM3BuscotPredictor( args.sam3_model_path, checkpoint_path=args.checkpoint_path, prompt_type="semantic_text", prompt_text=prompt, image_size=args.image_size, device=args.device, encoder_trainable=args.encoder_trainable, lora_rank=args.lora_rank, lora_alpha=args.lora_alpha, threshold=args.threshold, ) def save_mask(mask: np.ndarray, path: Path) -> str: path.parent.mkdir(parents=True, exist_ok=True) Image.fromarray(mask.astype(np.uint8) * 255).save(path) return str(path) def save_overlay(image_path: str, gt: np.ndarray, pred: np.ndarray, out_path: Path): img = Image.open(image_path).convert("RGB") gt_r = Image.fromarray(gt.astype(np.uint8) * 255).resize(img.size, Image.NEAREST) pred_r = Image.fromarray(pred.astype(np.uint8) * 255).resize(img.size, Image.NEAREST) base = np.asarray(img).copy() gt_m = np.asarray(gt_r) > 0 pred_m = np.asarray(pred_r) > 0 base[gt_m] = (0.55 * base[gt_m] + np.array([0, 255, 0]) * 0.45).astype(np.uint8) base[pred_m] = (0.55 * base[pred_m] + np.array([255, 0, 0]) * 0.45).astype(np.uint8) out_path.parent.mkdir(parents=True, exist_ok=True) Image.fromarray(base).save(out_path) def summarize(df: pd.DataFrame, prompt: str) -> dict: valid = df[df["valid_mask"].astype(bool)] return { "prompt": prompt, "N_total": int(len(df)), "N_evaluated": int(len(valid)), "valid_mask_count": int(df["valid_mask"].sum()) if len(df) else 0, "failure_rate": float((~df["valid_mask"].astype(bool)).mean()) if len(df) else np.nan, "mean Dice": float(valid["dice"].mean()) if len(valid) else np.nan, "median Dice": float(valid["dice"].median()) if len(valid) else np.nan, "mean IoU": float(valid["mask_iou"].mean()) if len(valid) else np.nan, "median IoU": float(valid["mask_iou"].median()) if len(valid) else np.nan, "Dice > 0.3 rate": float((valid["dice"] > 0.3).mean()) if len(valid) else np.nan, "Dice > 0.5 rate": float((valid["dice"] > 0.5).mean()) if len(valid) else np.nan, "Dice > 0.7 rate": float((valid["dice"] > 0.7).mean()) if len(valid) else np.nan, "Dice > 0.8 rate": float((valid["dice"] > 0.8).mean()) if len(valid) else np.nan, "mean predicted area": float(valid["pred_area"].mean()) if len(valid) else np.nan, "median predicted area": float(valid["pred_area"].median()) if len(valid) else np.nan, "mean area ratio": float(valid["area_ratio"].mean()) if len(valid) else np.nan, "mean bbox IoU": float(valid["bbox_iou_predmaskbbox_vs_gtmaskbbox"].mean()) if len(valid) else np.nan, "median bbox IoU": float(valid["bbox_iou_predmaskbbox_vs_gtmaskbbox"].median()) if len(valid) else np.nan, } def main(): args = parse_args() out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) index = pd.read_csv(args.index_csv) index = index[index["has_original_image"].astype(bool) & index["has_gt_mask"].astype(bool)].copy() if args.max_samples: index = index.head(args.max_samples) summaries = [] for prompt in args.prompts: ps = slug(prompt) pred_csv = out_dir / f"predictions_{ps}.csv" print(f"[prompt] start {prompt} -> {pred_csv}", flush=True) if pred_csv.exists() and args.resume and not args.overwrite: df = pd.read_csv(pred_csv) expected = len(index) if len(df) >= expected: print(f"[prompt] resume skip complete {prompt}: {len(df)}/{expected}", flush=True) summaries.append(summarize(df, prompt)) continue print(f"[prompt] resume found incomplete {prompt}: {len(df)}/{expected}; rerunning", flush=True) rows = [] predictor = None predictor_error = "" try: predictor = load_predictor(args, prompt) except Exception as exc: predictor_error = f"{type(exc).__name__}: {exc}" for _, row in index.iterrows(): sid = str(row["sample_id"]) gt = read_mask(row["gt_mask_path"]) gt_box = json.loads(row["gt_mask_bbox"]) if isinstance(row["gt_mask_bbox"], str) and row["gt_mask_bbox"] else mask_bbox(gt) rec = { "sample_id": sid, "dataset": row.get("dataset", ""), "source_split": row.get("source_split", row.get("split", "")), "original_image_path": row["original_image_path"], "gt_mask_path": row["gt_mask_path"], "prompt": prompt, } try: if predictor is None: raise RuntimeError(predictor_error or "SAM3 predictor unavailable") pred, details = predictor.predict(row["original_image_path"]) pred = pred > 0 valid = bool(pred.sum() > 0) failure = "" if valid else "empty_mask" pred_path = save_mask(pred, out_dir / f"pred_masks_{ps}" / f"{sid.replace('/', '_')}_mask.png") om = overlap(pred, gt) pred_box = mask_bbox(pred) rec.update({ "pred_mask_path": pred_path, "valid_mask": valid, "pred_mask_bbox": json.dumps(pred_box) if pred_box else "", "gt_mask_bbox": json.dumps(gt_box) if gt_box else "", "bbox_iou_predmaskbbox_vs_gtmaskbbox": bbox_iou(pred_box, gt_box), "failure_type": failure, **om, **details, }) except Exception as exc: rec.update({ "pred_mask_path": "", "valid_mask": False, "dice": np.nan, "mask_iou": np.nan, "pred_area": 0, "gt_area": int(gt.sum()), "area_ratio": np.nan, "pred_mask_bbox": "", "gt_mask_bbox": json.dumps(gt_box) if gt_box else "", "bbox_iou_predmaskbbox_vs_gtmaskbbox": np.nan, "failure_type": f"{type(exc).__name__}: {exc}", }) rows.append(rec) df = pd.DataFrame(rows) df.to_csv(pred_csv, index=False) print(f"[prompt] wrote {prompt}: {len(df)} rows", flush=True) summary = summarize(df, prompt) summaries.append(summary) pd.DataFrame([summary]).to_csv(out_dir / f"metrics_{ps}.csv", index=False) if "dataset" in df.columns and df["dataset"].astype(str).str.len().gt(0).any(): by_dataset = [] for dataset, group in df.groupby("dataset"): by_dataset.append({"dataset": dataset, **summarize(group, prompt)}) pd.DataFrame(by_dataset).to_csv(out_dir / f"metrics_by_dataset_{ps}.csv", index=False) examples = pd.concat([df.sort_values("dice").head(20), df.sort_values("dice", ascending=False).head(20)]).drop_duplicates("sample_id") for _, ex in examples.iterrows(): if not ex.get("pred_mask_path"): continue gt = read_mask(ex["gt_mask_path"]) pred = read_mask(ex["pred_mask_path"]) save_overlay(ex["original_image_path"], gt, pred, out_dir / f"overlay_examples_{ps}" / f"{ex['sample_id']}.png") all_df = pd.DataFrame(summaries) all_df.to_csv(out_dir / "all_prompt_metrics.csv", index=False) all_df.sort_values("mean Dice", ascending=False).to_csv(out_dir / "best_prompt_summary.csv", index=False) print(f"wrote SAM3 text-prompt results to {out_dir}", flush=True) if __name__ == "__main__": main()