337 lines
12 KiB
Python
337 lines
12 KiB
Python
"""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
|
|
|
|
|
|
@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)
|