feat(fc2b): add tag_and_embed + backfill Celery tasks
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) <noreply@anthropic.com>
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user