diff --git a/app/tasks/ml.py b/app/tasks/ml.py index 807052a..a43e1ce 100644 --- a/app/tasks/ml.py +++ b/app/tasks/ml.py @@ -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)) diff --git a/app/utils/image_importer.py b/app/utils/image_importer.py index 7e071b1..f025ac3 100644 --- a/app/utils/image_importer.py +++ b/app/utils/image_importer.py @@ -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.