scripts/train_separate.py

"""
Separate (one-model-per-class) MVTec training.

Saves metrics only — no checkpoints. After each class finishes, the model is deleted.
Compatible with the unified path: does not change train_unified behavior.
"""

import gc
import json
import os
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
import torch
from torch.amp import GradScaler

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))

from data.data_loader import build_dataloader
from scripts.evaluate import evaluate_category
from scripts.train_unified import build_model_and_optimizer, build_scheduler, run_one_epoch
from utils.checkpointing import append_epoch_metrics, ensure_run_dir, save_metrics_json
from utils.metrics import compute_image_auroc, compute_pixel_ap, compute_pixel_auroc


def _eval_one_category(model, cfg: dict, category: str, device: str) -> Dict[str, float]:
    root = cfg["data"]["mvtec_root"]
    size = tuple(cfg["data"]["mvtec_size"])
    dataset_name = cfg["train"].get("dataset_name", "mvtec")
    smooth_k = cfg["model"].get("anomaly_smooth_kernel", 7)
    smooth_s = cfg["model"].get("anomaly_smooth_sigma", 1.5)

    res = evaluate_category(
        model, root, category, size, device, dataset_name,
        smooth_kernel=smooth_k, smooth_sigma=smooth_s,
    )
    if len(np.unique(res["image_labels"])) < 2:
        return {"I-AUROC": 0.0, "P-AUROC": 0.0, "P-AP": 0.0}

    return {
        "I-AUROC": compute_image_auroc(res["image_scores"], res["image_labels"]),
        "P-AUROC": compute_pixel_auroc(res["pixel_scores"], res["pixel_labels"]),
        "P-AP": compute_pixel_ap(res["pixel_scores"], res["pixel_labels"]),
    }


def _free_model(*objs) -> None:
    for o in objs:
        del o
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def train_one_category(
    category: str,
    cfg: dict,
    device: str,
    run_dir: str,
) -> Dict[str, Any]:
    """
    Train a fresh model on one category. Metrics only (no checkpoints).
    Deletes the model before returning.
    """
    cat_dir = str(ensure_run_dir(os.path.join(run_dir, category)))
    max_epochs = int(cfg["train"].get("max_epochs", cfg["train"].get("epochs", 400)))
    patience = int(cfg["train"].get("early_stopping_patience", 10))
    monitor = cfg["train"].get("early_stopping_monitor", "I-AUROC")
    if monitor.startswith("avg_"):
        monitor = monitor[len("avg_"):]
    use_amp = bool(cfg["train"].get("use_amp", True))

    train_loader, train_dataset = build_dataloader(
        cfg, dataset_name="mvtec", shuffle=True, categories=[category]
    )
    if len(train_dataset) == 0:
        print(f"[WARN] {category}: no training images; skipping.")
        return {"category": category, "skipped": True}

    model, optimizer, criterion, device = build_model_and_optimizer(cfg, device=device)
    scheduler = build_scheduler(optimizer, cfg)

    amp_on = use_amp and device != "cpu" and torch.cuda.is_available()
    scaler = GradScaler("cuda", enabled=amp_on) if amp_on else GradScaler(enabled=False)

    history: Dict[str, Any] = {"train": [], "eval": []}
    best_metric = float("-inf")
    best_eval: Dict[str, float] = {}
    epochs_without_improvement = 0
    last_epoch = 0

    print(
        f"\n=== Separate train: {category} | samples={len(train_dataset)} | "
        f"batch={train_loader.batch_size} | max_epochs={max_epochs} | patience={patience} | "
        f"wd={cfg['train'].get('weight_decay', 1e-4)} | lr={optimizer.param_groups[0]['lr']} ===",
        flush=True,
    )

    for epoch in range(1, max_epochs + 1):
        last_epoch = epoch
        t0 = time.time()
        train_losses = run_one_epoch(
            model, train_loader, optimizer, criterion, device,
            train=True, use_amp=use_amp, scaler=scaler,
            desc=f"{category} {epoch}/{max_epochs}",
        )
        eval_m = _eval_one_category(model, cfg, category, device)
        elapsed = time.time() - t0
        monitor_val = float(eval_m.get(monitor, eval_m.get("I-AUROC", float("-inf"))))
        if scheduler is not None:
            scheduler.step(monitor_val)
        current_lr = optimizer.param_groups[0]["lr"]

        history["train"].append({"epoch": epoch, **train_losses})
        history["eval"].append({"epoch": epoch, **eval_m})

        epoch_record = {
            "epoch": epoch,
            "category": category,
            "train": train_losses,
            "eval": eval_m,
            "monitor": monitor,
            "monitor_value": monitor_val,
            "lr": current_lr,
            "elapsed_sec": round(elapsed, 2),
        }
        append_epoch_metrics(cat_dir, epoch_record)
        save_metrics_json(cat_dir, history)

        print(
            f"[{category} Epoch {epoch}/{max_epochs}] "
            f"train_total={train_losses['total']:.4f} | "
            f"I-AUROC={eval_m.get('I-AUROC', 0):.4f} | "
            f"P-AUROC={eval_m.get('P-AUROC', 0):.4f} | "
            f"P-AP={eval_m.get('P-AP', 0):.4f} | "
            f"{monitor}={monitor_val:.4f} | lr={current_lr:.2e} | {elapsed:.1f}s",
            flush=True,
        )

        if monitor_val > best_metric:
            best_metric = monitor_val
            best_eval = dict(eval_m)
            epochs_without_improvement = 0
            print(f"  -> new best {monitor}={best_metric:.4f}", flush=True)
        else:
            epochs_without_improvement += 1
            print(f"  -> no improvement ({epochs_without_improvement}/{patience})", flush=True)

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        if epochs_without_improvement >= patience:
            print(f"Early stopping {category} at epoch {epoch}.", flush=True)
            break

    result = {
        "category": category,
        "skipped": False,
        "last_epoch": last_epoch,
        "best_metric": best_metric,
        "best_eval": best_eval,
        "history": history,
    }
    with open(os.path.join(cat_dir, "metrics", "best_summary.json"), "w", encoding="utf-8") as f:
        json.dump(
            {k: v for k, v in result.items() if k != "history"},
            f, indent=2, ensure_ascii=False,
        )

    _free_model(model, optimizer, criterion, scaler, scheduler, train_loader, train_dataset)
    return result


def train_all_separate(
    cfg: dict,
    run_dir: str,
    device: str = None,
    categories: Optional[List[str]] = None,
    skip_done: bool = True,
    hf_repo_id: Optional[str] = None,
    hf_token: Optional[str] = None,
    hf_run_name: Optional[str] = None,
    hf_push_every_category: bool = True,
) -> Dict[str, Any]:
    """
    Train one model per MVTec category. Metrics-only; models are deleted after each class.
    """
    device = device or cfg["train"].get("device", "cuda")
    device = device if torch.cuda.is_available() else "cpu"
    run_dir = str(ensure_run_dir(run_dir))
    cats = categories or list(cfg["data"]["mvtec_categories"])

    summary_path = os.path.join(run_dir, "summary.json")
    summary: Dict[str, Any] = {"categories": {}, "average": {}}
    if os.path.isfile(summary_path):
        try:
            with open(summary_path, "r", encoding="utf-8") as f:
                summary = json.load(f)
        except Exception:
            pass

    print(
        f"Separate run dir: {run_dir} | classes={len(cats)} | "
        f"device={device} | amp={cfg['train'].get('use_amp', True)}",
        flush=True,
    )

    for cat in cats:
        done_file = os.path.join(run_dir, cat, "metrics", "best_summary.json")
        if skip_done and os.path.isfile(done_file):
            with open(done_file, "r", encoding="utf-8") as f:
                prev = json.load(f)
            summary["categories"][cat] = prev.get("best_eval", prev)
            print(f"=== Skip {cat} (already has metrics) ===", flush=True)
            continue

        result = train_one_category(cat, cfg, device, run_dir)
        if not result.get("skipped"):
            summary["categories"][cat] = result["best_eval"]

        # Update running average over finished classes
        finished = [v for v in summary["categories"].values() if isinstance(v, dict) and "I-AUROC" in v]
        if finished:
            summary["average"] = {
                k: float(np.mean([c[k] for c in finished]))
                for k in ["I-AUROC", "P-AUROC", "P-AP"]
            }
        with open(summary_path, "w", encoding="utf-8") as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)

        print(
            f"Running average over {len(finished)} classes: "
            f"I-AUROC={summary['average'].get('I-AUROC', 0):.4f}  "
            f"P-AUROC={summary['average'].get('P-AUROC', 0):.4f}  "
            f"P-AP={summary['average'].get('P-AP', 0):.4f}",
            flush=True,
        )

        if hf_repo_id and hf_push_every_category:
            try:
                from utils.hf_sync import push_folder_to_hub
                url = push_folder_to_hub(
                    run_dir, hf_repo_id,
                    path_in_repo=f"runs/{hf_run_name or Path(run_dir).name}",
                    token=hf_token,
                    commit_message=f"Separate metrics after {cat}",
                )
                print(f"  -> HF metrics push OK: {url}", flush=True)
            except Exception as e:
                print(f"  -> HF push failed: {e}", flush=True)

    print("\n=== Separate training finished ===")
    for cat, m in sorted(summary.get("categories", {}).items()):
        if isinstance(m, dict) and "I-AUROC" in m:
            print(
                f"  {cat:15s} I-AUROC={m['I-AUROC']:.3f}  "
                f"P-AUROC={m['P-AUROC']:.3f}  P-AP={m['P-AP']:.3f}"
            )
    if summary.get("average"):
        a = summary["average"]
        print(
            f"  {'AVERAGE':15s} I-AUROC={a.get('I-AUROC', 0):.3f}  "
            f"P-AUROC={a.get('P-AUROC', 0):.3f}  P-AP={a.get('P-AP', 0):.3f}"
        )

    return summary