Add tools

tools/create_amos_ct_mr_manifest.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python
+"""Create a SACFlow manifest for AMOS22 CT->MRI or MRI->CT.
+
+Expected AMOS root:
+  dataset.json
+  labeled_data_meta_0000_0599.csv
+  imagesTr/ labelsTr/ imagesVa/ labelsVa/ imagesTs/ labelsTs/
+
+This v2 script handles the common ambiguity in AMOS case numbering:
+  - zero_based:  amos_0000..amos_0499 = CT, amos_0500..amos_0599 = MRI
+  - one_based:   amos_0001..amos_0500 = CT, amos_0501..amos_0600 = MRI
+  - auto:        infers zero_based vs one_based from the observed min/max IDs
+
+It also tries metadata first when a modality column exists. Always verify printed counts.
+"""
+from __future__ import annotations
+import argparse, csv, json, re
+from pathlib import Path
+from typing import Optional
+
+NIFTI_EXTS = [".nii.gz", ".nii", ".mha", ".mhd"]
+
+
+def strip_ext(p: Path) -> str:
+    name = p.name
+    for ext in NIFTI_EXTS:
+        if name.endswith(ext):
+            return name[:-len(ext)]
+    return p.stem
+
+
+def case_num(s: str) -> Optional[int]:
+    m = re.search(r"(\d+)$", str(s))
+    return int(m.group(1)) if m else None
+
+
+def list_imgs(d: Path):
+    out = []
+    for ext in NIFTI_EXTS:
+        out += list(d.glob(f"*{ext}"))
+    return sorted(out)
+
+
+def load_modality_map(meta_csv: Optional[Path]):
+    if not meta_csv or not meta_csv.exists():
+        return {}, None, None
+    with open(meta_csv, newline="") as f:
+        rows = list(csv.DictReader(f))
+    if not rows:
+        return {}, None, None
+    cols = list(rows[0].keys())
+    mod_col = None
+    # Look for a column whose values contain explicit modality strings.
+    for c in cols:
+        allvals = {str(r.get(c, "")).strip().lower() for r in rows}
+        if any(v in {"ct", "mri", "mr"} for v in allvals):
+            mod_col = c
+            break
+    id_cols = [c for c in cols if any(k in c.lower() for k in ["amos", "case", "id", "name", "file", "image"])]
+    if not id_cols:
+        id_cols = list(cols)
+    mapping = {}
+    if mod_col:
+        for r in rows:
+            mod = str(r.get(mod_col, "")).strip().lower()
+            if mod == "mr":
+                mod = "mri"
+            if mod not in {"ct", "mri"}:
+                continue
+            for c in id_cols:
+                v = str(r.get(c, "")).strip()
+                if not v:
+                    continue
+                stem = strip_ext(Path(v))
+                keys = {v, Path(v).name, stem}
+                n = case_num(v)
+                if n is not None:
+                    keys.update({str(n), f"{n:04d}", f"amos_{n:04d}"})
+                for k in keys:
+                    mapping[str(k).lower()] = mod
+    return mapping, mod_col, id_cols
+
+
+def infer_id_scheme(root: Path) -> str:
+    ids = []
+    for split in ["Tr", "Va", "Ts"]:
+        d = root / f"images{split}"
+        for img in list_imgs(d):
+            n = case_num(strip_ext(img))
+            if n is not None:
+                ids.append(n)
+    if not ids:
+        return "zero_based"
+    mn, mx = min(ids), max(ids)
+    # AMOS labeled release is 600 cases. If IDs run 1..600, use one_based.
+    # If IDs run 0..599, use zero_based.
+    if mn >= 1 and mx >= 600:
+        return "one_based"
+    return "zero_based"
+
+
+def infer_modality(stem: str, mod_map: dict, id_scheme: str, ct_max_id: Optional[int], mri_min_id: Optional[int]):
+    keys = [stem.lower(), Path(stem).name.lower()]
+    n = case_num(stem)
+    if n is not None:
+        keys += [str(n), f"{n:04d}", f"amos_{n:04d}"]
+    for k in keys:
+        if k in mod_map:
+            return mod_map[k]
+    if n is None:
+        return None
+    if ct_max_id is not None:
+        return "ct" if n <= ct_max_id else "mri"
+    if mri_min_id is not None:
+        return "mri" if n >= mri_min_id else "ct"
+    if id_scheme == "one_based":
+        return "ct" if 1 <= n <= 500 else "mri"
+    # zero_based
+    return "ct" if 0 <= n <= 499 else "mri"
+
+
+def make_items(root: Path, split: str, modality: str, mod_map: dict, prefix: str, id_scheme: str, ct_max_id: Optional[int], mri_min_id: Optional[int]):
+    img_dir = root / f"images{split}"
+    lab_dir = root / f"labels{split}"
+    items = []
+    if not img_dir.exists():
+        return items
+    for img in list_imgs(img_dir):
+        stem = strip_ext(img)
+        mod = infer_modality(stem, mod_map, id_scheme, ct_max_id, mri_min_id)
+        if mod != modality:
+            continue
+        lab = None
+        if lab_dir.exists():
+            for ext in NIFTI_EXTS:
+                cand = lab_dir / f"{stem}{ext}"
+                if cand.exists():
+                    lab = cand
+                    break
+        item = {"id": f"{prefix}_{split}_{stem}", "case_id": stem, "modality": modality, "image": str(img.resolve())}
+        if lab is not None:
+            item["label"] = str(lab.resolve())
+        items.append(item)
+    return sorted(items, key=lambda x: x["id"])
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--amos-root", required=True)
+    ap.add_argument("--output", required=True)
+    ap.add_argument("--source", choices=["ct", "mri"], default="ct")
+    ap.add_argument("--target", choices=["ct", "mri"], default="mri")
+    ap.add_argument("--meta", default=None)
+    ap.add_argument("--name", default=None)
+    ap.add_argument("--id-scheme", choices=["auto", "zero_based", "one_based"], default="auto")
+    ap.add_argument("--ct-max-id", type=int, default=None, help="Override: IDs <= this are CT")
+    ap.add_argument("--mri-min-id", type=int, default=None, help="Override: IDs >= this are MRI")
+    args = ap.parse_args()
+    root = Path(args.amos_root).resolve()
+    meta = Path(args.meta).resolve() if args.meta else root / "labeled_data_meta_0000_0599.csv"
+    mod_map, mod_col, id_cols = load_modality_map(meta)
+    id_scheme = infer_id_scheme(root) if args.id_scheme == "auto" else args.id_scheme
+    print(f"Metadata: {meta if meta.exists() else 'not found'}")
+    print(f"Detected modality column: {mod_col}; candidate id columns: {id_cols}")
+    if not mod_map:
+        print(f"WARNING: no explicit modality map from metadata; using numeric ID scheme: {id_scheme}")
+        if args.ct_max_id is not None:
+            print(f"Override active: IDs <= {args.ct_max_id} treated as CT")
+        if args.mri_min_id is not None:
+            print(f"Override active: IDs >= {args.mri_min_id} treated as MRI")
+    manifest = {"name": args.name or f"amos_{args.source}2{args.target}"}
+    for split_name, code in [("train", "Tr"), ("val", "Va"), ("test", "Ts")]:
+        manifest[f"source_{split_name}"] = make_items(root, code, args.source, mod_map, "source", id_scheme, args.ct_max_id, args.mri_min_id)
+        manifest[f"target_{split_name}"] = make_items(root, code, args.target, mod_map, "target", id_scheme, args.ct_max_id, args.mri_min_id)
+    out = Path(args.output)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    with open(out, "w") as f:
+        json.dump(manifest, f, indent=2)
+    print(f"Wrote {out}")
+    total_source = 0
+    total_target = 0
+    for k, v in manifest.items():
+        if isinstance(v, list):
+            with_labels = sum(1 for item in v if "label" in item)
+            print(f"{k}: {len(v)} items, {with_labels} labels")
+            if k.startswith("source_"):
+                total_source += len(v)
+            if k.startswith("target_"):
+                total_target += len(v)
+    print(f"TOTAL source ({args.source}): {total_source}")
+    print(f"TOTAL target ({args.target}): {total_target}")
+    print("\nIf using AMOS22 labeled 600 cases, expect about 500 CT and 100 MRI. If you see 499/101, rerun with --id-scheme one_based or --ct-max-id 500.")
+
+
+if __name__ == "__main__":
+    main()

tools/create_dataset_json.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+from __future__ import annotations
+import argparse
+import json
+from pathlib import Path
+import random
+
+NIFTI_EXTS = [".nii.gz", ".nii", ".mha", ".mhd"]
+
+def list_images(folder):
+    folder = Path(folder)
+    files = []
+    for ext in NIFTI_EXTS:
+        files.extend(folder.glob(f"*{ext}"))
+    return sorted(files)
+
+def stem_nii(p: Path):
+    name = p.name
+    for ext in NIFTI_EXTS:
+        if name.endswith(ext):
+            return name[:-len(ext)]
+    return p.stem
+
+def match_labels(images, labels_dir):
+    if labels_dir is None:
+        return {stem_nii(p): None for p in images}
+    labels = list_images(labels_dir)
+    lab_map = {stem_nii(p): p for p in labels}
+    out = {}
+    for img in images:
+        sid = stem_nii(img)
+        if sid in lab_map:
+            out[sid] = lab_map[sid]
+        else:
+            # fuzzy: strip common prefixes
+            cand = None
+            for k, v in lab_map.items():
+                if sid in k or k in sid:
+                    cand = v; break
+            out[sid] = cand
+    return out
+
+def make_items(images_dir, labels_dir=None, domain=""):
+    imgs = list_images(images_dir)
+    labels = match_labels(imgs, labels_dir)
+    items = []
+    for p in imgs:
+        sid = stem_nii(p)
+        item = {"id": f"{domain}_{sid}" if domain else sid, "image": str(p)}
+        if labels.get(sid):
+            item["label"] = str(labels[sid])
+        items.append(item)
+    return items
+
+def split_items(items, val_fraction, test_fraction, seed):
+    items = list(items)
+    random.Random(seed).shuffle(items)
+    n = len(items)
+    n_test = int(round(n * test_fraction))
+    n_val = int(round(n * val_fraction))
+    test = items[:n_test]
+    val = items[n_test:n_test+n_val]
+    train = items[n_test+n_val:]
+    return train, val, test
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--name", required=True)
+    ap.add_argument("--source-images", required=True)
+    ap.add_argument("--source-labels", default=None)
+    ap.add_argument("--target-images", required=True)
+    ap.add_argument("--target-labels", default=None)
+    ap.add_argument("--output", required=True)
+    ap.add_argument("--val-fraction", type=float, default=0.2)
+    ap.add_argument("--test-fraction", type=float, default=0.2)
+    ap.add_argument("--seed", type=int, default=1337)
+    args = ap.parse_args()
+    src = make_items(args.source_images, args.source_labels, "source")
+    tgt = make_items(args.target_images, args.target_labels, "target")
+    src_train, src_val, src_test = split_items(src, args.val_fraction, args.test_fraction, args.seed)
+    tgt_train, tgt_val, tgt_test = split_items(tgt, args.val_fraction, args.test_fraction, args.seed)
+    man = {
+        "name": args.name,
+        "source_train": src_train,
+        "source_val": src_val,
+        "source_test": src_test,
+        "target_train": tgt_train,
+        "target_val": tgt_val,
+        "target_test": tgt_test,
+    }
+    out = Path(args.output)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    with open(out, "w") as f:
+        json.dump(man, f, indent=2)
+    print(f"Wrote {out}")
+    print({k: len(v) for k, v in man.items() if isinstance(v, list)})
+if __name__ == "__main__":
+    main()

tools/eval.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+from __future__ import annotations
+import argparse
+import json
+from pathlib import Path
+import torch
+from sacflow.utils.config import load_yaml
+from sacflow.utils.misc import seed_everything, ensure_dir
+from sacflow.utils.distributed import init_distributed, cleanup, is_main_process
+from sacflow.data.loader import build_loader
+from sacflow.models.unet3d import build_model
+from sacflow.engine.train_loop import evaluate
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", required=True)
+    ap.add_argument("--checkpoint", default=None)
+    ap.add_argument("--split", default=None)
+    args = ap.parse_args()
+    cfg = load_yaml(args.config)
+    if args.checkpoint:
+        cfg.setdefault("eval", {})["checkpoint"] = args.checkpoint
+    if args.split:
+        cfg.setdefault("eval", {})["split"] = args.split
+    seed_everything(int(cfg.get("seed", 1337)))
+    device = init_distributed(cfg.get("distributed", {}).get("backend", "nccl"))
+    model = build_model(cfg).to(device)
+    ckpt_path = cfg.get("eval", {}).get("checkpoint") or cfg.get("train", {}).get("source_checkpoint")
+    if ckpt_path is None:
+        ckpt_path = str(Path(cfg["output_dir"]) / "checkpoints" / "best.pt")
+    ckpt = torch.load(ckpt_path, map_location="cpu")
+    model.load_state_dict(ckpt.get("model", ckpt), strict=False)
+    split = cfg.get("eval", {}).get("split", "target_test")
+    loader = build_loader(cfg, split=split, training=False, require_label=True)
+    metrics = evaluate(model, loader, cfg, device)
+    if is_main_process():
+        print(json.dumps(metrics, indent=2))
+        out = ensure_dir(Path(cfg["output_dir"]) / "eval")
+        with open(out / f"metrics_{split}.json", "w") as f:
+            json.dump(metrics, f, indent=2)
+    cleanup()
+
+if __name__ == "__main__":
+    main()

tools/export_source_memory.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+from __future__ import annotations
+import argparse
+from pathlib import Path
+import copy
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+from sacflow.utils.config import load_yaml
+from sacflow.utils.misc import seed_everything, ensure_dir, move_to_device
+from sacflow.utils.distributed import init_distributed, cleanup, is_main_process
+from sacflow.data.loader import build_loader
+from sacflow.models.unet3d import build_model
+from sacflow.methods.task_space import centered_classifier_basis, project_task_and_residual
+from sacflow.methods.source_memory import hard_onehot, class_moments, save_source_memory
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", required=True)
+    ap.add_argument("--checkpoint", required=True)
+    ap.add_argument("--output", required=True)
+    ap.add_argument("--num-passes", type=int, default=3, help="Number of random-crop passes over source_train for memory estimation")
+    args = ap.parse_args()
+    cfg = load_yaml(args.config)
+    # Source memory should estimate source feature statistics, not augmentation noise.
+    cfg = copy.deepcopy(cfg)
+    cfg.setdefault("data", {}).setdefault("augmentation", {})
+    cfg["data"]["augmentation"] = {"random_flip": False, "random_intensity_shift": 0.0, "random_intensity_scale": 0.0}
+    seed_everything(int(cfg.get("seed", 1337)))
+    device = init_distributed(cfg.get("distributed", {}).get("backend", "nccl"))
+    model = build_model(cfg).to(device)
+    ckpt = torch.load(args.checkpoint, map_location="cpu")
+    model.load_state_dict(ckpt.get("model", ckpt), strict=False)
+    model.eval()
+    loader = build_loader(cfg, split="source_train", training=True, require_label=True)
+    W = model.final_classifier_weight().detach().to(device)
+    Q = centered_classifier_basis(W)
+    C = cfg["data"]["num_classes"]
+    mu_acc = None
+    var_acc = None
+    feat_mu_acc = None
+    feat_var_acc = None
+    count_acc = None
+    with torch.no_grad():
+        for pass_idx in range(max(1, args.num_passes)):
+            for batch in tqdm(loader, desc=f"source memory pass {pass_idx+1}/{max(1,args.num_passes)}", disable=not is_main_process()):
+                batch = move_to_device(batch, device)
+                logits, feats = model(batch["image"], return_features=True)
+                feat = feats["prelogit"]
+                task, residual = project_task_and_residual(feat, Q)
+                labels = batch["label"]
+                if labels.shape[-3:] != residual.shape[-3:]:
+                    labels = F.interpolate(labels[:,None].float(), size=residual.shape[-3:], mode="nearest")[:,0].long()
+                probs = hard_onehot(labels, C)
+                mu, std, counts = class_moments(residual, probs)
+                fmu, fstd, _ = class_moments(feat, probs)
+                var = std.pow(2)
+                fvar = fstd.pow(2)
+                if mu_acc is None:
+                    mu_acc = mu * counts[:,None]
+                    var_acc = (var + mu.pow(2)) * counts[:,None]
+                    feat_mu_acc = fmu * counts[:,None]
+                    feat_var_acc = (fvar + fmu.pow(2)) * counts[:,None]
+                    count_acc = counts
+                else:
+                    mu_acc += mu * counts[:,None]
+                    var_acc += (var + mu.pow(2)) * counts[:,None]
+                    feat_mu_acc += fmu * counts[:,None]
+                    feat_var_acc += (fvar + fmu.pow(2)) * counts[:,None]
+                    count_acc += counts
+    count = count_acc.clamp_min(1e-6)
+    mu = mu_acc / count[:,None]
+    second = var_acc / count[:,None]
+    var = (second - mu.pow(2)).clamp_min(1e-6)
+    fmu = feat_mu_acc / count[:,None]
+    fsecond = feat_var_acc / count[:,None]
+    fvar = (fsecond - fmu.pow(2)).clamp_min(1e-6)
+    mem = {
+        "classifier_weight": W.detach().cpu(),
+        "task_basis_Q": Q.detach().cpu(),
+        "residual_mu": mu.detach().cpu(),
+        "residual_std": torch.sqrt(var).detach().cpu(),
+        "feature_mu": fmu.detach().cpu(),
+        "feature_std": torch.sqrt(fvar).detach().cpu(),
+        "class_counts": count.detach().cpu(),
+        "num_classes": C,
+        "feature_dim": W.shape[1],
+        "note": "Compact source memory. No raw source images stored. Estimated from source_train random crops with augmentation disabled.",
+        "num_passes": args.num_passes,
+    }
+    if is_main_process():
+        save_source_memory(args.output, mem)
+        print(f"Saved source memory to {args.output}")
+    cleanup()
+
+if __name__ == "__main__":
+    main()

tools/inspect_dataset.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+from __future__ import annotations
+import argparse, json
+from pathlib import Path
+import nibabel as nib
+import numpy as np
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--manifest", required=True)
+    args = ap.parse_args()
+    man = json.load(open(args.manifest))
+    for split, items in man.items():
+        if not isinstance(items, list):
+            continue
+        print(split, len(items))
+        for it in items[:3]:
+            img = nib.load(it["image"])
+            print(" ", it.get("id"), img.shape, img.header.get_zooms()[:3], "label", bool(it.get("label")))
+            if it.get("label"):
+                lab = np.asanyarray(nib.load(it["label"]).dataobj)
+                print("    labels", np.unique(lab)[:20])
+if __name__ == "__main__":
+    main()

tools/make_sfda_manifest.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+"""Create an SFDA-safe manifest by dropping labels from specified target splits.
+
+This prevents accidental target-label leakage during adaptation. Validation/test labels are
+kept by default so metrics can be computed.
+"""
+from __future__ import annotations
+import argparse
+import json
+from pathlib import Path
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--input", required=True, help="Input manifest JSON")
+    ap.add_argument("--output", required=True, help="Output manifest JSON")
+    ap.add_argument("--drop-splits", nargs="+", default=["target_train"], help="Splits whose labels should be removed")
+    args = ap.parse_args()
+    with open(args.input, "r") as f:
+        man = json.load(f)
+    for split in args.drop_splits:
+        for item in man.get(split, []):
+            item.pop("label", None)
+    man.setdefault("notes", {})["sfda_safe"] = {
+        "dropped_label_splits": args.drop_splits,
+        "warning": "Target adaptation labels removed to avoid SFDA leakage."
+    }
+    out = Path(args.output)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    with open(out, "w") as f:
+        json.dump(man, f, indent=2)
+    print(f"Wrote {out}")
+    for k, v in man.items():
+        if isinstance(v, list):
+            print(f"{k}: {len(v)} items, {sum('label' in x for x in v)} labels")
+
+
+if __name__ == "__main__":
+    main()

tools/remap_labels_in_manifest.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python
+"""Remap labels referenced by a SACFlow manifest and write a new manifest.
+
+Useful for AMOS CT->MRI shared-label protocols. Example:
+  labels 14 and 15 -> background 0, num_classes=14.
+"""
+from __future__ import annotations
+import argparse
+import json
+from pathlib import Path
+import nibabel as nib
+import numpy as np
+from tqdm import tqdm
+
+
+def parse_map(map_items):
+    out = {}
+    for item in map_items or []:
+        if ":" not in item:
+            raise ValueError(f"Bad --map entry {item}; expected old:new")
+        a, b = item.split(":", 1)
+        out[int(a)] = int(b)
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--input", required=True)
+    ap.add_argument("--output", required=True)
+    ap.add_argument("--label-root", required=True, help="Directory for remapped labels")
+    ap.add_argument("--map", nargs="*", default=[], help="Label mappings old:new, e.g. 14:0 15:0")
+    ap.add_argument("--num-classes", type=int, default=None)
+    ap.add_argument("--protocol-name", default="remapped")
+    ap.add_argument("--overwrite", action="store_true")
+    args = ap.parse_args()
+    mapping = parse_map(args.map)
+    with open(args.input, "r") as f:
+        man = json.load(f)
+    label_root = Path(args.label_root)
+    label_root.mkdir(parents=True, exist_ok=True)
+    seen = {}
+
+    def remap_one(path_str: str) -> str:
+        path = Path(path_str)
+        if path_str in seen:
+            return seen[path_str]
+        out = label_root / f"{path.parent.name}_{path.name}"
+        if out.exists() and not args.overwrite:
+            seen[path_str] = str(out.resolve())
+            return seen[path_str]
+        img = nib.load(str(path))
+        arr = img.get_fdata().astype(np.int16)
+        for old, new in mapping.items():
+            arr[arr == old] = new
+        if args.num_classes is not None:
+            arr[arr >= args.num_classes] = 0
+            arr[arr < 0] = 0
+        out_img = nib.Nifti1Image(arr.astype(np.uint8), img.affine, img.header)
+        out_img.set_data_dtype(np.uint8)
+        nib.save(out_img, str(out))
+        seen[path_str] = str(out.resolve())
+        return seen[path_str]
+
+    count = 0
+    for split, items in man.items():
+        if not isinstance(items, list):
+            continue
+        for item in tqdm(items, desc=f"remap {split}"):
+            if item.get("label"):
+                item["label"] = remap_one(item["label"])
+                count += 1
+    man["label_protocol"] = {
+        "name": args.protocol_name,
+        "mapping": {str(k): v for k, v in mapping.items()},
+        "num_classes": args.num_classes,
+        "description": "Labels remapped using tools/remap_labels_in_manifest.py",
+    }
+    out_path = Path(args.output)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(out_path, "w") as f:
+        json.dump(man, f, indent=2)
+    print(f"Wrote {out_path}; remapped {count} label references into {label_root}")
+    for split, items in man.items():
+        if isinstance(items, list):
+            print(f"{split}: {len(items)} items, {sum('label' in x for x in items)} labels")
+
+
+if __name__ == "__main__":
+    main()

tools/train.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python
+from __future__ import annotations
+import argparse
+from pathlib import Path
+import yaml
+from sacflow.utils.config import load_yaml, save_yaml, set_by_path
+from sacflow.utils.misc import seed_everything, ensure_dir
+from sacflow.utils.distributed import init_distributed, cleanup, is_main_process, barrier
+from sacflow.utils.wandb_utils import init_wandb, wandb_finish
+from sacflow.engine.train_loop import run_training
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", required=True)
+    ap.add_argument("--name", default=None)
+    ap.add_argument("--resume", nargs="?", const="auto", default=None, help="Resume checkpoint path, or use --resume without a value for auto-detection.")
+    ap.add_argument("--opts", nargs="*", default=[], help="Override config values: key=value, e.g. train.epochs=220 optim.lr=1e-4")
+    args = ap.parse_args()
+    cfg = load_yaml(args.config)
+    for opt in args.opts:
+        if "=" not in opt:
+            raise ValueError(f"Invalid --opts entry {opt!r}; expected key=value")
+        key, value = opt.split("=", 1)
+        try:
+            parsed = yaml.safe_load(value)
+        except Exception:
+            parsed = value
+        set_by_path(cfg, key, parsed)
+    if args.resume is not None:
+        cfg.setdefault("train", {})["resume_checkpoint"] = args.resume
+    seed_everything(int(cfg.get("seed", 1337)))
+    device = init_distributed(cfg.get("distributed", {}).get("backend", "nccl"))
+    out_dir = ensure_dir(cfg["output_dir"])
+    if is_main_process():
+        save_yaml(cfg, out_dir / "resolved_config.yaml")
+    run = init_wandb(cfg, run_name=args.name or Path(cfg["output_dir"]).name)
+    try:
+        run_training(cfg, device, wandb_run=run)
+    finally:
+        barrier()
+        wandb_finish(run)
+        cleanup()
+
+if __name__ == "__main__":
+    main()