| |
| """Evaluate deployable multi-prompt mask stability as an offline QC signal.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import itertools |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| from PIL import Image |
|
|
|
|
| def parse_args(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--predictions_csv", default="outputs/sam3_multiprompt_lora_qc/multiprompt_lora_eval/all_prompt_predictions.csv") |
| p.add_argument("--index_csv", default="data/sam3_multiprompt_lora_qc/index/test_index.csv") |
| p.add_argument("--output_dir", default="outputs/sam3_multiprompt_lora_qc/stability_qc") |
| p.add_argument("--default_prompt", default="breast lesion") |
| p.add_argument("--overwrite", action="store_true") |
| return p.parse_args() |
|
|
|
|
| def read(path): |
| return np.asarray(Image.open(path).convert("L")) > 0 |
|
|
|
|
| def align(a, shape): |
| return a if a.shape == shape else np.asarray(Image.fromarray(a.astype(np.uint8) * 255).resize((shape[1], shape[0]), Image.Resampling.NEAREST)) > 0 |
|
|
|
|
| def overlap(a, b): |
| b = align(b, a.shape) |
| inter = int(np.logical_and(a, b).sum()) |
| den = int(a.sum() + b.sum()) |
| union = int(np.logical_or(a, b).sum()) |
| return (2 * inter / den if den else 1.0, inter / union if union else 1.0) |
|
|
|
|
| def bbox(mask): |
| y, x = np.where(mask) |
| return None if not len(x) else (x.min(), y.min(), x.max() + 1, y.max() + 1) |
|
|
|
|
| def bbox_iou(a, b): |
| if a is None or b is None: |
| return 0.0 |
| ix1, iy1, ix2, iy2 = max(a[0], b[0]), max(a[1], b[1]), min(a[2], b[2]), min(a[3], b[3]) |
| inter = max(0, ix2 - ix1) * max(0, iy2 - iy1) |
| aa, bb = (a[2] - a[0]) * (a[3] - a[1]), (b[2] - b[0]) * (b[3] - b[1]) |
| return inter / (aa + bb - inter) if aa + bb > inter else 0.0 |
|
|
|
|
| def auc(y, score): |
| y = np.asarray(y, dtype=bool) |
| score = np.asarray(score, dtype=float) |
| keep = np.isfinite(score) |
| y, score = y[keep], score[keep] |
| positives, negatives = int(y.sum()), int((~y).sum()) |
| if not positives or not negatives: |
| return np.nan |
| ranks = pd.Series(score).rank(method="average").to_numpy() |
| return float((ranks[y].sum() - positives * (positives + 1) / 2) / (positives * negatives)) |
|
|
|
|
| def save_stability_overlay(group, output_path): |
| valid = group[group["valid_mask"].astype(bool)] |
| if valid.empty: |
| return |
| image = Image.open(valid.iloc[0]["original_image_path"]).convert("RGB") |
| masks = [align(read(path), (image.height, image.width)) for path in valid["pred_mask_path"]] |
| agreement = np.mean(np.stack(masks), axis=0) |
| base = np.asarray(image).copy() |
| high, uncertain = agreement >= 0.8, (agreement > 0.2) & (agreement < 0.8) |
| base[high] = (0.55 * base[high] + np.array([0, 220, 0]) * 0.45).astype(np.uint8) |
| base[uncertain] = (0.55 * base[uncertain] + np.array([255, 170, 0]) * 0.45).astype(np.uint8) |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| Image.fromarray(base).save(output_path) |
|
|
|
|
| def main(): |
| args = parse_args() |
| out = Path(args.output_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| predictions = pd.read_csv(args.predictions_csv) |
| rows = [] |
| for sid, group in predictions.groupby("sample_id"): |
| valid = group[group["valid_mask"].astype(bool)] |
| masks = [(r["prompt"], read(r["pred_mask_path"])) for _, r in valid.iterrows()] |
| pair = [(overlap(a, b), bbox_iou(bbox(a), bbox(b))) for (_, a), (_, b) in itertools.combinations(masks, 2)] |
| areas = np.array([m.sum() for _, m in masks], dtype=float) |
| centers = np.array([[(b[0] + b[2]) / 2, (b[1] + b[3]) / 2] for _, m in masks if (b := bbox(m)) is not None]) |
| selected = valid[valid["prompt"] == args.default_prompt] |
| selected = selected.iloc[0] if len(selected) else valid.iloc[0] |
| rows.append({ |
| "sample_id": sid, |
| "selected_prompt": selected["prompt"], |
| "selected_mask_path": selected["pred_mask_path"], |
| "selected_mask_dice": selected["dice"], |
| "mean_pairwise_dice": np.mean([x[0][0] for x in pair]) if pair else np.nan, |
| "min_pairwise_dice": np.min([x[0][0] for x in pair]) if pair else np.nan, |
| "mean_pairwise_iou": np.mean([x[0][1] for x in pair]) if pair else np.nan, |
| "mask_area_mean": areas.mean() if len(areas) else np.nan, |
| "mask_area_std": areas.std() if len(areas) else np.nan, |
| "mask_area_cv": areas.std() / areas.mean() if len(areas) and areas.mean() else np.nan, |
| "mask_center_std": centers.std(axis=0).mean() if len(centers) else np.nan, |
| "mask_bbox_iou_mean": np.mean([x[1] for x in pair]) if pair else np.nan, |
| "mask_bbox_iou_min": np.min([x[1] for x in pair]) if pair else np.nan, |
| "valid_prompt_count": len(valid), |
| "failure_prompt_count": len(group) - len(valid), |
| }) |
| features = pd.DataFrame(rows) |
| features["good_mask"] = features["selected_mask_dice"] >= 0.75 |
| features["bad_mask"] = features["selected_mask_dice"] < 0.50 |
| features["quality_label"] = np.where(features["good_mask"], "good", np.where(features["bad_mask"], "bad", "medium")) |
| features.to_csv(out / "stability_features.csv", index=False) |
| feature_cols = ["mean_pairwise_dice", "min_pairwise_dice", "mean_pairwise_iou", "mask_area_cv", "mask_center_std", "mask_bbox_iou_mean", "mask_bbox_iou_min"] |
| summary = [] |
| for col in feature_cols: |
| sign = -1 if col in {"mask_area_cv", "mask_center_std"} else 1 |
| summary.append({"feature": col, "dice_correlation": features[[col, "selected_mask_dice"]].corr().iloc[0, 1], "auc_good": auc(features["good_mask"], sign * features[col]), "auc_bad": auc(features["bad_mask"], -sign * features[col])}) |
| pd.DataFrame(summary).to_csv(out / "qc_metrics_summary.csv", index=False) |
| auc_good = auc(features["good_mask"], features["mean_pairwise_dice"]) |
| auc_bad = auc(features["bad_mask"], -features["mean_pairwise_dice"]) |
| sweeps = [] |
| for threshold in [0.7, 0.8, 0.9]: |
| passed = (features["mean_pairwise_dice"] >= threshold) & (features["mask_area_cv"] <= 0.3) & (features["failure_prompt_count"] == 0) |
| p, f = features[passed], features[~passed] |
| sweeps.append({"threshold": threshold, "pass_rate": passed.mean(), "pass_subset_mean_dice": p["selected_mask_dice"].mean(), "fail_subset_mean_dice": f["selected_mask_dice"].mean(), "bad_mask_rate_in_pass": p["bad_mask"].mean(), "good_mask_rate_in_pass": p["good_mask"].mean(), "auc_good": auc_good, "auc_bad": auc_bad, "bad_detection_sensitivity": ((~passed) & features["bad_mask"]).sum() / max(1, features["bad_mask"].sum()), "bad_detection_specificity": (passed & ~features["bad_mask"]).sum() / max(1, (~features["bad_mask"]).sum())}) |
| pd.DataFrame(sweeps).to_csv(out / "qc_rule_threshold_sweep.csv", index=False) |
| features.sort_values("selected_mask_dice", ascending=False).head(30).to_csv(out / "pass_subset_examples.csv", index=False) |
| features.sort_values("selected_mask_dice").head(30).to_csv(out / "fail_subset_examples.csv", index=False) |
| features.sort_values("mean_pairwise_dice").head(30).to_csv(out / "worst_instability_examples.csv", index=False) |
| examples = pd.concat([features.nlargest(10, "mean_pairwise_dice"), features.nsmallest(10, "mean_pairwise_dice")]).drop_duplicates("sample_id") |
| for _, example in examples.iterrows(): |
| group = predictions[predictions["sample_id"] == example["sample_id"]] |
| save_stability_overlay(group, out / "overlay_examples" / f"{example['sample_id']}.png") |
| print(pd.DataFrame(sweeps).to_string(index=False)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
