FabledCurator/backend/app/tasks/ml.py

"""ML Celery tasks: per-image inference, backfill discovery, centroid
recompute, allowlist auto-apply, model self-heal.

All run on the ml-worker (queue 'ml') except recompute_centroids and
apply_allowlist_tags sweeps which are 'maintenance' lane. Sync sessions
(Celery workers are sync processes), same pattern as FC-2a tasks.
"""

from pathlib import Path

from sqlalchemy import create_engine, select
from sqlalchemy.orm import sessionmaker

from ..celery_app import celery
from ..config import get_config
from ..models import ImageRecord, MLSettings

IMAGES_ROOT = Path("/images")
VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".wmv", ".flv"}


def _sync_session_factory():
    cfg = get_config()
    engine = create_engine(cfg.database_url_sync, future=True, pool_pre_ping=True)
    return sessionmaker(engine, expire_on_commit=False)


def _is_video(path: Path) -> bool:
    return path.suffix.lower() in VIDEO_EXTS


@celery.task(name="backend.app.tasks.ml.tag_and_embed", bind=True)
def tag_and_embed(self, image_id: int) -> dict:
    """Run Camie + SigLIP on one image; store predictions + embedding;
    then enqueue per-image allowlist application.

    Video: sample frames between 10% and 90% of duration (VIDEO_ML_FRAMES,
    default 10). Max-pool tagger confidences across frames, mean-pool the
    SigLIP embeddings. On no-frames returns status='no_frames' (not an error).
    """
    import os

    from ..services.ml.embedder import get_embedder
    from ..services.ml.tagger import get_tagger

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        record = session.get(ImageRecord, image_id)
        if record is None:
            return {"status": "missing", "image_id": image_id}
        settings = session.execute(
            select(MLSettings).where(MLSettings.id == 1)
        ).scalar_one()

        src = Path(record.path)
        if not src.is_file():
            return {"status": "file_missing", "image_id": image_id}

        tagger = get_tagger()
        embedder = get_embedder()

        if _is_video(src):
            frames = _sample_video_frames(
                src, int(os.environ.get("VIDEO_ML_FRAMES", "10"))
            )
            if not frames:
                return {"status": "no_frames", "image_id": image_id}
            preds = _maxpool_predictions([tagger.infer(f) for f in frames])
            import numpy as np

            embedding = np.mean(
                [embedder.infer(f) for f in frames], axis=0
            ).astype("float32")
            for f in frames:
                f.unlink(missing_ok=True)
        else:
            raw = tagger.infer(src)
            preds = {
                name: {"category": p.category, "confidence": p.confidence}
                for name, p in raw.items()
            }
            embedding = embedder.infer(src)

        record.tagger_predictions = preds
        record.tagger_model_version = settings.tagger_model_version
        record.siglip_embedding = embedding.tolist()
        record.siglip_model_version = settings.embedder_model_version
        session.add(record)
        session.commit()

    apply_allowlist_tags.delay(image_id=image_id)
    return {"status": "ok", "image_id": image_id, "tags": len(preds)}


def _sample_video_frames(src: Path, n: int) -> list[Path]:
    """Extract n frames evenly between 10% and 90% of duration via ffmpeg.
    Returns temp file paths (caller deletes). Empty list on failure."""
    import json
    import subprocess
    import tempfile

    try:
        probe = subprocess.run(
            [
                "ffprobe", "-v", "quiet", "-print_format", "json",
                "-show_format", str(src),
            ],
            check=True, capture_output=True, timeout=30,
        )
        duration = float(json.loads(probe.stdout)["format"]["duration"])
    except Exception:
        return []
    if duration <= 0:
        return []

    start, end = duration * 0.10, duration * 0.90
    step = (end - start) / max(n - 1, 1)
    out: list[Path] = []
    tmpdir = Path(tempfile.mkdtemp(prefix="fc_vid_"))
    for i in range(n):
        ts = start + i * step
        dest = tmpdir / f"frame_{i:02d}.jpg"
        try:
            subprocess.run(
                [
                    "ffmpeg", "-ss", f"{ts:.2f}", "-i", str(src),
                    "-frames:v", "1", "-q:v", "3", "-y", str(dest),
                ],
                check=True, capture_output=True, timeout=60,
            )
            if dest.is_file():
                out.append(dest)
        except Exception:
            continue
    return out


def _maxpool_predictions(per_frame: list[dict]) -> dict:
    """Aggregate per-frame {name: TagPrediction} dicts by max confidence."""
    merged: dict[str, dict] = {}
    for frame_preds in per_frame:
        for name, p in frame_preds.items():
            cur = merged.get(name)
            if cur is None or p.confidence > cur["confidence"]:
                merged[name] = {
                    "category": p.category,
                    "confidence": p.confidence,
                }
    return merged


@celery.task(name="backend.app.tasks.ml.backfill", bind=True)
def backfill(self) -> int:
    """Enqueue tag_and_embed for images missing predictions/embeddings for
    the current model versions. Keyset pagination by id ASC (restart-safe).
    """
    SessionLocal = _sync_session_factory()
    enqueued = 0
    last_id = 0
    with SessionLocal() as session:
        settings = session.execute(
            select(MLSettings).where(MLSettings.id == 1)
        ).scalar_one()
        while True:
            rows = session.execute(
                select(ImageRecord.id)
                .where(ImageRecord.id > last_id)
                .where(
                    (ImageRecord.tagger_predictions.is_(None))
                    | (
                        ImageRecord.tagger_model_version
                        != settings.tagger_model_version
                    )
                    | (ImageRecord.siglip_embedding.is_(None))
                    | (
                        ImageRecord.siglip_model_version
                        != settings.embedder_model_version
                    )
                )
                .order_by(ImageRecord.id.asc())
                .limit(500)
            ).scalars().all()
            if not rows:
                break
            for image_id in rows:
                tag_and_embed.delay(image_id)
                enqueued += 1
            last_id = rows[-1]
    return enqueued


# --- Defined fully in Task 9/10. Stub so tag_and_embed's .delay() resolves
#     and the module imports cleanly between commits. Replaced in Task 9. ---
@celery.task(name="backend.app.tasks.ml.apply_allowlist_tags", bind=True)
def apply_allowlist_tags(self, tag_id: int | None = None,
                         image_id: int | None = None) -> int:
    return 0  # replaced in Task 9