From ac7e0d13bc18d6418e0c031f9202b1a4888ab55a Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Fri, 15 May 2026 07:40:27 -0400 Subject: [PATCH] feat(fc2b): add tag_and_embed + backfill Celery tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tag_and_embed: Camie + SigLIP on one image (video → 10-frame sample, max-pool tags, mean-pool embeddings), stores predictions/embedding with model versions, then enqueues per-image allowlist apply. backfill: keyset-paginated discovery of images missing predictions/embeddings for the current model versions (restart-safe). apply_allowlist_tags stub included so .delay() resolves between commits (filled in Task 9). Co-Authored-By: Claude Opus 4.7 (1M context) --- backend/app/celery_app.py | 1 + backend/app/tasks/ml.py | 197 ++++++++++++++++++++++++++++++++++++++ tests/test_tasks_ml.py | 54 +++++++++++ 3 files changed, 252 insertions(+) create mode 100644 backend/app/tasks/ml.py create mode 100644 tests/test_tasks_ml.py diff --git a/backend/app/celery_app.py b/backend/app/celery_app.py index 1f10df3..d20222b 100644 --- a/backend/app/celery_app.py +++ b/backend/app/celery_app.py @@ -28,6 +28,7 @@ def make_celery() -> Celery: "backend.app.tasks.import_file", "backend.app.tasks.thumbnail", "backend.app.tasks.maintenance", + "backend.app.tasks.ml", ], ) app.conf.update( diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py new file mode 100644 index 0000000..5305644 --- /dev/null +++ b/backend/app/tasks/ml.py @@ -0,0 +1,197 @@ +"""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 diff --git a/tests/test_tasks_ml.py b/tests/test_tasks_ml.py new file mode 100644 index 0000000..fe383df --- /dev/null +++ b/tests/test_tasks_ml.py @@ -0,0 +1,54 @@ +"""tag_and_embed / backfill task tests. Models aren't in CI, so we test +the pure helpers (_maxpool_predictions, _is_video) as unit tests, and the +DB-touching backfill query as an integration test with monkeypatched +inference. +""" + +from pathlib import Path + +import pytest + +from backend.app.services.ml.tagger import TagPrediction +from backend.app.tasks.ml import _is_video, _maxpool_predictions + + +def test_is_video(): + assert _is_video(Path("a.mp4")) is True + assert _is_video(Path("a.MKV")) is True + assert _is_video(Path("a.jpg")) is False + + +def test_maxpool_predictions(): + f1 = {"smile": TagPrediction("smile", "general", 0.6)} + f2 = { + "smile": TagPrediction("smile", "general", 0.9), + "sword": TagPrediction("sword", "general", 0.7), + } + merged = _maxpool_predictions([f1, f2]) + assert merged["smile"]["confidence"] == 0.9 + assert merged["sword"]["confidence"] == 0.7 + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_backfill_enqueues_missing(db, monkeypatch): + from backend.app.models import ImageRecord + from backend.app.tasks import ml as ml_tasks + + calls = [] + monkeypatch.setattr( + ml_tasks.tag_and_embed, "delay", lambda image_id: calls.append(image_id) + ) + + img = ImageRecord( + path="/images/n.jpg", sha256="n" * 64, size_bytes=1, + mime="image/jpeg", width=1, height=1, + origin="imported_filesystem", integrity_status="unknown", + tagger_predictions=None, siglip_embedding=None, + ) + db.add(img) + await db.commit() + + count = ml_tasks.backfill() + assert count >= 1 + assert img.id in calls