bvandeusen
bc6b0a4a58
Videos now route through a 10-frame sampling branch (configurable via VIDEO_ML_FRAMES env) instead of the previous unsupported_format early return. WD14 predictions are aggregated by max-confidence per (name, category) across frames so sparse signals aren't diluted; SigLIP embeddings are mean-pooled for a representative shot. Also generates a fallback thumbnail when the record is missing one, and removes the video_filter from backfill so videos get enqueued. Celery soft/hard limits bumped to 240s/360s to accommodate 10x inference. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
304 lines
12 KiB
Python
304 lines
12 KiB
Python
"""ML inference Celery tasks (WD14 tagging, SigLIP embedding, centroid recompute)."""
|
|
from __future__ import annotations
|
|
import logging
|
|
import os
|
|
import time
|
|
|
|
import numpy as np
|
|
from sqlalchemy import and_, func
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
|
|
from app.celery_app import celery
|
|
from app import db
|
|
from app.models import (
|
|
ImageRecord,
|
|
Tag,
|
|
ImageTagPrediction,
|
|
ImageEmbedding,
|
|
TagReferenceEmbedding,
|
|
TagSuggestionConfig,
|
|
image_tags,
|
|
)
|
|
|
|
log = logging.getLogger('celery.tasks.ml')
|
|
|
|
# Minimum raw WD14 confidence to bother storing. Below this, rows are noise.
|
|
WD14_STORE_FLOOR = float(os.environ.get('WD14_STORE_FLOOR', '0.05'))
|
|
|
|
# Number of frames to sample from each video for ML inference.
|
|
VIDEO_ML_FRAMES = int(os.environ.get('VIDEO_ML_FRAMES', '10'))
|
|
|
|
|
|
def _config_value(key: str, default: str) -> str:
|
|
row = TagSuggestionConfig.query.filter_by(key=key).first()
|
|
return row.value if row else default
|
|
|
|
|
|
def _run_video_inference(image, wd14, siglip) -> tuple[list[dict], np.ndarray] | None:
|
|
"""Sample VIDEO_ML_FRAMES frames from the video and aggregate predictions.
|
|
|
|
WD14 predictions are aggregated by taking the max confidence per (name, category)
|
|
across frames — a character appearing clearly in one frame shouldn't be diluted by
|
|
frames where it's absent. SigLIP embeddings are mean-pooled, which preserves cosine
|
|
distance behavior since the pgvector index normalizes as needed.
|
|
|
|
Returns (predictions, embedding) or None if frame extraction fails.
|
|
"""
|
|
import shutil
|
|
from app.utils.image_importer import extract_video_frames
|
|
|
|
frame_paths = extract_video_frames(image.filepath, count=VIDEO_ML_FRAMES)
|
|
if not frame_paths:
|
|
return None
|
|
|
|
tmpdir = os.path.dirname(frame_paths[0])
|
|
try:
|
|
best: dict[tuple[str, str], dict] = {}
|
|
embeddings: list[np.ndarray] = []
|
|
for fp in frame_paths:
|
|
raw = wd14.infer_filtered(fp, min_any=WD14_STORE_FLOOR)
|
|
for pred in raw:
|
|
key = (pred['name'], pred['category'])
|
|
prev = best.get(key)
|
|
if prev is None or pred['confidence'] > prev['confidence']:
|
|
best[key] = pred
|
|
embeddings.append(siglip.infer(fp))
|
|
|
|
if not embeddings:
|
|
return None
|
|
mean_emb = np.stack(embeddings).mean(axis=0).astype(np.float32)
|
|
return list(best.values()), mean_emb
|
|
finally:
|
|
shutil.rmtree(tmpdir, ignore_errors=True)
|
|
|
|
|
|
def _ensure_video_thumb(image) -> None:
|
|
"""Generate a thumbnail if the video record is missing one, and persist the path."""
|
|
from app.utils.image_importer import generate_video_thumbnail_mirrored
|
|
|
|
if image.thumb_path and os.path.exists(image.thumb_path):
|
|
return
|
|
new_thumb = generate_video_thumbnail_mirrored(image.filepath)
|
|
if new_thumb:
|
|
image.thumb_path = new_thumb
|
|
|
|
|
|
@celery.task(bind=True, name='app.tasks.ml.tag_and_embed',
|
|
max_retries=2, default_retry_delay=60,
|
|
soft_time_limit=240, time_limit=360)
|
|
def tag_and_embed(self, image_id: int):
|
|
"""Run WD14 + SigLIP on one image (or sampled video frames) and persist results."""
|
|
from app.ml import wd14, siglip # lazy import so web process never loads torch
|
|
from app.utils.image_importer import VIDEO_EXTS
|
|
|
|
image = ImageRecord.query.get(image_id)
|
|
if image is None:
|
|
log.warning(f"tag_and_embed: image {image_id} not found")
|
|
return {'status': 'missing'}
|
|
|
|
if not os.path.exists(image.filepath):
|
|
log.warning(f"tag_and_embed: file missing at {image.filepath}")
|
|
return {'status': 'file_missing'}
|
|
|
|
is_video = image.filepath.lower().endswith(VIDEO_EXTS)
|
|
|
|
try:
|
|
t0 = time.time()
|
|
if is_video:
|
|
result = _run_video_inference(image, wd14, siglip)
|
|
if result is None:
|
|
log.warning(f"tag_and_embed: video {image_id} produced no frames")
|
|
return {'status': 'no_frames'}
|
|
raw, embedding = result
|
|
_ensure_video_thumb(image)
|
|
else:
|
|
raw = wd14.infer_filtered(image.filepath, min_any=WD14_STORE_FLOOR)
|
|
embedding = siglip.infer(image.filepath)
|
|
t_inf = time.time() - t0
|
|
|
|
# Remove any pre-existing predictions/embedding for this image+model_version pair
|
|
db.session.query(ImageTagPrediction).filter(
|
|
ImageTagPrediction.image_id == image_id,
|
|
ImageTagPrediction.model_version == wd14.MODEL_VERSION,
|
|
).delete(synchronize_session=False)
|
|
db.session.query(ImageEmbedding).filter(
|
|
ImageEmbedding.image_id == image_id,
|
|
ImageEmbedding.model_version == siglip.MODEL_VERSION,
|
|
).delete(synchronize_session=False)
|
|
|
|
for pred in raw:
|
|
db.session.add(ImageTagPrediction(
|
|
image_id=image_id,
|
|
tag_name=pred['name'],
|
|
tag_category=pred['category'],
|
|
confidence=pred['confidence'],
|
|
model_version=wd14.MODEL_VERSION,
|
|
))
|
|
|
|
db.session.add(ImageEmbedding(
|
|
image_id=image_id,
|
|
model_version=siglip.MODEL_VERSION,
|
|
embedding=embedding.tolist(),
|
|
))
|
|
|
|
db.session.commit()
|
|
kind = 'video' if is_video else 'image'
|
|
log.info(f"tag_and_embed: {kind} {image_id} done in {t_inf:.2f}s ({len(raw)} predictions)")
|
|
return {'status': 'ok', 'predictions': len(raw), 'duration_s': t_inf, 'kind': kind}
|
|
except Exception as e:
|
|
db.session.rollback()
|
|
log.error(f"tag_and_embed failed for image {image_id}: {e}", exc_info=True)
|
|
raise self.retry(exc=e)
|
|
|
|
|
|
@celery.task(bind=True, name='app.tasks.ml.backfill',
|
|
soft_time_limit=None, time_limit=None)
|
|
def backfill(self, batch_size: int = 50, pause_seconds: float = 0.5):
|
|
"""Enqueue tag_and_embed for every image missing predictions or embeddings for the current model versions.
|
|
|
|
Uses keyset pagination on image_record.id so the loop moves forward monotonically.
|
|
Re-running after some tag_and_embed tasks have completed is safe — those images
|
|
simply drop out of the filter on the next run.
|
|
|
|
Runs on the ml queue with concurrency=1, so enqueued tag_and_embed tasks only
|
|
start executing after this task finishes. Keyset pagination prevents the loop
|
|
from seeing its own pending enqueues as "still missing" and re-enqueueing them.
|
|
"""
|
|
from app.ml.wd14 import MODEL_VERSION as WD14_VER
|
|
from app.ml.siglip import MODEL_VERSION as SIGLIP_VER
|
|
|
|
enqueued_total = 0
|
|
last_id = 0
|
|
while True:
|
|
q = (
|
|
db.session.query(ImageRecord.id)
|
|
.outerjoin(
|
|
ImageTagPrediction,
|
|
and_(
|
|
ImageTagPrediction.image_id == ImageRecord.id,
|
|
ImageTagPrediction.model_version == WD14_VER,
|
|
),
|
|
)
|
|
.outerjoin(
|
|
ImageEmbedding,
|
|
and_(
|
|
ImageEmbedding.image_id == ImageRecord.id,
|
|
ImageEmbedding.model_version == SIGLIP_VER,
|
|
),
|
|
)
|
|
.filter(ImageRecord.id > last_id)
|
|
.filter(
|
|
(ImageTagPrediction.image_id.is_(None))
|
|
| (ImageEmbedding.image_id.is_(None))
|
|
)
|
|
.order_by(ImageRecord.id)
|
|
.limit(batch_size)
|
|
)
|
|
ids = [row[0] for row in q.all()]
|
|
if not ids:
|
|
break
|
|
for image_id in ids:
|
|
tag_and_embed.apply_async(args=[image_id], queue='ml')
|
|
enqueued_total += len(ids)
|
|
last_id = ids[-1]
|
|
log.info(f"backfill: enqueued batch of {len(ids)} (total {enqueued_total}, last_id {last_id})")
|
|
time.sleep(pause_seconds)
|
|
log.info(f"backfill: complete, enqueued {enqueued_total} images")
|
|
return {'status': 'ok', 'enqueued': enqueued_total}
|
|
|
|
|
|
@celery.task(bind=True, name='app.tasks.ml.recompute_centroid',
|
|
soft_time_limit=60, time_limit=120)
|
|
def recompute_centroid(self, tag_name: str):
|
|
"""Recompute the mean embedding for an eligible tag from its currently-associated images.
|
|
|
|
Eligible tag kinds are defined by ELIGIBLE_CENTROID_KINDS in
|
|
app.services.tag_suggestions. Tags of any other kind return early with
|
|
status='ineligible_kind' and do not write a centroid row.
|
|
"""
|
|
from app.ml.siglip import MODEL_VERSION as SIGLIP_VER
|
|
from app.services.tag_suggestions import ELIGIBLE_CENTROID_KINDS
|
|
|
|
tag = Tag.query.filter_by(name=tag_name).first()
|
|
if tag is None:
|
|
log.warning(f"recompute_centroid: tag {tag_name!r} not found")
|
|
return {'status': 'missing_tag'}
|
|
|
|
if tag.kind not in ELIGIBLE_CENTROID_KINDS:
|
|
log.info(f"recompute_centroid: tag {tag_name!r} has ineligible kind {tag.kind!r}")
|
|
return {'status': 'ineligible_kind'}
|
|
|
|
rows = (
|
|
db.session.query(ImageEmbedding.embedding)
|
|
.join(image_tags, image_tags.c.image_id == ImageEmbedding.image_id)
|
|
.filter(image_tags.c.tag_id == tag.id)
|
|
.filter(ImageEmbedding.model_version == SIGLIP_VER)
|
|
.all()
|
|
)
|
|
if not rows:
|
|
log.info(f"recompute_centroid: no embeddings yet for {tag_name}")
|
|
return {'status': 'no_embeddings'}
|
|
|
|
vectors = np.array([np.array(r[0], dtype=np.float32) for r in rows])
|
|
centroid = vectors.mean(axis=0)
|
|
|
|
stmt = pg_insert(TagReferenceEmbedding).values(
|
|
tag_name=tag_name,
|
|
tag_kind=tag.kind,
|
|
model_version=SIGLIP_VER,
|
|
centroid=centroid.tolist(),
|
|
reference_count=len(rows),
|
|
computed_at=func.now(),
|
|
)
|
|
stmt = stmt.on_conflict_do_update(
|
|
index_elements=['tag_name', 'model_version'],
|
|
set_={
|
|
'tag_kind': stmt.excluded.tag_kind,
|
|
'centroid': stmt.excluded.centroid,
|
|
'reference_count': stmt.excluded.reference_count,
|
|
'computed_at': func.now(),
|
|
},
|
|
)
|
|
db.session.execute(stmt)
|
|
db.session.commit()
|
|
log.info(f"recompute_centroid: {tag_name} (kind={tag.kind}) -> n={len(rows)}")
|
|
return {'status': 'ok', 'reference_count': len(rows), 'tag_kind': tag.kind}
|
|
|
|
|
|
@celery.task(bind=True, name='app.tasks.ml.recompute_all_centroids',
|
|
soft_time_limit=None, time_limit=None)
|
|
def recompute_all_centroids(self):
|
|
"""Enqueue recompute_centroid for every eligible tag with enough reference images.
|
|
|
|
Uses a single aggregate query to find tags with >= min_reference_images applied
|
|
images, then enqueues one recompute_centroid task per tag on the ml queue.
|
|
"""
|
|
from app.services.tag_suggestions import ELIGIBLE_CENTROID_KINDS
|
|
|
|
min_refs_row = TagSuggestionConfig.query.filter_by(key='min_reference_images').first()
|
|
min_refs = int(min_refs_row.value) if min_refs_row else 5
|
|
|
|
# Build an IS NULL / IN filter that covers ELIGIBLE_CENTROID_KINDS including None.
|
|
kinds_not_null = [k for k in ELIGIBLE_CENTROID_KINDS if k is not None]
|
|
allow_null = None in ELIGIBLE_CENTROID_KINDS
|
|
|
|
kind_filter = Tag.kind.in_(kinds_not_null)
|
|
if allow_null:
|
|
kind_filter = kind_filter | Tag.kind.is_(None)
|
|
|
|
rows = (
|
|
db.session.query(Tag.name, func.count(image_tags.c.image_id).label('n'))
|
|
.join(image_tags, image_tags.c.tag_id == Tag.id)
|
|
.filter(kind_filter)
|
|
.group_by(Tag.name)
|
|
.having(func.count(image_tags.c.image_id) >= min_refs)
|
|
.all()
|
|
)
|
|
|
|
enqueued = 0
|
|
for tag_name, n in rows:
|
|
recompute_centroid.apply_async(args=[tag_name], queue='ml')
|
|
enqueued += 1
|
|
log.info(f"recompute_all_centroids: enqueued {enqueued} tags (min_refs={min_refs})")
|
|
return {'status': 'ok', 'enqueued': enqueued, 'min_refs': min_refs}
|