scripts/train_sam3_decoder_segmentation.py

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