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 select
from sqlalchemy.exc import DBAPIError, OperationalError

from ..celery_app import celery
from ..models import ImageRecord, MLSettings
from ._sync_engine import sync_session_factory as _sync_session_factory

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


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,
    autoretry_for=(OperationalError, DBAPIError, OSError),
    retry_backoff=5,
    retry_backoff_max=60,
    retry_jitter=True,
    max_retries=3,
    # Sized for the video branch: sample 10 frames, run tagger +
    # embedder on each (≈20 GPU ops vs 2 for an image). A loaded
    # ml-worker can take 5-10 min on a long video; bumped from
    # 5min/7min on 2026-05-28 after operator-flagged image 6288 (a
    # .mp4) hit the recovery sweep at 5 min while still legitimately
    # processing. Image runs return in seconds; the bump doesn't
    # affect their UX.
    soft_time_limit=900,   # 15 min
    time_limit=1200,       # 20 min hard
)
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


@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:
    """Retroactively apply allowlisted tags.

    Modes:
    - tag_id only      : scan all images for this tag.
    - image_id only    : scan all allowlisted tags for this image.
    - both             : just the (image, tag) pair.
    - neither          : full sweep (daily beat).

    Skips: already-applied, rejected (tag_suggestion_rejection), or
    confidence below the tag's allowlist min_confidence. Applied with
    source='ml_auto'.
    """
    from sqlalchemy import and_
    from sqlalchemy import select as sa_select
    from sqlalchemy.dialects.postgresql import insert as pg_insert

    from ..models import TagAllowlist, TagSuggestionRejection
    from ..models.tag import image_tag

    SessionLocal = _sync_session_factory()
    applied = 0
    with SessionLocal() as session:
        allow_rows = session.execute(
            sa_select(TagAllowlist.tag_id, TagAllowlist.min_confidence)
            if tag_id is None
            else sa_select(
                TagAllowlist.tag_id, TagAllowlist.min_confidence
            ).where(TagAllowlist.tag_id == tag_id)
        ).all()
        allow = {r[0]: r[1] for r in allow_rows}
        if not allow:
            return 0

        img_query = sa_select(
            ImageRecord.id, ImageRecord.tagger_predictions
        ).where(ImageRecord.tagger_predictions.is_not(None))
        if image_id is not None:
            img_query = img_query.where(ImageRecord.id == image_id)

        for img_id, preds in session.execute(img_query).all():
            preds = preds or {}
            for a_tag_id, min_conf in allow.items():
                exists = session.execute(
                    sa_select(image_tag.c.tag_id).where(
                        and_(
                            image_tag.c.image_record_id == img_id,
                            image_tag.c.tag_id == a_tag_id,
                        )
                    )
                ).scalar_one_or_none()
                if exists is not None:
                    continue
                rej = session.get(
                    TagSuggestionRejection, (img_id, a_tag_id)
                )
                if rej is not None:
                    continue
                from ..models import Tag

                tag = session.get(Tag, a_tag_id)
                if tag is None:
                    continue
                conf = _confidence_for_tag(session, tag, preds)
                if conf is None or conf < min_conf:
                    continue
                stmt = pg_insert(image_tag).values(
                    image_record_id=img_id,
                    tag_id=a_tag_id,
                    source="ml_auto",
                )
                stmt = stmt.on_conflict_do_nothing(
                    index_elements=["image_record_id", "tag_id"]
                )
                session.execute(stmt)
                applied += 1
        session.commit()
    return applied


def _confidence_for_tag(session, tag, preds: dict) -> float | None:
    """Highest confidence among predictions that resolve to `tag` —
    either the prediction name equals the tag name, or an alias maps
    (prediction name, category) -> tag.id.
    """
    from sqlalchemy import select as sa_select

    from ..models import TagAlias

    best: float | None = None
    direct = preds.get(tag.name)
    if direct is not None:
        best = float(direct.get("confidence", 0.0))
    alias_rows = session.execute(
        sa_select(TagAlias.alias_string, TagAlias.alias_category).where(
            TagAlias.canonical_tag_id == tag.id
        )
    ).all()
    for alias_string, alias_category in alias_rows:
        p = preds.get(alias_string)
        if p is None:
            continue
        if p.get("category") != alias_category:
            continue
        c = float(p.get("confidence", 0.0))
        if best is None or c > best:
            best = c
    return best


@celery.task(name="backend.app.tasks.ml.recompute_centroid", bind=True)
def recompute_centroid(self, tag_id: int) -> bool:
    import asyncio

    from ..extensions import get_session
    from ..services.ml.centroids import CentroidService

    async def _run() -> bool:
        async with get_session() as session:
            svc = CentroidService(session)
            result = await svc.recompute_for_tag(tag_id)
            await session.commit()
            return result

    return asyncio.run(_run())


@celery.task(name="backend.app.tasks.ml.recompute_centroids", bind=True)
def recompute_centroids(self) -> int:
    """Daily: find drifted centroids, enqueue recompute_centroid for each."""
    import asyncio

    from ..extensions import get_session
    from ..services.ml.centroids import CentroidService

    async def _list() -> list[int]:
        async with get_session() as session:
            return await CentroidService(session).list_drifted()

    drifted = asyncio.run(_list())
    for tid in drifted:
        recompute_centroid.delay(tid)
    return len(drifted)