feat(ml): video tagging + embedding via frame sampling
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>
This commit is contained in:
+70
-17
@@ -5,7 +5,7 @@ import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
from sqlalchemy import and_, func, or_
|
||||
from sqlalchemy import and_, func
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
from app.celery_app import celery
|
||||
@@ -25,18 +25,71 @@ 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=120, time_limit=180)
|
||||
soft_time_limit=240, time_limit=360)
|
||||
def tag_and_embed(self, image_id: int):
|
||||
"""Run WD14 + SigLIP on one image and persist predictions and embedding."""
|
||||
"""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:
|
||||
@@ -47,15 +100,20 @@ def tag_and_embed(self, image_id: int):
|
||||
log.warning(f"tag_and_embed: file missing at {image.filepath}")
|
||||
return {'status': 'file_missing'}
|
||||
|
||||
# WD14 + SigLIP are image-only; videos bypass inference entirely.
|
||||
from app.utils.image_importer import VIDEO_EXTS
|
||||
if image.filepath.lower().endswith(VIDEO_EXTS):
|
||||
return {'status': 'unsupported_format'}
|
||||
is_video = image.filepath.lower().endswith(VIDEO_EXTS)
|
||||
|
||||
try:
|
||||
t0 = time.time()
|
||||
raw = wd14.infer_filtered(image.filepath, min_any=WD14_STORE_FLOOR)
|
||||
embedding = siglip.infer(image.filepath)
|
||||
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
|
||||
@@ -84,8 +142,9 @@ def tag_and_embed(self, image_id: int):
|
||||
))
|
||||
|
||||
db.session.commit()
|
||||
log.info(f"tag_and_embed: image {image_id} done in {t_inf:.2f}s ({len(raw)} predictions)")
|
||||
return {'status': 'ok', 'predictions': len(raw), 'duration_s': t_inf}
|
||||
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)
|
||||
@@ -107,11 +166,6 @@ def backfill(self, batch_size: int = 50, pause_seconds: float = 0.5):
|
||||
"""
|
||||
from app.ml.wd14 import MODEL_VERSION as WD14_VER
|
||||
from app.ml.siglip import MODEL_VERSION as SIGLIP_VER
|
||||
from app.utils.image_importer import VIDEO_EXTS
|
||||
|
||||
video_filter = ~or_(*[
|
||||
func.lower(ImageRecord.filepath).like(f'%{ext}') for ext in VIDEO_EXTS
|
||||
])
|
||||
|
||||
enqueued_total = 0
|
||||
last_id = 0
|
||||
@@ -133,7 +187,6 @@ def backfill(self, batch_size: int = 50, pause_seconds: float = 0.5):
|
||||
),
|
||||
)
|
||||
.filter(ImageRecord.id > last_id)
|
||||
.filter(video_filter)
|
||||
.filter(
|
||||
(ImageTagPrediction.image_id.is_(None))
|
||||
| (ImageEmbedding.image_id.is_(None))
|
||||
|
||||
@@ -556,6 +556,72 @@ def generate_video_thumbnail_mirrored(video_path: str, time_position: str = "00:
|
||||
return generate_video_thumbnail(str(video_path_p), str(thumb_path), time_position=time_position)
|
||||
|
||||
|
||||
def _probe_video_duration(video_path: str) -> float | None:
|
||||
"""Return duration in seconds via ffprobe, or None on failure."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffprobe", "-v", "error",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "default=noprint_wrappers=1:nokey=1",
|
||||
str(video_path),
|
||||
],
|
||||
capture_output=True, text=True, timeout=30, check=True,
|
||||
)
|
||||
return float(result.stdout.strip())
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def extract_video_frames(video_path: str, count: int = 10) -> list[str]:
|
||||
"""Extract `count` evenly-spaced JPEG frames from a video into tempfiles.
|
||||
|
||||
Frames are sampled between 10% and 90% of the video's duration so title
|
||||
cards and black-fade endings don't dominate the sample. Returns a list of
|
||||
absolute paths; the caller is responsible for deleting the containing
|
||||
temp directory when done. Returns [] if duration probe fails or no frames
|
||||
could be extracted.
|
||||
"""
|
||||
duration = _probe_video_duration(video_path)
|
||||
if duration is None or duration <= 0:
|
||||
return []
|
||||
|
||||
if count <= 1:
|
||||
fractions = [0.5]
|
||||
else:
|
||||
fractions = [0.1 + 0.8 * (i / (count - 1)) for i in range(count)]
|
||||
|
||||
tmpdir = tempfile.mkdtemp(prefix="vframes_")
|
||||
frame_paths: list[str] = []
|
||||
for i, frac in enumerate(fractions):
|
||||
t = duration * frac
|
||||
out = os.path.join(tmpdir, f"frame_{i:02d}.jpg")
|
||||
try:
|
||||
subprocess.run(
|
||||
[
|
||||
"ffmpeg", "-y",
|
||||
"-ss", f"{t:.3f}",
|
||||
"-i", str(video_path),
|
||||
"-vframes", "1",
|
||||
out,
|
||||
],
|
||||
check=True,
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||||
timeout=FFMPEG_THUMB_TIMEOUT,
|
||||
)
|
||||
if os.path.exists(out) and os.path.getsize(out) > 0:
|
||||
frame_paths.append(out)
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
|
||||
continue
|
||||
|
||||
if not frame_paths:
|
||||
try:
|
||||
os.rmdir(tmpdir)
|
||||
except OSError:
|
||||
pass
|
||||
return frame_paths
|
||||
|
||||
|
||||
def transcode_video_to_mp4(input_path: str, output_path: str = None) -> str | None:
|
||||
"""
|
||||
Transcode a video to H.264 MP4 for universal browser playback.
|
||||
|
||||
Reference in New Issue
Block a user