#!/usr/bin/env python """ GCP-side preprocessing. Enumerates a GCS bucket of movie mp4s, downloads missing ones, extracts frames + clips, and captions with Gemini in both `frames` and `video_clip` modes. Steps (all run by default): frames — uniform + fps frame extraction clip_embeddings — CLIP features (needed by VideoAgent) clips — mp4 clip extraction (shot_grouped) caption_frames — Gemini frames-mode captioning (fixed 1-min chunks) caption_clips — Gemini video-clip-mode captioning (over extracted clips) Usage: python run_gemini.py python run_gemini.py --step caption_frames python run_gemini.py --bucket my-bucket --project my-project """ from __future__ import annotations import argparse import io import json import math import os import subprocess import sys import time from pathlib import Path _PROJECT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(_PROJECT)) import numpy as np from PIL import Image from tqdm.auto import tqdm from util.backbones import call_backbone, init_backbone from util.frames import extract_frames_fps, extract_frames_uniform, get_clip_embeddings from util.paths import ( VIDEO_DIR, GCS_BUCKET, get_captions_dir, get_clip_path, get_clips_dir, get_frames_dir, ensure_dir, ) from util.prompts import BASE_CAPTION_PROMPT from util.shots import extract_all_clips, load_clips def list_bucket(bucket_name: str, project: str) -> list[str]: from google.cloud import storage client = storage.Client(project=project) bucket = client.bucket(bucket_name) movies = [] for blob in bucket.list_blobs(): if blob.name.endswith('.mp4'): movies.append(blob.name[:-4]) return sorted(set(movies)) def resolve_video_path(vid: str, bucket_name: str, project: str) -> Path | None: local = VIDEO_DIR / f'{vid}.mp4' if local.exists(): return local local.parent.mkdir(parents=True, exist_ok=True) try: from google.cloud import storage client = storage.Client(project=project) blob = client.bucket(bucket_name).blob(f'{vid}.mp4') if blob.exists(): blob.download_to_filename(str(local)) return local except Exception as e: print(f' GCS download error for {vid}: {e}') # fallback: gcloud CLI try: r = subprocess.run( ['gcloud', 'storage', 'cp', f'gs://{bucket_name}/{vid}.mp4', str(local)], capture_output=True, text=True, ) if r.returncode == 0: return local except Exception: pass return None def ffprobe_duration(video_path: Path) -> float: try: r = subprocess.run( ['ffprobe', '-v', 'quiet', '-print_format', 'json', '-show_format', str(video_path)], capture_output=True, text=True, ) return float(json.loads(r.stdout)['format']['duration']) except Exception: from decord import VideoReader, cpu vr = VideoReader(str(video_path), ctx=cpu(0), num_threads=1) d = len(vr) / vr.get_avg_fps() del vr return d def step_frames(movies: list[str], args) -> dict[str, float]: durations: dict[str, float] = {} for vid in tqdm(movies, desc='frames'): v = resolve_video_path(vid, args.bucket, args.project) if v is None: print(f' SKIP {vid}: not in GCS and not local') continue dur = ffprobe_duration(v) durations[vid] = dur if args.fps > 0: d = get_frames_dir(vid, 'fps', args.fps) expected = int(np.ceil(dur * args.fps)) if len(list(d.glob('*.jpg'))) < expected * 0.96: extract_frames_fps(str(v), fps=args.fps, cache_dir=str(d), low_memory=True) for n in args.uniform_n: d = get_frames_dir(vid, 'uniform', n) if len(list(d.glob('*.jpg'))) < n: extract_frames_uniform(str(v), n, cache_dir=str(d), low_memory=True) return durations def step_clip_embeddings(movies: list[str], args) -> None: clip_n = max(args.uniform_n) for vid in tqdm(movies, desc=f'CLIP (uniform_{clip_n})'): out = get_clip_path(vid, 'uniform', clip_n) if out.exists(): continue uni_dir = get_frames_dir(vid, 'uniform', clip_n) jpgs = sorted(uni_dir.glob('*.jpg')) if not jpgs: continue images = [Image.open(f) for f in jpgs] get_clip_embeddings(images, cache_path=str(out)) def step_clips(movies: list[str], durations: dict[str, float], args) -> None: for chunk_mode in args.chunk_modes: for vid in tqdm(movies, desc=f'clips ({chunk_mode})'): v = VIDEO_DIR / f'{vid}.mp4' if not v.exists(): continue dur = durations.get(vid, ffprobe_duration(v)) clips_dir = get_clips_dir(vid, chunk_mode) manifest = clips_dir / 'manifest.json' if manifest.exists(): try: with open(manifest) as f: m = json.load(f) if all((clips_dir / item['clip']).exists() for item in m): continue except Exception: pass try: clips = extract_all_clips(vid, str(v), args.assume_fps, dur, chunk_mode=chunk_mode, with_audio=True) print(f' {vid}: {len(clips)} clips') except Exception as e: print(f' Error {vid}: {e}') # ---------- Captioning ---------- def _pil_list(chunk_start: float, chunk_end: float, fps_frames: list[Image.Image], fps: float): return [fps_frames[j] for j in range(len(fps_frames)) if chunk_start <= (j / fps) < chunk_end] def step_caption_frames(movies: list[str], durations: dict[str, float], args) -> None: """Frames-mode captioning: 1-minute fixed chunks, caption each with Gemini over fps frames.""" cap_dir_name = 'gemini_frames' print(f'=== {cap_dir_name} (frames) ===') for vid in tqdm(movies, desc=cap_dir_name): dur = durations.get(vid) if dur is None: v = VIDEO_DIR / f'{vid}.mp4' if not v.exists(): continue dur = ffprobe_duration(v) n_chunks = int(math.ceil(dur / 60)) chunks = [(i * 60, min((i + 1) * 60, dur)) for i in range(n_chunks)] cap_dir = ensure_dir(get_captions_dir(cap_dir_name, vid)) existing = {int(p.stem) for p in cap_dir.glob('*.txt') if p.stem.isdigit()} if all(i in existing for i in range(len(chunks))): continue fps_dir = get_frames_dir(vid, 'fps', args.fps) jpgs = sorted(fps_dir.glob('*.jpg')) if not jpgs: print(f' SKIP {vid}: no fps frames under {fps_dir}') continue fps_frames = [] for f in jpgs: img = Image.open(f); img.load(); fps_frames.append(img) for chunk_idx, (s, e) in enumerate(tqdm(chunks, desc=vid, leave=False)): out = cap_dir / f'{chunk_idx}.txt' if out.exists(): continue frames_in_chunk = _pil_list(s, e, fps_frames, args.fps) if not frames_in_chunk: continue try: caption = call_backbone( 'gemini', BASE_CAPTION_PROMPT, images=frames_in_chunk, system_instruction='You caption video segments.', json_mode=False, max_output_tokens=2048, temperature=0.1, thinking_level='LOW', ) out.write_text(caption) except Exception as ex: print(f' Error {vid} chunk {chunk_idx}: {type(ex).__name__}: {ex!r}') time.sleep(2) def step_caption_clips(movies: list[str], args) -> None: """Video-clip-mode captioning: feed each mp4 clip directly to Gemini.""" from google.genai import types from google.genai.types import GenerateContentConfig # Direct Gemini client; can't route clip bytes through call_backbone's image path. from google import genai as _genai _client = _genai.Client(vertexai=True, project=args.project, location=args.gemini_location) for chunk_mode in args.chunk_modes: cap_dir_name = f'gemini_video_clip_{chunk_mode}' if chunk_mode != 'fixed' else 'gemini_video_clip' print(f'=== {cap_dir_name} (clips) ===') for vid in tqdm(movies, desc=cap_dir_name): cached_clips = load_clips(vid, chunk_mode=chunk_mode) if not cached_clips: continue cap_dir = ensure_dir(get_captions_dir(cap_dir_name, vid)) existing = {int(p.stem) for p in cap_dir.glob('*.txt') if p.stem.isdigit()} if all(i in existing for i in range(len(cached_clips))): continue for chunk_idx, (clip_path, _, _) in enumerate(tqdm(cached_clips, desc=vid, leave=False)): out = cap_dir / f'{chunk_idx}.txt' if out.exists(): continue try: with open(clip_path, 'rb') as cf: clip_bytes = cf.read() response = _client.models.generate_content( model=args.gemini_model, contents=[ types.Part.from_bytes(data=clip_bytes, mime_type='video/mp4'), BASE_CAPTION_PROMPT, ], config=GenerateContentConfig( thinking_config=types.ThinkingConfig(thinking_level=types.ThinkingLevel.LOW), ), ) out.write_text((response.text or '').strip()) except Exception as ex: print(f' Error {vid} chunk {chunk_idx}: {type(ex).__name__}: {ex!r}') time.sleep(2) def main() -> None: from util.config import DEFAULT_MODELS, DEFAULT_VERTEX_LOCATION, GOOGLE_CLOUD_PROJECT_ENV ap = argparse.ArgumentParser(description=__doc__.split('\n\n', 1)[0]) ap.add_argument('--step', default='all', choices=['frames', 'clip_embeddings', 'clips', 'caption_frames', 'caption_clips', 'all']) ap.add_argument('--bucket', default=GCS_BUCKET) ap.add_argument('--project', default=os.environ.get(GOOGLE_CLOUD_PROJECT_ENV)) ap.add_argument('--gemini-model', default=DEFAULT_MODELS['gemini']) ap.add_argument('--gemini-location', default=DEFAULT_VERTEX_LOCATION) ap.add_argument('--uniform-n', type=int, nargs='+', default=[64, 128, 256]) ap.add_argument('--fps', type=float, default=0.5) ap.add_argument('--chunk-modes', nargs='+', default=['shot_grouped']) ap.add_argument('--assume-fps', type=float, default=24.0) ap.add_argument('--vids', default=None, help='Comma-separated list of video_file_ids to scope captioning ' 'to. Default: all bucket movies.') args = ap.parse_args() if args.vids: wanted = [v.strip() for v in args.vids.split(',') if v.strip()] movies = [] for v in wanted: local = VIDEO_DIR / f'{v}.mp4' if local.exists(): movies.append(v) else: print(f' requested --vids not local: {v} (expected {local})') if not movies: sys.exit(f'No vids found locally for --vids={args.vids!r}; aborting.') print(f'Scope: {len(movies)} movie(s) -> {movies}') else: movies = list_bucket(args.bucket, args.project) print(f'Found {len(movies)} .mp4 in gs://{args.bucket}') if args.step in ('caption_frames', 'all'): init_backbone('gemini', project_id=args.project, location=args.gemini_location, model_id=args.gemini_model) durations: dict[str, float] = {} if args.step in ('frames', 'all'): durations = step_frames(movies, args) if args.step in ('clip_embeddings', 'all'): step_clip_embeddings(movies, args) if args.step in ('clips', 'all'): if not durations: for vid in tqdm(movies, desc='ffprobe'): v = VIDEO_DIR / f'{vid}.mp4' if v.exists(): durations[vid] = ffprobe_duration(v) step_clips(movies, durations, args) if args.step in ('caption_frames', 'all'): if not durations: for vid in tqdm(movies, desc='ffprobe'): v = VIDEO_DIR / f'{vid}.mp4' if v.exists(): durations[vid] = ffprobe_duration(v) step_caption_frames(movies, durations, args) if args.step in ('caption_clips', 'all'): step_caption_clips(movies, args) if __name__ == '__main__': main()