scripts/evaluate_sam3_test_text_prompt_segmentation.py

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