feedmill-auto-label/backend/training.py

"""Fine-tune the project's base model on its master dataset (REQ-060…065).

The default is old + new together: the master dataset already accumulates every
merged batch, so a run sees the whole history. Training on the newest batch
alone is what makes a model quietly forget what it used to know, so it is not
what happens here.
"""

import json
import os
import shutil
import time
from typing import Optional

from backend import config, dataset, db, evaluate, hardware, jobs, projects

PRETRAINED = {"bbox": "yolo11n.pt", "polygon": "yolo11n-seg.pt"}


class TrainingError(Exception):
    pass


def models_dir(project_slug: str) -> str:
    return os.path.join(config.project_dir(project_slug), "models")


def start(project_id: int, epochs: int = 50, overrides: Optional[dict] = None, batch_ids: Optional[list] = None, class_ids: Optional[list] = None) -> dict:
    project = projects.get(project_id)
    if project is None:
        raise TrainingError("No such project")
    counts = dataset.summary(project_id)["splits"]
    if counts["train"] == 0:
        raise TrainingError(
            "The master dataset is empty — approve and merge a batch before training"
        )

    settings = hardware.resolve(overrides, epochs)
    job = jobs.create(
        "train",
        params={"project_id": project_id, "settings": settings, "batch_ids": batch_ids, "class_ids": class_ids},
        project_id=project_id,
        message=f"{counts['train']} train / {counts['val']} val",
    )
    return job.to_dict()


def listing(project_id: int) -> list:
    with db.cursor() as cur:
        cur.execute(
            "SELECT * FROM model_versions WHERE project_id = ? ORDER BY version DESC",
            (project_id,),
        )
        rows = []
        for row in cur.fetchall():
            item = dict(row)
            item["metrics"] = json.loads(item["metrics"] or "null")
            item["base_metrics"] = json.loads(item["base_metrics"] or "null")
            rows.append(item)
        return rows


def get_version(model_id: int) -> Optional[dict]:
    with db.cursor() as cur:
        cur.execute("SELECT * FROM model_versions WHERE id = ?", (model_id,))
        row = cur.fetchone()
        return dict(row) if row else None


def promote(model_id: int) -> dict:
    """Make a trained version the project's base model for the next round (REQ-064)."""
    version = get_version(model_id)
    if version is None:
        raise TrainingError("No such model version")
    project = projects.get(version["project_id"])
    base_path = os.path.join(config.project_dir(project["slug"]), "base", "model.pt")
    os.makedirs(os.path.dirname(base_path), exist_ok=True)
    shutil.copyfile(version["weights_path"], base_path)
    with db.cursor() as cur:
        cur.execute(
            "UPDATE projects SET base_model_path = ?, base_model_kind = 'trained' WHERE id = ?",
            (base_path, project["id"]),
        )
    return projects.get(project["id"])


def _next_version(cur, project_id: int) -> int:
    cur.execute(
        "SELECT COALESCE(MAX(version), 0) + 1 FROM model_versions WHERE project_id = ?",
        (project_id,),
    )
    return cur.fetchone()[0]


@jobs.handler("train")
def _run_train(job) -> None:
    from ultralytics import YOLO

    os.environ["ULTRALYTICS_OFFLINE"] = "true"
    os.environ["YOLO_OFFLINE"] = "true"

    project = projects.get(job.params["project_id"])
    settings = job.params["settings"]
    batch_ids = job.params.get("batch_ids")
    class_ids = job.params.get("class_ids")
    data_yaml = dataset.write_data_yaml(project, batch_ids=batch_ids, selected_class_ids=class_ids)

    # SAM3 and a training run must not hold VRAM at the same time (REQ-065).
    from backend.sam3_engine import release_engine

    if release_engine():
        job.log("Released SAM3 from VRAM before training")

    start_point = project["base_model_path"] or PRETRAINED[project["label_type"]]
    if not project["base_model_path"]:
        job.log(f"No base model on this project — starting from {start_point}")
    job.log(f"Fine-tuning {os.path.basename(start_point)} for {settings['epochs']} epoch(s) "
            f"(batch={settings['batch']}, imgsz={settings['imgsz']}, "
            f"device={settings['device']})")

    with db.cursor() as cur:
        version = _next_version(cur, project["id"])
    out_dir = os.path.join(models_dir(project["slug"]), str(version))
    os.makedirs(out_dir, exist_ok=True)

    model = YOLO(start_point)

    def on_epoch(trainer):
        # trainer.epoch is 0-based; report a human-facing 1-based count.
        epoch = getattr(trainer, 'epoch', 0) + 1
        total = getattr(trainer, 'epochs', settings["epochs"])
        job.progress(epoch, total, f"epoch {epoch}/{total}")

    model.add_callback("on_fit_epoch_end", on_epoch)
    job.progress(0, settings["epochs"])

    import torch
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True

    keep_run_dir = False
    try:
        model.train(
            data=data_yaml,
            epochs=settings["epochs"],
            imgsz=settings["imgsz"],
            batch=settings["batch"],
            device=settings["device"],
            workers=settings.get("workers", 8),
            cache="ram",
            project=os.path.join(out_dir, "runs"),
            name="train",
            exist_ok=True,
            amp=True,
            plots=False,
            verbose=False,
        )

        produced = os.path.join(out_dir, "runs", "train", "weights", "best.pt")
        if not os.path.isfile(produced):
            raise TrainingError("Training finished without producing best.pt")
        weights = os.path.join(out_dir, "best.pt")
        shutil.copyfile(produced, weights)

        job.log("Validating the base model and the new one on the same val set…")
        comparison = evaluate.compare(
            project["base_model_path"] or None, weights, data_yaml,
            [item["name"] for item in project["classes"]],
            imgsz=settings["imgsz"], device=settings["device"], batch=settings["batch"],
        )
        with open(os.path.join(out_dir, "metrics.json"), "w", encoding="utf-8") as handle:
            json.dump(comparison, handle, indent=2)

        with db.cursor() as cur:
            cur.execute(
                """INSERT INTO model_versions (project_id, version, weights_path,
                                               parent_model_path, metrics, base_metrics,
                                               created_at)
                   VALUES (?, ?, ?, ?, ?, ?, ?)""",
                (project["id"], version, weights, project["base_model_path"],
                 json.dumps(comparison["new"]), json.dumps(comparison["base"]), time.time()),
            )

        new = comparison["new"]
        if comparison["delta"]:
            delta = comparison["delta"]
            job.log(f"v{version}: mAP50 {new['map50']:.4f} ({delta['map50']:+.4f} vs base), "
                    f"mAP50-95 {new['map50_95']:.4f} ({delta['map50_95']:+.4f})")
        else:
            job.log(f"v{version}: mAP50 {new['map50']:.4f}, mAP50-95 {new['map50_95']:.4f} "
                    f"— {comparison['skipped']}")
    except Exception:
        # Keep the runs directory on failure: results.csv and logs are the only record
        # of why training failed (REQ-006, REQ-064).
        keep_run_dir = True
        raise
    finally:
        # Delete run directory on success or cancellation. job.cancelled finishes training
        # without raising an exception, so it takes this delete path.
        if not keep_run_dir:
            shutil.rmtree(os.path.join(out_dir, "runs"), ignore_errors=True)