Delete train_model_v2.py

train_model_v2.py

@@ -1,552 +0,0 @@
-"""
-Train MyoSeg v2.
-
-This trains a stronger Attention Residual U-Net for two-channel input:
-channel 0 = MyHC/myotube signal, channel 1 = DAPI/nuclei signal.
-
-Expected data layout:
-    data/
-      images/
-      masks/
-        Nuclei_m/
-        Myotubes_m/
-
-Optional app corrections are also accepted from:
-    corrections/<sample_id>/image.png
-    corrections/<sample_id>/nuclei_mask.png
-    corrections/<sample_id>/myotube_mask.png
-"""
-
-from __future__ import annotations
-
-import argparse
-import csv
-import json
-import math
-import os
-import random
-from dataclasses import asdict, dataclass
-from pathlib import Path
-
-import numpy as np
-from PIL import Image
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.utils.data import DataLoader, Dataset
-
-try:
-    from huggingface_hub import HfApi
-except Exception:
-    HfApi = None
-
-
-ARCH = "attention_resunet_v2"
-IMG_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff"}
-MASK_EXTS = {".png", ".tif", ".tiff"}
-
-
-def seed_everything(seed: int) -> None:
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
-
-
-def pil_resample(kind: str):
-    try:
-        return Image.Resampling.NEAREST if kind == "nearest" else Image.Resampling.BILINEAR
-    except AttributeError:
-        return Image.NEAREST if kind == "nearest" else Image.BILINEAR
-
-
-def norm_key(path_or_name: str) -> str:
-    stem = Path(path_or_name).stem.lower()
-    keep = [c if c.isalnum() else "_" for c in stem]
-    return "_".join("".join(keep).split("_"))
-
-
-def find_matching_file(folder: Path, image_path: Path) -> Path | None:
-    for ext in MASK_EXTS:
-        p = folder / f"{image_path.stem}{ext}"
-        if p.exists():
-            return p
-
-    wanted = norm_key(image_path.name)
-    for p in folder.iterdir() if folder.exists() else []:
-        if p.suffix.lower() in MASK_EXTS and norm_key(p.name) == wanted:
-            return p
-    return None
-
-
-def mask_has_pixels(path: Path) -> bool:
-    arr = np.array(Image.open(path).convert("L"))
-    return bool((arr > 0).any())
-
-
-@dataclass
-class Sample:
-    sample_id: str
-    image: str
-    nuclei_mask: str
-    myotube_mask: str
-    source: str
-
-
-def collect_samples(data_root: Path, corrections_dir: Path | None, allow_empty_masks: bool) -> list[Sample]:
-    image_dir = data_root / "images"
-    nuc_dir = data_root / "masks" / "Nuclei_m"
-    myo_dir = data_root / "masks" / "Myotubes_m"
-    if not image_dir.exists():
-        raise FileNotFoundError(f"Missing image folder: {image_dir}")
-    if not nuc_dir.exists() or not myo_dir.exists():
-        raise FileNotFoundError(f"Missing mask folders under: {data_root / 'masks'}")
-
-    samples: list[Sample] = []
-    skipped: list[str] = []
-    for image_path in sorted(p for p in image_dir.iterdir() if p.suffix.lower() in IMG_EXTS):
-        nuc = find_matching_file(nuc_dir, image_path)
-        myo = find_matching_file(myo_dir, image_path)
-        if not nuc or not myo:
-            skipped.append(f"{image_path.name}: missing mask")
-            continue
-        if not allow_empty_masks and (not mask_has_pixels(nuc) or not mask_has_pixels(myo)):
-            skipped.append(f"{image_path.name}: empty mask")
-            continue
-        samples.append(Sample(norm_key(image_path.name), str(image_path), str(nuc), str(myo), "base"))
-
-    if corrections_dir and corrections_dir.exists():
-        for folder in sorted(p for p in corrections_dir.iterdir() if p.is_dir()):
-            image = folder / "image.png"
-            nuc = folder / "nuclei_mask.png"
-            myo = folder / "myotube_mask.png"
-            if not (image.exists() and nuc.exists() and myo.exists()):
-                continue
-            if not allow_empty_masks and (not mask_has_pixels(nuc) or not mask_has_pixels(myo)):
-                skipped.append(f"{folder.name}: empty correction mask")
-                continue
-            samples.append(Sample(norm_key(folder.name), str(image), str(nuc), str(myo), "correction"))
-
-    if skipped:
-        print(f"Skipped {len(skipped)} unusable sample(s). First few:")
-        for item in skipped[:12]:
-            print(f"  - {item}")
-
-    if len(samples) < 12:
-        raise RuntimeError(
-            f"Only {len(samples)} usable labelled samples found. "
-            "For v2, add more base labelled images before training."
-        )
-    return samples
-
-
-def make_splits(samples: list[Sample], out_dir: Path, seed: int, val_split: float, test_split: float) -> dict:
-    split_path = out_dir / "splits_v2.json"
-    if split_path.exists():
-        return json.loads(split_path.read_text())
-
-    ids = [s.sample_id for s in samples]
-    random.Random(seed).shuffle(ids)
-    n = len(ids)
-    n_test = max(1, int(round(n * test_split))) if n >= 20 and test_split > 0 else 0
-    n_val = max(2, int(round(n * val_split)))
-    n_train = n - n_val - n_test
-    if n_train < 8:
-        n_test = 0
-        n_val = max(2, min(n - 8, n_val))
-        n_train = n - n_val
-    if n_train < 4:
-        raise RuntimeError("Not enough data after split. Add more labelled samples.")
-
-    splits = {
-        "train": ids[:n_train],
-        "val": ids[n_train:n_train + n_val],
-        "test": ids[n_train + n_val:],
-        "seed": seed,
-    }
-    split_path.write_text(json.dumps(splits, indent=2))
-    return splits
-
-
-def load_image_pair(image_path: str, size: int) -> np.ndarray:
-    rgb = np.array(Image.open(image_path).convert("RGB"), dtype=np.uint8)
-    red = Image.fromarray(rgb[..., 0], "L").resize((size, size), pil_resample("bilinear"))
-    blue = Image.fromarray(rgb[..., 2], "L").resize((size, size), pil_resample("bilinear"))
-    x = np.stack([np.array(red), np.array(blue)], axis=0).astype(np.float32) / 255.0
-    return x
-
-
-def load_mask_pair(nuc_path: str, myo_path: str, size: int) -> np.ndarray:
-    nuc = Image.open(nuc_path).convert("L").resize((size, size), pil_resample("nearest"))
-    myo = Image.open(myo_path).convert("L").resize((size, size), pil_resample("nearest"))
-    y = np.stack([(np.array(nuc) > 0), (np.array(myo) > 0)], axis=0).astype(np.float32)
-    return y
-
-
-class MyoSegV2Dataset(Dataset):
-    def __init__(self, samples: list[Sample], selected_ids: set[str], size: int, augment: bool):
-        self.samples = [s for s in samples if s.sample_id in selected_ids]
-        self.size = size
-        self.augment = augment
-
-    def __len__(self) -> int:
-        return len(self.samples)
-
-    def __getitem__(self, index: int):
-        s = self.samples[index]
-        x = load_image_pair(s.image, self.size)
-        y = load_mask_pair(s.nuclei_mask, s.myotube_mask, self.size)
-
-        if self.augment:
-            x, y = augment_xy(x, y)
-
-        return torch.from_numpy(x.copy()), torch.from_numpy(y.copy()), s.sample_id
-
-
-def augment_xy(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-    if random.random() < 0.5:
-        x = x[:, :, ::-1]
-        y = y[:, :, ::-1]
-    if random.random() < 0.5:
-        x = x[:, ::-1, :]
-        y = y[:, ::-1, :]
-    k = random.randint(0, 3)
-    if k:
-        x = np.rot90(x, k, axes=(1, 2))
-        y = np.rot90(y, k, axes=(1, 2))
-
-    scale = np.random.uniform(0.85, 1.20, size=(2, 1, 1)).astype(np.float32)
-    shift = np.random.uniform(-0.05, 0.05, size=(2, 1, 1)).astype(np.float32)
-    gamma = float(np.random.uniform(0.85, 1.20))
-    noise = np.random.normal(0.0, 0.015, size=x.shape).astype(np.float32)
-    x = np.clip((x * scale + shift), 0.0, 1.0)
-    x = np.clip(x ** gamma + noise, 0.0, 1.0)
-    return x.astype(np.float32), y.astype(np.float32)
-
-
-def group_count(channels: int) -> int:
-    for g in (8, 4, 2, 1):
-        if channels % g == 0:
-            return g
-    return 1
-
-
-class ResidualBlock(nn.Module):
-    def __init__(self, in_ch: int, out_ch: int):
-        super().__init__()
-        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False)
-        self.norm1 = nn.GroupNorm(group_count(out_ch), out_ch)
-        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False)
-        self.norm2 = nn.GroupNorm(group_count(out_ch), out_ch)
-        self.skip = nn.Identity() if in_ch == out_ch else nn.Conv2d(in_ch, out_ch, 1, bias=False)
-
-    def forward(self, x):
-        identity = self.skip(x)
-        x = F.silu(self.norm1(self.conv1(x)))
-        x = self.norm2(self.conv2(x))
-        return F.silu(x + identity)
-
-
-class AttentionGate(nn.Module):
-    def __init__(self, gate_ch: int, skip_ch: int, inter_ch: int):
-        super().__init__()
-        self.gate = nn.Sequential(nn.Conv2d(gate_ch, inter_ch, 1, bias=True), nn.GroupNorm(group_count(inter_ch), inter_ch))
-        self.skip = nn.Sequential(nn.Conv2d(skip_ch, inter_ch, 1, bias=True), nn.GroupNorm(group_count(inter_ch), inter_ch))
-        self.psi = nn.Sequential(nn.Conv2d(inter_ch, 1, 1, bias=True), nn.Sigmoid())
-
-    def forward(self, gate, skip):
-        alpha = self.psi(F.silu(self.gate(gate) + self.skip(skip)))
-        return skip * alpha
-
-
-def center_pad_like(x, ref):
-    diff_y = ref.size(2) - x.size(2)
-    diff_x = ref.size(3) - x.size(3)
-    if diff_x or diff_y:
-        x = F.pad(x, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2])
-    return x
-
-
-class AttentionResUNetV2(nn.Module):
-    def __init__(self, in_ch: int = 2, out_ch: int = 2, base: int = 32):
-        super().__init__()
-        self.e1 = ResidualBlock(in_ch, base)
-        self.e2 = ResidualBlock(base, base * 2)
-        self.e3 = ResidualBlock(base * 2, base * 4)
-        self.e4 = ResidualBlock(base * 4, base * 8)
-        self.pool = nn.MaxPool2d(2)
-        self.bridge = ResidualBlock(base * 8, base * 16)
-
-        self.up4 = nn.ConvTranspose2d(base * 16, base * 8, 2, 2)
-        self.att4 = AttentionGate(base * 8, base * 8, base * 4)
-        self.dec4 = ResidualBlock(base * 16, base * 8)
-        self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, 2)
-        self.att3 = AttentionGate(base * 4, base * 4, base * 2)
-        self.dec3 = ResidualBlock(base * 8, base * 4)
-        self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, 2)
-        self.att2 = AttentionGate(base * 2, base * 2, base)
-        self.dec2 = ResidualBlock(base * 4, base * 2)
-        self.up1 = nn.ConvTranspose2d(base * 2, base, 2, 2)
-        self.att1 = AttentionGate(base, base, max(base // 2, 1))
-        self.dec1 = ResidualBlock(base * 2, base)
-        self.out = nn.Conv2d(base, out_ch, 1)
-
-    def forward(self, x):
-        e1 = self.e1(x)
-        e2 = self.e2(self.pool(e1))
-        e3 = self.e3(self.pool(e2))
-        e4 = self.e4(self.pool(e3))
-        b = self.bridge(self.pool(e4))
-
-        x = center_pad_like(self.up4(b), e4)
-        x = self.dec4(torch.cat([x, self.att4(x, e4)], dim=1))
-        x = center_pad_like(self.up3(x), e3)
-        x = self.dec3(torch.cat([x, self.att3(x, e3)], dim=1))
-        x = center_pad_like(self.up2(x), e2)
-        x = self.dec2(torch.cat([x, self.att2(x, e2)], dim=1))
-        x = center_pad_like(self.up1(x), e1)
-        x = self.dec1(torch.cat([x, self.att1(x, e1)], dim=1))
-        return self.out(x)
-
-
-class ComboLoss(nn.Module):
-    def __init__(self, dice_weight: float = 0.45, bce_weight: float = 0.45, focal_weight: float = 0.10):
-        super().__init__()
-        self.dice_weight = dice_weight
-        self.bce_weight = bce_weight
-        self.focal_weight = focal_weight
-
-    def forward(self, logits, target):
-        bce_map = F.binary_cross_entropy_with_logits(logits, target, reduction="none")
-        bce = bce_map.mean()
-
-        prob = torch.sigmoid(logits)
-        dims = (0, 2, 3)
-        inter = (prob * target).sum(dim=dims)
-        union = prob.sum(dim=dims) + target.sum(dim=dims)
-        dice = 1.0 - ((2.0 * inter + 1.0) / (union + 1.0)).mean()
-
-        pt = prob * target + (1.0 - prob) * (1.0 - target)
-        focal = ((1.0 - pt).pow(2.0) * bce_map).mean()
-        return self.bce_weight * bce + self.dice_weight * dice + self.focal_weight * focal
-
-
-@torch.no_grad()
-def dice_from_logits(logits, target, thresholds=(0.5, 0.5)):
-    prob = torch.sigmoid(logits)
-    pred = torch.zeros_like(prob)
-    pred[:, 0] = prob[:, 0] > thresholds[0]
-    pred[:, 1] = prob[:, 1] > thresholds[1]
-    inter = (pred * target).sum(dim=(0, 2, 3))
-    union = pred.sum(dim=(0, 2, 3)) + target.sum(dim=(0, 2, 3))
-    return ((2.0 * inter + 1.0) / (union + 1.0)).cpu().numpy()
-
-
-@torch.no_grad()
-def evaluate(model, loader, device, thresholds=(0.5, 0.5)) -> dict:
-    model.eval()
-    dices = []
-    losses = []
-    loss_fn = ComboLoss()
-    for x, y, _ in loader:
-        x = x.to(device)
-        y = y.to(device)
-        logits = model(x)
-        losses.append(float(loss_fn(logits, y).item()))
-        dices.append(dice_from_logits(logits, y, thresholds=thresholds))
-    arr = np.array(dices)
-    d_nuc = float(arr[:, 0].mean())
-    d_myo = float(arr[:, 1].mean())
-    return {
-        "loss": float(np.mean(losses)) if losses else math.nan,
-        "dice_nuclei": d_nuc,
-        "dice_myotube": d_myo,
-        "mean_dice": (d_nuc + d_myo) / 2.0,
-    }
-
-
-@torch.no_grad()
-def tune_thresholds(model, loader, device) -> tuple[tuple[float, float], dict]:
-    best = ((0.5, 0.5), {"mean_dice": -1.0})
-    grid = [round(v, 2) for v in np.linspace(0.25, 0.75, 11)]
-    for tn in grid:
-        for tm in grid:
-            metrics = evaluate(model, loader, device, thresholds=(tn, tm))
-            if metrics["mean_dice"] > best[1]["mean_dice"]:
-                best = ((tn, tm), metrics)
-    return best
-
-
-def write_metrics_csv(path: Path, rows: list[dict]) -> None:
-    if not rows:
-        return
-    fields = sorted({k for row in rows for k in row})
-    with path.open("w", newline="") as f:
-        writer = csv.DictWriter(f, fieldnames=fields)
-        writer.writeheader()
-        writer.writerows(rows)
-
-
-def push_to_hub(model_path: Path, repo_id: str, filename: str, token: str | None) -> None:
-    if HfApi is None:
-        raise RuntimeError("huggingface_hub is not installed.")
-    api = HfApi(token=token)
-    api.upload_file(
-        path_or_fileobj=str(model_path),
-        path_in_repo=filename,
-        repo_id=repo_id,
-        repo_type="model",
-        commit_message=f"Upload MyoSeg v2 model ({filename})",
-    )
-
-
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--data_root", default="data")
-    ap.add_argument("--corrections_dir", default="corrections")
-    ap.add_argument("--out_dir", default="runs_v2")
-    ap.add_argument("--epochs", type=int, default=100)
-    ap.add_argument("--batch_size", type=int, default=4)
-    ap.add_argument("--image_size", type=int, default=512)
-    ap.add_argument("--base", type=int, default=32)
-    ap.add_argument("--lr", type=float, default=3e-4)
-    ap.add_argument("--weight_decay", type=float, default=1e-4)
-    ap.add_argument("--val_split", type=float, default=0.15)
-    ap.add_argument("--test_split", type=float, default=0.15)
-    ap.add_argument("--patience", type=int, default=18)
-    ap.add_argument("--seed", type=int, default=42)
-    ap.add_argument("--allow_empty_masks", action="store_true")
-    ap.add_argument("--push_to_hub", action="store_true")
-    ap.add_argument("--hf_repo", default=os.environ.get("HF_REPO_ID", "skarugu/myotube-unet"))
-    ap.add_argument("--hf_filename", default=os.environ.get("HF_FILENAME", "model_v2.pt"))
-    ap.add_argument("--hf_token", default=os.environ.get("HF_TOKEN"))
-    args = ap.parse_args()
-
-    seed_everything(args.seed)
-    out_dir = Path(args.out_dir)
-    out_dir.mkdir(parents=True, exist_ok=True)
-
-    samples = collect_samples(
-        Path(args.data_root),
-        Path(args.corrections_dir) if args.corrections_dir else None,
-        allow_empty_masks=args.allow_empty_masks,
-    )
-    splits = make_splits(samples, out_dir, args.seed, args.val_split, args.test_split)
-    sample_by_id = {s.sample_id: s for s in samples}
-    for split_name in ("train", "val", "test"):
-        missing = [sid for sid in splits.get(split_name, []) if sid not in sample_by_id]
-        if missing:
-            raise RuntimeError(f"Split file references missing sample(s): {missing[:5]}")
-
-    train_ds = MyoSegV2Dataset(samples, set(splits["train"]), args.image_size, augment=True)
-    val_ds = MyoSegV2Dataset(samples, set(splits["val"]), args.image_size, augment=False)
-    test_ds = MyoSegV2Dataset(samples, set(splits.get("test", [])), args.image_size, augment=False)
-
-    train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=0)
-    val_loader = DataLoader(val_ds, batch_size=max(1, min(args.batch_size, 4)), shuffle=False, num_workers=0)
-    test_loader = DataLoader(test_ds, batch_size=max(1, min(args.batch_size, 4)), shuffle=False, num_workers=0) if len(test_ds) else None
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = AttentionResUNetV2(in_ch=2, out_ch=2, base=args.base).to(device)
-    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
-    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(args.epochs, 1), eta_min=args.lr * 0.05)
-    loss_fn = ComboLoss()
-
-    best_score = -1.0
-    best_epoch = 0
-    history: list[dict] = []
-    best_path = out_dir / "model_v2_best.pt"
-    final_path = out_dir / "model_v2.pt"
-
-    print(f"Training {ARCH} on {device}")
-    print(f"Samples: train={len(train_ds)} val={len(val_ds)} test={len(test_ds)}")
-
-    for epoch in range(1, args.epochs + 1):
-        model.train()
-        train_losses = []
-        for x, y, _ in train_loader:
-            x = x.to(device)
-            y = y.to(device)
-            opt.zero_grad(set_to_none=True)
-            loss = loss_fn(model(x), y)
-            loss.backward()
-            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
-            opt.step()
-            train_losses.append(float(loss.item()))
-        scheduler.step()
-
-        val_metrics = evaluate(model, val_loader, device)
-        row = {
-            "epoch": epoch,
-            "train_loss": float(np.mean(train_losses)),
-            **{f"val_{k}": v for k, v in val_metrics.items()},
-            "lr": scheduler.get_last_lr()[0],
-        }
-        history.append(row)
-        print(
-            f"epoch {epoch:03d} | loss={row['train_loss']:.4f} | "
-            f"val_nuc={val_metrics['dice_nuclei']:.4f} | "
-            f"val_myo={val_metrics['dice_myotube']:.4f} | "
-            f"mean={val_metrics['mean_dice']:.4f}"
-        )
-
-        if val_metrics["mean_dice"] > best_score:
-            best_score = val_metrics["mean_dice"]
-            best_epoch = epoch
-            checkpoint = {
-                "arch": ARCH,
-                "model": model.state_dict(),
-                "args": vars(args),
-                "metrics": {
-                    "best_epoch": best_epoch,
-                    "val": val_metrics,
-                    "n_train": len(train_ds),
-                    "n_val": len(val_ds),
-                    "n_test": len(test_ds),
-                },
-                "class_names": ["nuclei", "myotubes"],
-                "input_channels": ["myhc_red", "dapi_blue"],
-            }
-            torch.save(checkpoint, best_path)
-
-        if epoch - best_epoch >= args.patience:
-            print(f"Early stopping at epoch {epoch}; best epoch was {best_epoch}.")
-            break
-
-    write_metrics_csv(out_dir / "training_history_v2.csv", history)
-
-    ckpt = torch.load(best_path, map_location=device)
-    model.load_state_dict(ckpt["model"])
-    tuned_thresholds, tuned_val = tune_thresholds(model, val_loader, device)
-    final_metrics = {
-        "arch": ARCH,
-        "best_epoch": best_epoch,
-        "threshold_nuclei": tuned_thresholds[0],
-        "threshold_myotube": tuned_thresholds[1],
-        "val": tuned_val,
-        "n_train": len(train_ds),
-        "n_val": len(val_ds),
-        "n_test": len(test_ds),
-    }
-    if test_loader is not None:
-        final_metrics["test"] = evaluate(model, test_loader, device, thresholds=tuned_thresholds)
-
-    ckpt["metrics"] = final_metrics
-    ckpt["args"]["thr_nuc"] = tuned_thresholds[0]
-    ckpt["args"]["thr_myo"] = tuned_thresholds[1]
-    torch.save(ckpt, final_path)
-    (out_dir / "metrics_v2.json").write_text(json.dumps(final_metrics, indent=2))
-    (out_dir / "samples_v2.json").write_text(json.dumps([asdict(s) for s in samples], indent=2))
-
-    print(f"Saved best v2 checkpoint: {final_path}")
-    print(json.dumps(final_metrics, indent=2))
-
-    if args.push_to_hub:
-        push_to_hub(final_path, args.hf_repo, args.hf_filename, args.hf_token)
-        print(f"Pushed {final_path} to {args.hf_repo}/{args.hf_filename}")
-
-
-if __name__ == "__main__":
-    main()