Files
FabledCurator/backend/app/tasks/ml.py
T
bvandeusen 369e3de684
CI / lint (push) Successful in 2s
CI / frontend-build (push) Successful in 18s
CI / backend-lint-and-test (push) Successful in 32s
CI / integration (push) Successful in 3m19s
feat(ml): cadence-based video frame sampling + min-frame tag aggregation (#747)
Video tag noise root cause: frames were a FIXED count (6) max-pooled — a tag
firing on one frame survived at peak confidence, and a fixed count under-samples
long multi-scene videos so real scene-local tags looked like noise.

Redesign (operator-steered):
- Sample at a fixed CADENCE — one frame every `video_frame_interval_seconds`
  (default 4) across the 5–95% window — so a tag's frame-presence reflects real
  screen time independent of video length. Capped at `video_max_frames` (default
  64): a long video stretches the spacing instead of exploding into hundreds of
  inferences, bounding per-video cost on the single ml-worker (per-frame ffmpeg
  timeout also cut 60s→30s).
- Aggregate with `_aggregate_video_predictions`: keep a tag only if it appears in
  >= `video_min_tag_frames` sampled frames (≈ that many × interval seconds on
  screen — duration-independent noise rejection), with confidence = MEAN over the
  frames it appears in (not max). Clamps the threshold to the sample count so a
  1–2-frame short video still tags.
- All three knobs are DB-backed ml_settings (migration 0053), patchable via
  /api/ml/settings + sliders in the ML settings card — replaces the
  VIDEO_ML_FRAMES env var (product-not-project).

Tests: aggregation drops one-frame noise + means corroborated tags + clamps on
short videos; settings round-trip + min>max validation. Replaced the
_maxpool_predictions unit test.

NOTE: this is the QUALITY half of #747. The perf half — the ml-worker runs
CPU-only — is GPU enablement, tracked separately in #872.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-16 11:07:00 -04:00

529 lines
20 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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.
"""
import logging
from pathlib import Path
from celery.exceptions import SoftTimeLimitExceeded
from sqlalchemy import delete, select
from sqlalchemy.exc import DBAPIError, OperationalError
from ..celery_app import celery
from ..models import ImagePrediction, ImageRecord, MLSettings
from ._sync_engine import sync_session_factory as _sync_session_factory
log = logging.getLogger(__name__)
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 6 frames, run tagger +
# embedder on each (≈12 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 (#747): sample frames at a fixed cadence (ml_settings
video_frame_interval_seconds, capped at video_max_frames), keep a tag only if
it appears in >= video_min_tag_frames frames and average its confidence over
those frames (mean-pool, not max — kills one-frame noise); mean-pool the
SigLIP embeddings. On no-frames returns status='no_frames' (not an error).
"""
import time
from ..services.ml.embedder import get_embedder
from ..services.ml.tagger import get_tagger
# Phase + file context, so a timeout/crash names WHICH file and WHERE it
# died instead of a bare SoftTimeLimitExceeded() (operator-flagged 2026-06-08:
# the activity told them nothing about the file or why). `ctx` is enriched
# once the record is loaded; both feed the worker log AND the re-raised
# exception message (which becomes the activity's error_message).
started = time.monotonic()
phase = "open_session"
ctx = f"image_id={image_id}"
def _elapsed() -> float:
return time.monotonic() - started
try:
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)
is_vid = _is_video(src)
ctx = (
f"image_id={image_id} path={record.path} mime={record.mime} "
f"bytes={record.size_bytes} video={is_vid}"
)
log.info("tag_and_embed start: %s", ctx)
if not src.is_file():
log.warning("tag_and_embed file missing on disk: %s", ctx)
return {"status": "file_missing", "image_id": image_id}
phase = "load_models"
tagger = get_tagger()
embedder = get_embedder()
if is_vid:
# Layer-3 isolation: ffprobe (a separate process) validates
# the container before we burn ~20 GPU ops sampling frames
# from it. A corrupt video that would crash the frame
# decoder is rejected cleanly here instead of taking down
# the ml-worker. Operator-flagged 2026-05-28.
phase = "video_probe"
from ..utils import safe_probe
vprobe = safe_probe.probe_video(src)
if not vprobe.ok:
log.warning(
"tag_and_embed bad video (%s): %s", vprobe.reason, ctx
)
return {
"status": "bad_video", "image_id": image_id,
"reason": vprobe.reason,
}
phase = "video_sample_frames"
t0 = time.monotonic()
frames = _sample_video_frames(
src,
interval=settings.video_frame_interval_seconds,
max_frames=settings.video_max_frames,
)
log.info(
"tag_and_embed sampled %d frame(s) in %.1fs: %s",
len(frames), time.monotonic() - t0, ctx,
)
if not frames:
return {"status": "no_frames", "image_id": image_id}
phase = "video_infer"
import numpy as np
preds = _aggregate_video_predictions(
[tagger.infer(f, store_floor=settings.tagger_store_floor)
for f in frames],
min_frames=settings.video_min_tag_frames,
)
embedding = np.mean(
[embedder.infer(f) for f in frames], axis=0
).astype("float32")
log.info(
"tag_and_embed video aggregated %d tag(s) from %d frame(s) "
"(min_frames=%d): %s",
len(preds), len(frames), settings.video_min_tag_frames, ctx,
)
for f in frames:
f.unlink(missing_ok=True)
else:
phase = "tag"
t0 = time.monotonic()
raw = tagger.infer(src, store_floor=settings.tagger_store_floor)
log.info(
"tag_and_embed tagged in %.1fs (%d tags): %s",
time.monotonic() - t0, len(raw), ctx,
)
preds = {
name: {"category": p.category, "confidence": p.confidence}
for name, p in raw.items()
}
phase = "embed"
t0 = time.monotonic()
embedding = embedder.infer(src)
log.info(
"tag_and_embed embedded in %.1fs: %s",
time.monotonic() - t0, ctx,
)
phase = "persist"
record.tagger_model_version = settings.tagger_model_version
record.siglip_embedding = embedding.tolist()
record.siglip_model_version = settings.embedder_model_version
session.add(record)
# Write the normalized image_prediction rows (#768) — the sole home
# for predictions now (image_record.tagger_predictions was dropped in
# migration 0046). Delete-then-insert keeps a re-tag idempotent;
# tagger_store_floor was already applied in tagger.infer, so preds is
# the >=floor set.
session.execute(
delete(ImagePrediction).where(
ImagePrediction.image_record_id == image_id
)
)
session.add_all([
ImagePrediction(
image_record_id=image_id, raw_name=name,
category=p.get("category", "general"),
score=float(p.get("confidence", 0.0)),
)
for name, p in preds.items()
])
session.commit()
except SoftTimeLimitExceeded:
log.error(
"tag_and_embed TIMED OUT after %.0fs in phase=%s: %s",
_elapsed(), phase, ctx,
)
# Re-raise as SoftTimeLimitExceeded (preserves the 'timeout' status in
# the task_run signal) but WITH context, so the activity error_message
# names the file + phase instead of being empty.
raise SoftTimeLimitExceeded(
f"timed out in phase={phase} after {_elapsed():.0f}s ({ctx})"
) from None
except Exception:
# OSError/DBAPIError/OperationalError are autoretried — re-raise the
# ORIGINAL so the type is preserved; just make sure it's logged with
# context first.
log.exception(
"tag_and_embed FAILED in phase=%s after %.0fs: %s",
phase, _elapsed(), ctx,
)
raise
log.info(
"tag_and_embed ok in %.1fs (%d tags): %s", _elapsed(), len(preds), ctx
)
apply_allowlist_tags.delay(image_id=image_id)
return {"status": "ok", "image_id": image_id, "tags": len(preds)}
def _sample_video_frames(
src: Path, *, interval: float, max_frames: int,
) -> list[Path]:
"""Sample frames at a fixed CADENCE — one every `interval` seconds — so a
tag's frame-presence reflects real screen time regardless of video length
(#747). The count is capped at `max_frames`: a video longer than
interval×max_frames stretches the spacing instead of exploding the frame
count (keeps cost bounded so a long video can't hog the single ml-worker).
Frames are taken across the 5%95% window (skip intro/outro black/cards) via
per-frame fast-seek. Returns temp file paths (caller deletes); [] 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.05, duration * 0.95
span = max(end - start, 0.0)
# Cadence count, clamped to [1, max_frames]. int(span/interval)+1 ≈ one frame
# per `interval` seconds across the window; the cap stretches spacing on very
# long videos.
n = max(1, min(int(span / interval) + 1, max(1, max_frames)))
step = span / 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:04d}.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=30,
)
if dest.is_file():
out.append(dest)
except Exception:
continue
return out
def _aggregate_video_predictions(per_frame: list[dict], *, min_frames: int) -> dict:
"""Aggregate per-frame {name: TagPrediction} into one prediction set (#747).
A tag is kept only if it appears (≥ the tagger store floor, already applied)
in at least `min_frames` of the sampled frames — because sampling is at a
fixed cadence, that means it was on screen for roughly min_frames×interval
seconds, so a single-frame flicker / scene-transition artifact is dropped
while a genuine scene-local tag in a long video survives. Confidence is the
MEAN over the frames where the tag appears (not max — max re-inflated the
one-frame noise this whole change exists to remove).
`min_frames` is clamped to the number of frames actually sampled so a very
short video (12 frames) still tags instead of dropping everything.
"""
n = len(per_frame)
if n == 0:
return {}
threshold = max(1, min(min_frames, n))
agg: dict[str, dict] = {}
for frame_preds in per_frame:
for name, p in frame_preds.items():
cur = agg.get(name)
if cur is None:
agg[name] = {"category": p.category, "sum": p.confidence, "count": 1}
else:
cur["sum"] += p.confidence
cur["count"] += 1
return {
name: {"category": v["category"], "confidence": v["sum"] / v["count"]}
for name, v in agg.items()
if v["count"] >= threshold
}
@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_model_version.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,
# Audit 2026-06-02 — the full-sweep mode (neither tag_id nor image_id)
# is O(images × allowlist) and legitimately runs >5 min on large
# libraries. Cap matches the maintenance queue's recovery threshold.
soft_time_limit=1800, time_limit=2100,
)
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
# Images that have any predictions (#768: from image_prediction, not
# the old JSON column), optionally narrowed to one image.
img_ids_query = sa_select(ImagePrediction.image_record_id).distinct()
if image_id is not None:
img_ids_query = img_ids_query.where(
ImagePrediction.image_record_id == image_id
)
for (img_id,) in session.execute(img_ids_query).all():
preds = _load_predictions_sync(session, img_id)
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 _load_predictions_sync(session, image_id: int) -> dict:
"""Predictions for one image from image_prediction (#768), in the
{raw_name: {category, confidence}} shape _confidence_for_tag consumes —
keeps the allowlist resolution logic unchanged."""
from sqlalchemy import select as sa_select
rows = session.execute(
sa_select(
ImagePrediction.raw_name,
ImagePrediction.category,
ImagePrediction.score,
).where(ImagePrediction.image_record_id == image_id)
).all()
return {
r.raw_name: {"category": r.category, "confidence": r.score}
for r in rows
}
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,
# Audit 2026-06-02 — drifted-centroid rebuild over potentially
# hundreds of tags.
soft_time_limit=1800, time_limit=2100,
)
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)