train_final.py

"""
Train the final deployable model on ALL 10 images (no holdout).

LOOCV proved F1=0.94. This trains the production model using every
labeled particle for maximum generalization to new unseen images.

Usage:
    python train_final.py --config config/config.yaml --device cuda:0
    python train_final.py --config config/config.yaml --device mps
"""

import argparse
import random
import time
from pathlib import Path

import numpy as np
import torch
import yaml
from torch.utils.data import DataLoader

from src.dataset import ImmunogoldDataset
from src.model import ImmunogoldCenterNet
from src.loss import total_loss
from src.preprocessing import discover_synapse_data, load_synapse


def set_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def train_epoch(model, loader, optimizer, device):
    model.train()
    loss_sum = 0
    n = 0
    for batch in loader:
        imgs = batch["image"].to(device)
        optimizer.zero_grad()
        hm_pred, off_pred = model(imgs)
        loss, hm_l, off_l = total_loss(
            hm_pred, batch["heatmap"].to(device),
            off_pred, batch["offsets"].to(device),
            batch["offset_mask"].to(device),
            conf_weights=batch["conf_map"].to(device),
        )
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        optimizer.step()
        loss_sum += loss.item()
        n += 1
    return loss_sum / n


def main():
    parser = argparse.ArgumentParser(description="Train final deployable model")
    parser.add_argument("--config", default="config/config.yaml")
    parser.add_argument("--device", default="auto")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    with open(args.config) as f:
        cfg = yaml.safe_load(f)

    set_seed(args.seed)

    if args.device == "auto":
        device = torch.device(
            "cuda" if torch.cuda.is_available()
            else "mps" if torch.backends.mps.is_available()
            else "cpu"
        )
    else:
        device = torch.device(args.device)
    print(f"Device: {device}")

    # Load ALL data — no holdout
    records = discover_synapse_data(cfg["data"]["root"], cfg["data"]["synapse_ids"])

    # Dataset uses ALL images for training (fold_id=None means no exclusion)
    dataset = ImmunogoldDataset(
        records=records,
        fold_id="__NONE__",  # no image excluded
        mode="train",
        patch_size=cfg["data"]["patch_size"],
        stride=cfg["data"]["stride"],
        hard_mining_fraction=cfg["training"]["hard_mining_fraction"],
        copy_paste_per_class=cfg["training"]["copy_paste_per_class"],
        sigmas=cfg["heatmap"]["sigmas"],
        samples_per_epoch=500,
        seed=args.seed,
    )

    loader = DataLoader(
        dataset, batch_size=cfg["training"]["batch_size"],
        shuffle=True, num_workers=4, drop_last=True,
        worker_init_fn=ImmunogoldDataset.worker_init_fn,
    )

    print(f"Training on ALL {len(dataset.images)} images, "
          f"{sum(len(a['6nm'])+len(a['12nm']) for a in dataset.annotations.values())} particles")

    # Model
    pretrained = cfg["model"]["pretrained_weights"]
    if not Path(pretrained).exists():
        pretrained = None
        print("Warning: CEM500K weights not found, using ImageNet")

    model = ImmunogoldCenterNet(
        pretrained_path=pretrained,
        bifpn_channels=cfg["model"]["bifpn_channels"],
        bifpn_rounds=cfg["model"]["bifpn_rounds"],
    ).to(device)
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

    out_dir = Path("checkpoints/final")
    out_dir.mkdir(parents=True, exist_ok=True)
    start = time.time()

    # Phase 1: Frozen encoder (40 epochs — slightly shorter since more data)
    print("\n=== Phase 1: Frozen encoder (40 epochs) ===")
    model.freeze_encoder()
    opt = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=1e-3, weight_decay=1e-4,
    )
    sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=15, T_mult=2)

    for ep in range(1, 41):
        loss = train_epoch(model, loader, opt, device)
        sched.step()
        if ep % 10 == 0:
            elapsed = time.time() - start
            print(f"  Epoch {ep:3d} | loss={loss:.4f} | {elapsed:.0f}s")

    torch.save({"model_state_dict": model.state_dict(), "epoch": 40},
               out_dir / "phase1.pth")

    # Phase 2: Unfreeze deep layers (40 epochs)
    print("\n=== Phase 2: Unfreeze layer3+4 (40 epochs) ===")
    model.unfreeze_deep_layers()
    opt = torch.optim.AdamW([
        {"params": model.layer3.parameters(), "lr": 1e-5},
        {"params": model.layer4.parameters(), "lr": 5e-5},
        {"params": model.bifpn.parameters(), "lr": 5e-4},
        {"params": model.upsample.parameters(), "lr": 5e-4},
        {"params": model.heatmap_head.parameters(), "lr": 5e-4},
        {"params": model.offset_head.parameters(), "lr": 5e-4},
    ], weight_decay=1e-4)

    for ep in range(41, 81):
        loss = train_epoch(model, loader, opt, device)
        if ep % 10 == 0:
            elapsed = time.time() - start
            print(f"  Epoch {ep:3d} | loss={loss:.4f} | {elapsed:.0f}s")

    torch.save({"model_state_dict": model.state_dict(), "epoch": 80},
               out_dir / "phase2.pth")

    # Phase 3: Full fine-tune (60 epochs)
    print("\n=== Phase 3: Full fine-tune (60 epochs) ===")
    model.unfreeze_all()
    opt = torch.optim.AdamW([
        {"params": model.stem.parameters(), "lr": 1e-6},
        {"params": model.layer1.parameters(), "lr": 5e-6},
        {"params": model.layer2.parameters(), "lr": 1e-5},
        {"params": model.layer3.parameters(), "lr": 5e-5},
        {"params": model.layer4.parameters(), "lr": 1e-4},
        {"params": model.bifpn.parameters(), "lr": 2e-4},
        {"params": model.upsample.parameters(), "lr": 2e-4},
        {"params": model.heatmap_head.parameters(), "lr": 2e-4},
        {"params": model.offset_head.parameters(), "lr": 2e-4},
    ], weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=60, eta_min=1e-7)

    for ep in range(81, 141):
        loss = train_epoch(model, loader, opt, device)
        sched.step()
        if ep % 10 == 0:
            elapsed = time.time() - start
            print(f"  Epoch {ep:3d} | loss={loss:.4f} | {elapsed:.0f}s")
            torch.save({
                "model_state_dict": model.state_dict(),
                "epoch": ep,
            }, out_dir / f"phase3_{ep}.pth")

    # Save final model
    torch.save({
        "model_state_dict": model.state_dict(),
        "epoch": 140,
        "config": cfg,
    }, out_dir / "final_model.pth")

    elapsed = time.time() - start
    print(f"\n=== Done: 140 epochs in {elapsed:.0f}s ({elapsed/60:.1f}min) ===")
    print(f"Final model: {out_dir / 'final_model.pth'}")
    print(f"\nTo detect particles in a new image:")
    print(f"  python predict.py --image path/to/new_image.tif --checkpoint {out_dir / 'final_model.pth'}")


if __name__ == "__main__":
    main()