#!/usr/bin/env python3 import argparse import json import os import time from pathlib import Path import pandas as pd import torch from torch.utils.data import DataLoader from tqdm import tqdm from sam3_decoder_experiment_lib import ( PublicSegmentationDataset, SAM3FeatureModel, binary_metrics, bootstrap_ci, choose_visual_examples, dice_bce_loss, ensure_dir, save_overlay, save_prediction, seed_everything, split_dataset, summarize_metrics, write_csv, write_json, ) def parse_args(): p = argparse.ArgumentParser() p.add_argument("--dataset_csv", required=True) p.add_argument("--protocol", choices=["all_public_indomain", "leave_one_dataset_out", "source_to_target"], required=True) p.add_argument("--heldout_dataset", default=None) p.add_argument("--train_datasets", default=None, help="Comma-separated source datasets for source_to_target.") p.add_argument("--test_dataset", default=None, help="Target dataset for source_to_target.") p.add_argument("--output_dir", required=True) bundle_root = Path(__file__).resolve().parent.parent p.add_argument( "--sam3_checkpoint", default=os.environ.get("SAM3_CHECKPOINT", str(bundle_root / "model" / "sam3_base.pt")), ) p.add_argument("--encoder", default="sam3") p.add_argument("--encoder_trainable", choices=["frozen", "last_block", "lora"], default="frozen") p.add_argument("--decoder", choices=["sam3_native", "cnn", "unet", "segformer"], required=True) p.add_argument("--prompt_type", choices=["none", "semantic_text", "gt_bbox", "text"], default="none") p.add_argument("--prompt_text", default="breast tumor") p.add_argument("--prompt_column", default=None, help="Optional dataset column containing one semantic prompt per row.") p.add_argument("--lora_rank", type=int, default=8) p.add_argument("--lora_alpha", type=float, default=16) p.add_argument("--no_train_native_head", action="store_true", help="For encoder_trainable=lora, train only LoRA params and freeze native heads.") p.add_argument("--image_size", type=int, default=512) p.add_argument("--epochs", type=int, default=100) p.add_argument("--patience", type=int, default=15) p.add_argument("--batch_size", type=int, default=2) p.add_argument("--lr", type=float, default=1e-4) p.add_argument("--encoder_lr", type=float, default=1e-5) p.add_argument("--weight_decay", type=float, default=1e-4) p.add_argument("--num_workers", type=int, default=4) p.add_argument("--use_amp", action="store_true") p.add_argument("--seed", type=int, default=42) p.add_argument("--resume", action="store_true") p.add_argument("--debug", action="store_true") p.add_argument("--max_train_samples", type=int, default=None) p.add_argument("--max_val_samples", type=int, default=None) p.add_argument("--max_test_samples", type=int, default=None) p.add_argument("--bootstrap_resamples", type=int, default=2000) return p.parse_args() def make_loader(df, image_size, batch_size, num_workers, train): ds = PublicSegmentationDataset(df, image_size=image_size, augment=train) return DataLoader( ds, batch_size=batch_size, shuffle=train, num_workers=num_workers, pin_memory=torch.cuda.is_available(), drop_last=False, ) def make_optimizer(model, args): encoder_params = [] decoder_params = [] for name, p in model.named_parameters(): if not p.requires_grad: continue if "backbone" in name: encoder_params.append(p) else: decoder_params.append(p) groups = [] if decoder_params: groups.append({"params": decoder_params, "lr": args.lr}) if encoder_params: groups.append({"params": encoder_params, "lr": args.encoder_lr}) return torch.optim.AdamW(groups, weight_decay=args.weight_decay) def batch_prompt_texts(batch, args): if args.prompt_column: prompts = batch.get(args.prompt_column) if prompts is None: raise KeyError(f"Prompt column {args.prompt_column!r} is missing from dataset batch") return list(prompts) return None def train_one_epoch(model, loader, optimizer, scaler, device, use_amp, args): model.train() losses = [] for batch in tqdm(loader, desc="train", leave=False): images = batch["image"].to(device, non_blocking=True) masks = batch["mask"].to(device, non_blocking=True) optimizer.zero_grad(set_to_none=True) with torch.cuda.amp.autocast(enabled=use_amp and device.type == "cuda"): logits = model(images, masks, prompt_texts=batch_prompt_texts(batch, args)) loss = dice_bce_loss(logits, masks) if scaler is not None: scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() else: loss.backward() optimizer.step() losses.append(float(loss.detach().cpu())) return sum(losses) / max(1, len(losses)) @torch.no_grad() def run_eval(model, loader, device, threshold=0.5, save_dir=None, save_outputs=False, seed=42, args=None): model.eval() rows = [] pred_dir = Path(save_dir) / "masks" if save_dir and save_outputs else None vis_dir = Path(save_dir) / "overlays" if save_dir and save_outputs else None if pred_dir: ensure_dir(pred_dir) ensure_dir(vis_dir) cache_for_vis = [] for batch in tqdm(loader, desc="eval", leave=False): images = batch["image"].to(device, non_blocking=True) masks = batch["mask"].to(device, non_blocking=True) logits = model(images, masks, prompt_texts=batch_prompt_texts(batch, args) if args else None) probs = torch.sigmoid(logits).detach().cpu().numpy() gt = masks.detach().cpu().numpy() for i in range(probs.shape[0]): prob_i = probs[i, 0] gt_i = gt[i, 0] metrics = binary_metrics(prob_i, gt_i, threshold=threshold) row = { "dataset": batch["dataset"][i], "label": batch["label"][i], "label_id": int(batch["label_id"][i]), "case_id": batch["case_id"][i], "image_name": batch["image_name"][i], "mask_name": batch["mask_name"][i], "image_path": batch["image_path"][i], "threshold": float(threshold), **metrics, } rows.append(row) if save_outputs: safe_name = Path(row["image_name"]).stem pred_path = pred_dir / f"{safe_name}_pred.png" save_prediction(prob_i, pred_path) cache_for_vis.append((row, gt_i, prob_i >= threshold)) if save_outputs and cache_for_vis: by_name = {row["image_name"]: (row, gt_i, pred_i) for row, gt_i, pred_i in cache_for_vis} for group, examples in choose_visual_examples(rows, seed=seed).items(): group_dir = vis_dir / group ensure_dir(group_dir) for row in examples: cached = by_name.get(row["image_name"]) if cached is None: continue _, gt_i, pred_i = cached save_overlay(row["image_path"], gt_i, pred_i, row, group_dir / f"{Path(row['image_name']).stem}.png") return rows def select_threshold(model, loader, device, args=None): candidates = [0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7] best_t = 0.5 best_dice = -1.0 for t in candidates: rows = run_eval(model, loader, device, threshold=t, save_outputs=False, args=args) dice = pd.DataFrame(rows)["dice"].mean() if dice > best_dice: best_dice = float(dice) best_t = float(t) return best_t, best_dice def save_eval_outputs(rows, output_dir, prefix, bootstrap_resamples, seed): output_dir = Path(output_dir) write_csv(output_dir / f"{prefix}_per_image_metrics.csv", rows) overall = summarize_metrics(rows) overall.update(bootstrap_ci(rows, "dice", n_resamples=bootstrap_resamples, seed=seed)) overall.update(bootstrap_ci(rows, "iou", n_resamples=bootstrap_resamples, seed=seed + 1)) write_json(output_dir / f"{prefix}_overall_metrics.json", overall) by_dataset = summarize_metrics(rows, ["dataset"]) by_label = summarize_metrics(rows, ["label"]) by_dataset.to_csv(output_dir / f"{prefix}_metrics_by_dataset.csv", index=False) by_label.to_csv(output_dir / f"{prefix}_metrics_by_label.csv", index=False) write_json(output_dir / f"{prefix}_metrics_by_dataset.json", by_dataset.to_dict(orient="records")) write_json(output_dir / f"{prefix}_metrics_by_label.json", by_label.to_dict(orient="records")) return overall def main(): args = parse_args() if args.prompt_type == "text": args.prompt_type = "semantic_text" if args.prompt_type == "none": args.prompt_text = "" seed_everything(args.seed) output_dir = Path(args.output_dir) ensure_dir(output_dir) if (output_dir / "best_model.pt").exists() and (output_dir / "test_overall_metrics.json").exists() and not args.resume: print(f"Skip existing finished run: {output_dir}") return df = pd.read_csv(args.dataset_csv) train_df, val_df, test_df = split_dataset( df, args.protocol, heldout_dataset=args.heldout_dataset, train_datasets=args.train_datasets, test_dataset=args.test_dataset, seed=args.seed, max_train_samples=args.max_train_samples, max_val_samples=args.max_val_samples, max_test_samples=args.max_test_samples, ) config = vars(args).copy() config.update( { "train_datasets": sorted(train_df["dataset"].unique().tolist()), "val_datasets": sorted(val_df["dataset"].unique().tolist()), "test_datasets": sorted(test_df["dataset"].unique().tolist()), "n_train": int(len(train_df)), "n_val": int(len(val_df)), "n_test": int(len(test_df)), "loss": "BCEWithLogitsLoss + DiceLoss", "augmentation": "hflip, small rotation, brightness/contrast jitter", "checkpoint_path": args.sam3_checkpoint, "prompt_note": "Main decoder comparison should use prompt_type=none for all runs. Prompt-based sam3_native runs are reference baselines only.", "prompt_purpose": "semantic lesion localization, not OCR/text segmentation" if args.prompt_type != "none" else "no prompt", "lora_rank": args.lora_rank if args.encoder_trainable == "lora" else None, "lora_alpha": args.lora_alpha if args.encoder_trainable == "lora" else None, "train_native_head": (not args.no_train_native_head) if args.encoder_trainable == "lora" else None, } ) write_json(output_dir / "config.json", config) train_df.to_csv(output_dir / "train_split.csv", index=False) val_df.to_csv(output_dir / "val_split.csv", index=False) test_df.to_csv(output_dir / "test_split.csv", index=False) train_loader = make_loader(train_df, args.image_size, args.batch_size, args.num_workers, train=True) val_loader = make_loader(val_df, args.image_size, args.batch_size, args.num_workers, train=False) test_loader = make_loader(test_df, args.image_size, args.batch_size, args.num_workers, train=False) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = SAM3FeatureModel( args.sam3_checkpoint, image_size=args.image_size, encoder_trainable=args.encoder_trainable, decoder_name=args.decoder, prompt_type=args.prompt_type, prompt_text=args.prompt_text, lora_rank=args.lora_rank, lora_alpha=args.lora_alpha, train_native_head=not args.no_train_native_head, ).to(device) optimizer = make_optimizer(model, args) scaler = torch.cuda.amp.GradScaler(enabled=args.use_amp and device.type == "cuda") best_dice = -1.0 best_epoch = -1 wait = 0 log_rows = [] start_epoch = 1 if args.resume and (output_dir / "last_model.pt").exists(): ckpt = torch.load(output_dir / "last_model.pt", map_location=device) model.load_state_dict(ckpt["model"], strict=False) optimizer.load_state_dict(ckpt["optimizer"]) start_epoch = int(ckpt.get("epoch", 0)) + 1 best_dice = float(ckpt.get("best_dice", -1.0)) best_epoch = int(ckpt.get("best_epoch", -1)) for epoch in range(start_epoch, args.epochs + 1): tic = time.time() train_loss = train_one_epoch(model, train_loader, optimizer, scaler, device, args.use_amp, args) val_rows = run_eval(model, val_loader, device, threshold=0.5, args=args) val_dice = float(pd.DataFrame(val_rows)["dice"].mean()) row = {"epoch": epoch, "train_loss": train_loss, "val_mean_dice": val_dice, "seconds": time.time() - tic} log_rows.append(row) write_csv(output_dir / "train_log.csv", log_rows) torch.save( {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "epoch": epoch, "best_dice": best_dice, "best_epoch": best_epoch}, output_dir / "last_model.pt", ) print(json.dumps(row)) if val_dice > best_dice: best_dice = val_dice best_epoch = epoch wait = 0 torch.save({"model": model.state_dict(), "epoch": epoch, "best_dice": best_dice, "config": config}, output_dir / "best_model.pt") write_csv(output_dir / "val_metrics.csv", val_rows) save_eval_outputs(val_rows, output_dir, "val", args.bootstrap_resamples if not args.debug else 200, args.seed) else: wait += 1 if wait >= args.patience: break best = torch.load(output_dir / "best_model.pt", map_location=device) model.load_state_dict(best["model"], strict=False) best_t, best_t_dice = select_threshold(model, val_loader, device, args=args) threshold_info = {"threshold_0p5": 0.5, "val_selected_threshold": best_t, "val_selected_threshold_dice": best_t_dice} write_json(output_dir / "thresholds.json", threshold_info) test_dir_05 = output_dir / "test_predictions_threshold_0p5" test_rows = run_eval(model, test_loader, device, threshold=0.5, save_dir=test_dir_05, save_outputs=True, seed=args.seed, args=args) overall = save_eval_outputs(test_rows, output_dir, "test", args.bootstrap_resamples if not args.debug else 200, args.seed) test_dir_val = output_dir / "test_predictions_val_threshold" test_rows_val = run_eval(model, test_loader, device, threshold=best_t, save_dir=test_dir_val, save_outputs=True, seed=args.seed, args=args) save_eval_outputs(test_rows_val, output_dir, "test_val_threshold", args.bootstrap_resamples if not args.debug else 200, args.seed) print(json.dumps({"best_epoch": best_epoch, "best_val_dice": best_dice, "test_mean_dice": overall.get("mean_dice"), "test_mean_iou": overall.get("mean_iou")}, indent=2)) if __name__ == "__main__": main()