reTraining/backend/projects.py

"""Projects: the unit that makes this system reusable (REQ-001…006).

A project owns a base model, a locked class list, a video archive root, and its
own accumulating master dataset. Everything it produces lives under one folder,
so a project can be copied or backed up whole.

Classes come from the base model whenever there is one — `model.names` is the
only thing that keeps the master dataset, the auto-annotation prompts, and the
fine-tune consistent with each other (REQ-003).
"""

import os
import re
import shutil
import time
from typing import Dict, List, Optional

from backend import config, db

LABEL_TYPES = ("bbox", "polygon")

# Starting points when a project has no base model of its own (REQ-004).
PRETRAINED = {"bbox": "yolo11n.pt", "polygon": "yolo11n-seg.pt"}


class ProjectError(Exception):
    """Something the user can fix: a bad name, a missing folder, a locked field."""


def slugify(name: str) -> str:
    slug = re.sub(r"[^a-z0-9]+", "-", name.strip().lower()).strip("-")
    return slug or "project"


def _unique_slug(cur, name: str) -> str:
    base = slugify(name)
    slug, suffix = base, 2
    while True:
        cur.execute("SELECT 1 FROM projects WHERE slug = ?", (slug,))
        if cur.fetchone() is None:
            return slug
        slug, suffix = f"{base}-{suffix}", suffix + 1


def read_model_classes(weights_path: str) -> List[str]:
    """Class names in a YOLO checkpoint, in class-id order."""
    from ultralytics import YOLO

    try:
        names = YOLO(weights_path).names
    except Exception as exc:
        raise ProjectError(f"Could not read classes from that model: {exc}")
    if isinstance(names, dict):
        return [names[key] for key in sorted(names)]
    return list(names)


def _project_paths(slug: str) -> dict:
    root = config.project_dir(slug)
    return {
        "root": root,
        "base": os.path.join(root, "base"),
        "dataset": os.path.join(root, "dataset"),
        "batches": os.path.join(root, "batches"),
        "models": os.path.join(root, "models"),
    }


def create(name: str, label_type: str, video_root: str, classes: Optional[List[dict]] = None,
           base_model_path: Optional[str] = None, val_every: int = 5) -> dict:
    """Create a project. `classes` is [{"name": ..., "prompt": ...}, …] and is
    ignored when a base model is given — that model's names win."""
    if not name.strip():
        raise ProjectError("A project name is required")
    if label_type not in LABEL_TYPES:
        raise ProjectError(f"label_type must be one of {LABEL_TYPES}")

    video_root = os.path.abspath(os.path.expanduser(video_root))
    if not os.path.isdir(video_root):
        raise ProjectError(f"Video archive folder not found: {video_root}")

    if base_model_path:
        names = read_model_classes(base_model_path)
        classes = [{"name": n, "prompt": n} for n in names]
    if not classes:
        raise ProjectError("Give a base model to read classes from, or list the classes")

    cleaned = []
    for index, item in enumerate(classes):
        class_name = str(item.get("name", "")).strip()
        if not class_name:
            raise ProjectError(f"Class {index} has no name")
        cleaned.append({"name": class_name,
                        "prompt": str(item.get("prompt") or class_name).strip()})
    if len({c["name"] for c in cleaned}) != len(cleaned):
        raise ProjectError("Class names must be unique")

    with db.cursor() as cur:
        slug = _unique_slug(cur, name)
        paths = _project_paths(slug)
        for path in paths.values():
            os.makedirs(path, exist_ok=True)

        stored_model = ""
        kind = "pretrained"
        if base_model_path:
            stored_model = os.path.join(paths["base"], "model.pt")
            shutil.copyfile(base_model_path, stored_model)
            kind = "uploaded"

        cur.execute(
            """INSERT INTO projects (slug, name, label_type, base_model_path,
                                     base_model_kind, video_root, val_every, created_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
            (slug, name.strip(), label_type, config.rel_data_path(stored_model), kind,
             video_root, max(0, val_every), time.time()),
        )
        project_id = cur.lastrowid
        _write_classes(cur, project_id, cleaned)

    return get(project_id)


def _write_classes(cur, project_id: int, classes: List[dict]) -> None:
    cur.execute("DELETE FROM project_classes WHERE project_id = ?", (project_id,))
    cur.executemany(
        "INSERT INTO project_classes (project_id, class_id, name, prompt, container) "
        "VALUES (?, ?, ?, ?, ?)",
        [(project_id, index, item["name"], item["prompt"], item.get("container", 0))
         for index, item in enumerate(classes)],
    )


def _row_to_dict(cur, row) -> dict:
    cur.execute(
        "SELECT class_id, name, prompt, container FROM project_classes "
        "WHERE project_id = ? ORDER BY class_id",
        (row["id"],),
    )
    classes = [dict(item) for item in cur.fetchall()]
    # How many shapes hang off each class — the number the user needs before
    # agreeing to delete one (REQ-007).
    cur.execute(
        """SELECT a.class_id, COUNT(*) FROM annotations a
           JOIN frames f ON f.id = a.frame_id
           JOIN batches b ON b.id = f.batch_id
           WHERE b.project_id = ? GROUP BY a.class_id""",
        (row["id"],),
    )
    usage = dict(cur.fetchall())
    for item in classes:
        item["annotation_count"] = usage.get(item["class_id"], 0)
    cur.execute("SELECT COUNT(*) FROM batches WHERE project_id = ?", (row["id"],))
    batch_count = cur.fetchone()[0]
    cur.execute(
        "SELECT split, COUNT(*) FROM dataset_items WHERE project_id = ? GROUP BY split",
        (row["id"],),
    )
    dataset = {"train": 0, "val": 0}
    for split, count in cur.fetchall():
        dataset[split] = count

    sec_classes = []
    if "secondary_model_classes" in row.keys() and row["secondary_model_classes"]:
        try:
            sec_classes = json.loads(row["secondary_model_classes"])
        except Exception:
            sec_classes = []

    paths = _project_paths(row["slug"])
    return {
        "id": row["id"],
        "slug": row["slug"],
        "name": row["name"],
        "label_type": row["label_type"],
        "base_model_path": config.resolve_data_path(row["base_model_path"]),
        "base_model_kind": row["base_model_kind"],
        "secondary_model_path": row["secondary_model_path"] if "secondary_model_path" in row.keys() else None,
        "secondary_model_name": row["secondary_model_name"] if "secondary_model_name" in row.keys() else None,
        "secondary_model_classes": sec_classes,
        "base_model_fallback": PRETRAINED[row["label_type"]],
        "video_root": row["video_root"],
        "video_root_linux": config.archive_host_paths()["linux"],
        "video_root_windows": config.archive_host_paths()["windows"],
        "val_every": row["val_every"],
        "created_at": row["created_at"],
        "classes": classes,
        "batch_count": batch_count,
        "dataset": dataset,
        # Once anything has been merged the label type is settled (REQ-002).
        "label_type_locked": (dataset["train"] + dataset["val"]) > 0,
        "paths": paths,
    }


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


def listing() -> List[dict]:
    with db.cursor() as cur:
        cur.execute("SELECT * FROM projects ORDER BY created_at DESC")
        return [_row_to_dict(cur, row) for row in cur.fetchall()]


def update(project_id: int, prompts: Optional[dict] = None, val_every: Optional[int] = None,
           video_root: Optional[str] = None,
           containers: Optional[Dict[int, bool]] = None) -> dict:
    """Edit the things that are safe to change: prompts, container flags,
    split ratio, archive root. Class names and label type are not among them."""
    if get(project_id) is None:
        raise ProjectError("No such project")

    with db.cursor() as cur:
        if val_every is not None:
            cur.execute("UPDATE projects SET val_every = ? WHERE id = ?",
                        (max(0, val_every), project_id))
        if video_root is not None:
            resolved = os.path.abspath(os.path.expanduser(video_root))
            if not os.path.isdir(resolved):
                raise ProjectError(f"Video archive folder not found: {resolved}")
            cur.execute("UPDATE projects SET video_root = ? WHERE id = ?",
                        (resolved, project_id))
        for class_id, prompt in (prompts or {}).items():
            cur.execute(
                "UPDATE project_classes SET prompt = ? WHERE project_id = ? AND class_id = ?",
                (str(prompt).strip(), project_id, int(class_id)),
            )
        for class_id, flag in (containers or {}).items():
            cur.execute(
                "UPDATE project_classes SET container = ? WHERE project_id = ? AND class_id = ?",
                (1 if flag else 0, project_id, int(class_id)),
            )
    return get(project_id)



def add_class(project_id: int, name: str, prompt: Optional[str] = None) -> dict:
    """Append a class to an existing project (REQ-008).

    Appending is the easy direction: the new class takes the next id, so no
    existing annotation or label file means anything different afterwards. Only
    `data.yaml` has to be rewritten, because it carries `nc`.
    """
    project = get(project_id)
    if project is None:
        raise ProjectError("No such project")
    clean = name.strip()
    if not clean:
        raise ProjectError("A class needs a name")
    if any(item["name"] == clean for item in project["classes"]):
        raise ProjectError(f'This project already has a class called "{clean}"')

    next_id = max((item["class_id"] for item in project["classes"]), default=-1) + 1
    with db.cursor() as cur:
        cur.execute(
            "INSERT INTO project_classes (project_id, class_id, name, prompt) "
            "VALUES (?, ?, ?, ?)",
            (project_id, next_id, clean, (prompt or clean).strip()),
        )

    # No data.yaml to refresh here any more: it is assembled per training run
    # from the datasets that run picks, so it always reflects the current classes.
    return get(project_id)


def delete_class(project_id: int, class_id: int) -> dict:
    """Remove a class and renumber the ones above it, everywhere (REQ-007).

    "Everywhere" is the whole point: the database rows, and the label files
    already written into the master dataset. A YOLO label is an integer index,
    so a class list and a set of label files that disagree do not fail loudly —
    they train a model on the wrong names.
    """
    project = get(project_id)
    if project is None:
        raise ProjectError("No such project")
    target = next((c for c in project["classes"] if c["class_id"] == class_id), None)
    if target is None:
        raise ProjectError(f"This project has no class {class_id}")
    if len(project["classes"]) == 1:
        raise ProjectError("A project needs at least one class")

    from backend import dataset

    with db.cursor() as cur:
        frames_of_project = """
            SELECT f.id FROM frames f
            JOIN batches b ON b.id = f.batch_id
            WHERE b.project_id = ?
        """
        cur.execute(
            f"DELETE FROM annotations WHERE class_id = ? AND frame_id IN ({frames_of_project})",
            (class_id, project_id),
        )
        removed = cur.rowcount
        cur.execute(
            f"""UPDATE annotations SET class_id = class_id - 1
                WHERE class_id > ? AND frame_id IN ({frames_of_project})""",
            (class_id, project_id),
        )
        cur.execute("DELETE FROM project_classes WHERE project_id = ? AND class_id = ?",
                    (project_id, class_id))
        cur.execute(
            "UPDATE project_classes SET class_id = class_id - 1 "
            "WHERE project_id = ? AND class_id > ?",
            (project_id, class_id),
        )

    report = dataset.drop_class_from_labels(project, class_id)
    updated = get(project_id)

    return {
        "project": updated,
        "removed": {
            "class": target["name"],
            "annotations": removed,
            **report,
        },
    }


def set_base_model(project_id: int, weights_path: str) -> dict:
    """Point the project at a new base model and re-read its classes (REQ-003).

    Refused once the master dataset exists and the new model's classes differ —
    a dataset labelled against one class list cannot be trained against another.
    """
    project = get(project_id)
    if project is None:
        raise ProjectError("No such project")

    names = read_model_classes(weights_path)
    existing = [item["name"] for item in project["classes"]]
    if project["dataset"]["train"] + project["dataset"]["val"] > 0 and set(names) != set(existing):
        raise ProjectError(
            "That model's classes differ from the ones this project's dataset was "
            f"labelled with ({existing} vs {names}). Create a new project for it."
        )

    paths = _project_paths(project["slug"])
    os.makedirs(paths["base"], exist_ok=True)
    stored = os.path.join(paths["base"], "model.pt")
    if os.path.abspath(weights_path) != os.path.abspath(stored):
        shutil.copyfile(weights_path, stored)

    with db.cursor() as cur:
        cur.execute(
            "UPDATE projects SET base_model_path = ?, base_model_kind = 'uploaded' WHERE id = ?",
            (config.rel_data_path(stored), project_id),
        )
        _write_classes(cur, project_id, _kept_classes(project, names))
    return get(project_id)


def set_secondary_model(project_id: int, weights_path: str, name: str = "") -> dict:
    """Point the project at a secondary model for auto-annotation."""
    project = get(project_id)
    if project is None:
        raise ProjectError("No such project")

    paths = _project_paths(project["slug"])
    os.makedirs(paths["base"], exist_ok=True)
    stored = os.path.join(paths["base"], "secondary_model.pt")
    if os.path.abspath(weights_path) != os.path.abspath(stored):
        shutil.copyfile(weights_path, stored)

    names = read_model_classes(weights_path)
    model_label = name.strip() or os.path.basename(weights_path)
    with db.cursor() as cur:
        cur.execute(
            "UPDATE projects SET secondary_model_path = ?, secondary_model_name = ?, secondary_model_classes = ? WHERE id = ?",
            (stored, model_label, json.dumps(names), project_id),
        )
    return get(project_id)


def _kept_classes(project: dict, names: List[str]) -> List[dict]:
    """Keep the stored per-class attributes — prompt (REQ-171) and container
    flag (REQ-184) — for a class that survives a base-model swap; new classes
    fall back to the class name as prompt and no flag."""
    known = {item["name"]: item for item in project["classes"]}
    return [{"name": name,
             "prompt": known[name]["prompt"] if name in known else name,
             "container": known[name].get("container", 0) if name in known else 0}
            for name in names]


def training_start_point(project: dict) -> str:
    """The weights a training run should start from (REQ-060, REQ-004)."""
    path = config.resolve_data_path(project["base_model_path"])
    if path and os.path.isfile(path):
        return path
    return PRETRAINED[project["label_type"]]


def delete(project_id: int) -> bool:
    project = get(project_id)
    if project is None:
        return False
    with db.cursor() as cur:
        cur.execute("DELETE FROM projects WHERE id = ?", (project_id,))
    shutil.rmtree(project["paths"]["root"], ignore_errors=True)
    return True


def ensure_seed_project() -> None:
    """Ensure at least one project exists on startup using legacy data if available."""
    with db.cursor() as cur:
        cur.execute("SELECT COUNT(*) FROM projects")
        if cur.fetchone()[0] > 0:
            return

    legacy_model = "/data/base_models/best.pt"
    if not os.path.exists(legacy_model):
        legacy_model = "/videos/model/best.pt"
    if not os.path.exists(legacy_model):
        legacy_model = os.path.join(config.VIDEO_ROOT, "model", "best.pt")

    if os.path.exists(legacy_model):
        try:
            create(
                name="Cargo & Sack Detection",
                label_type="bbox",
                video_root=config.VIDEO_ROOT,
                base_model_path=legacy_model,
            )
            print("[startup] Seeded default project 'Cargo & Sack Detection' from legacy model")
            return
        except Exception as exc:
            print(f"[startup] Failed to seed project from legacy model: {exc}")

    try:
        create(
            name="Default Detection Project",
            label_type="bbox",
            video_root=config.VIDEO_ROOT,
            classes=[{"name": "sack", "prompt": "sack"}, {"name": "truck", "prompt": "truck"}],
        )
        print("[startup] Seeded 'Default Detection Project'")
    except Exception as exc:
        print(f"[startup] Seed project creation skipped: {exc}")