RAG with Video: Visual Question-Answering over Video Frames

!pip install -q google-genai Pillow==11.0.0 transformers==4.46.3 matplotlib opencv-python-headless git+https://github.com/KarAnalytics/llm_cascade.git

# =============================================================================
# Imports + LLM Cascade
# =============================================================================

import math
import numpy as np
from pathlib import Path

import cv2
from PIL import Image, ImageDraw, ImageFont

from llm_cascade import get_cascade

llm = get_cascade()


def has_llm_provider():
    """True if llm_cascade has at least one provider configured."""
    return True  # get_cascade() already validates


def generate_text(prompt, system_prompt=None):
    """Send a text prompt through llm_cascade.

    Returns (response_text, provider_string).
    """
    resp = llm.generate(prompt, system_prompt=system_prompt)
    return resp.text, f"{resp.provider}/{resp.model}"


print("LLM cascade ready.")

VIDEO_PATH = "product_showcase.mp4"
FPS = 10
WIDTH, HEIGHT = 640, 480

# Scene definitions
SCENES = [
    {
        "name": "title",
        "start": 0, "end": 3,
        "label": "Acme Corp\nProduct Showcase 2026",
        "bg_color": (40, 60, 120),
        "shape": None,
        "product": None,
    },
    {
        "name": "widget_alpha",
        "start": 3, "end": 7,
        "label": "Widget Alpha — $12.99",
        "bg_color": (255, 245, 238),
        "shape": "circle",
        "shape_color": (231, 76, 60),
        "product": {"name": "Widget Alpha", "price": 12.99, "color": "red", "shape": "circle"},
    },
    {
        "name": "gadget_beta",
        "start": 7, "end": 11,
        "label": "Gadget Beta — $8.75",
        "bg_color": (240, 255, 240),
        "shape": "triangle",
        "shape_color": (46, 204, 113),
        "product": {"name": "Gadget Beta", "price": 8.75, "color": "green", "shape": "triangle"},
    },
    {
        "name": "tool_gamma",
        "start": 11, "end": 15,
        "label": "Tool Gamma — $34.99",
        "bg_color": (235, 245, 255),
        "shape": "square",
        "shape_color": (52, 152, 219),
        "product": {"name": "Tool Gamma", "price": 34.99, "color": "blue", "shape": "square"},
    },
    {
        "name": "device_delta",
        "start": 15, "end": 19,
        "label": "Device Delta — $19.25",
        "bg_color": (255, 248, 230),
        "shape": "diamond",
        "shape_color": (230, 126, 34),
        "product": {"name": "Device Delta", "price": 19.25, "color": "orange", "shape": "diamond"},
    },
    {
        "name": "closing",
        "start": 19, "end": 22,
        "label": "Thank you!\nVisit acme-corp.example.com",
        "bg_color": (40, 60, 120),
        "shape": None,
        "product": None,
    },
]


def draw_frame_pil(scene, t, width=WIDTH, height=HEIGHT):
    """Draw a single video frame for a given scene and time offset using PIL."""
    bg = scene["bg_color"]
    img = Image.new("RGB", (width, height), bg)
    draw = ImageDraw.Draw(img)

    # Load font
    try:
        font_large = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 28)
        font_small = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 22)
    except (OSError, IOError):
        try:
            font_large = ImageFont.truetype("arial.ttf", 28)
            font_small = ImageFont.truetype("arial.ttf", 22)
        except (OSError, IOError):
            font_large = ImageFont.load_default()
            font_small = font_large

    cx, cy = width // 2, height // 2

    if scene["shape"] is None:
        # Title or closing card — centered text
        text_color = (255, 255, 255)
        lines = scene["label"].split("\n")
        total_h = len(lines) * 40
        y_start = cy - total_h // 2
        for j, line in enumerate(lines):
            bbox = draw.textbbox((0, 0), line, font=font_large)
            tw = bbox[2] - bbox[0]
            draw.text(((width - tw) // 2, y_start + j * 40), line, fill=text_color, font=font_large)
    else:
        # Product scene — animated shape (gentle bounce)
        bounce = int(15 * math.sin(t * 3))
        shape_cy = cy - 30 + bounce
        r = 80
        color = scene["shape_color"]

        if scene["shape"] == "circle":
            draw.ellipse([cx - r, shape_cy - r, cx + r, shape_cy + r], fill=color, outline=(0, 0, 0), width=3)
        elif scene["shape"] == "square":
            draw.rectangle([cx - r, shape_cy - r, cx + r, shape_cy + r], fill=color, outline=(0, 0, 0), width=3)
        elif scene["shape"] == "triangle":
            pts = [(cx, shape_cy - r), (cx - r, shape_cy + r), (cx + r, shape_cy + r)]
            draw.polygon(pts, fill=color, outline=(0, 0, 0), width=3)
        elif scene["shape"] == "diamond":
            pts = [(cx, shape_cy - r), (cx + r, shape_cy), (cx, shape_cy + r), (cx - r, shape_cy)]
            draw.polygon(pts, fill=color, outline=(0, 0, 0), width=3)

        # Product label at top
        bbox = draw.textbbox((0, 0), scene["label"], font=font_small)
        tw = bbox[2] - bbox[0]
        draw.text(((width - tw) // 2, 20), scene["label"], fill=(30, 30, 30), font=font_small)

        # Timestamp at bottom-right
        ts_text = f"t={t:.1f}s"
        draw.text((width - 90, height - 30), ts_text, fill=(120, 120, 120), font=font_small)

    return img


def generate_video(path, scenes, fps, width=WIDTH, height=HEIGHT):
    """Generate an MP4 video from scene definitions."""
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(path, fourcc, fps, (width, height))

    for scene in scenes:
        duration = scene["end"] - scene["start"]
        n_frames = int(duration * fps)
        for i in range(n_frames):
            t = scene["start"] + i / fps
            pil_img = draw_frame_pil(scene, t, width, height)
            # Convert PIL (RGB) -> OpenCV (BGR)
            frame = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
            writer.write(frame)

    writer.release()


generate_video(VIDEO_PATH, SCENES, FPS)
total_duration = SCENES[-1]["end"]
file_size_kb = Path(VIDEO_PATH).stat().st_size / 1024
print(f"Video generated: '{VIDEO_PATH}'")
print(f"  Duration: {total_duration}s | FPS: {FPS} | Resolution: {WIDTH}x{HEIGHT} | Size: {file_size_kb:.0f} KB")

def extract_keyframes(video_path, interval_sec=1.0):
    """Extract frames from a video at a fixed time interval.

    Returns a list of dicts: {"timestamp": float, "image": PIL.Image}
    """
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps
    frame_interval = int(fps * interval_sec)

    frames = []
    frame_idx = 0
    while True:
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        ret, frame = cap.read()
        if not ret:
            break
        # Convert BGR -> RGB -> PIL
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil_img = Image.fromarray(rgb)
        timestamp = frame_idx / fps
        frames.append({"timestamp": timestamp, "image": pil_img})
        frame_idx += frame_interval

    cap.release()
    return frames, duration


keyframes, video_duration = extract_keyframes(VIDEO_PATH, interval_sec=1.0)
print(f"Extracted {len(keyframes)} keyframes from {video_duration:.1f}s video (1 frame/sec)")
print(f"Timestamps: {[f'{kf["timestamp"]:.1f}s' for kf in keyframes]}")

import matplotlib.pyplot as plt

# Show a subset of keyframes (every other frame to save space)
display_frames = keyframes[::2]
n = len(display_frames)
cols = min(n, 6)
rows = math.ceil(n / cols)

fig, axes = plt.subplots(rows, cols, figsize=(4 * cols, 3.5 * rows))
axes_flat = axes.flat if hasattr(axes, 'flat') else [axes]
for i, ax in enumerate(axes_flat):
    if i < n:
        ax.imshow(display_frames[i]["image"])
        ax.set_title(f"t = {display_frames[i]['timestamp']:.1f}s", fontsize=10)
    ax.axis("off")
fig.suptitle("Keyframes from Product Showcase Video (every 2s)", fontsize=13, fontweight="bold")
plt.tight_layout()
plt.show()

from transformers import CLIPProcessor, CLIPModel
import torch

CLIP_MODEL_NAME = "openai/clip-vit-base-patch32"
clip_model = CLIPModel.from_pretrained(CLIP_MODEL_NAME)
clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME)

print(f"CLIP model loaded: {CLIP_MODEL_NAME}")
print(f"Embedding dimension: {clip_model.config.projection_dim}")

def embed_frames(keyframes, model, processor):
    """Compute CLIP embeddings for a list of keyframe images."""
    images = [kf["image"].convert("RGB") for kf in keyframes]
    inputs = processor(images=images, return_tensors="pt", padding=True)
    with torch.no_grad():
        embeddings = model.get_image_features(**inputs)
    embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
    return embeddings.numpy()


def embed_text(query, model, processor):
    """Compute CLIP embedding for a text query."""
    inputs = processor(text=[query], return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        embeddings = model.get_text_features(**inputs)
    embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
    return embeddings.numpy()


# Build the frame index
frame_embeddings = embed_frames(keyframes, clip_model, clip_processor)
print(f"Frame index built: {frame_embeddings.shape[0]} frames, {frame_embeddings.shape[1]}-dim embeddings")

def get_scene_for_timestamp(timestamp, scenes):
    """Look up which scene a timestamp belongs to."""
    for scene in scenes:
        if scene["start"] <= timestamp < scene["end"]:
            return scene
    return scenes[-1]  # default to last scene


def retrieve_frames(query, frame_embeddings, keyframes, model, processor, top_k=3):
    """Retrieve the top-k most relevant video frames for a text query."""
    query_emb = embed_text(query, model, processor)
    similarities = (query_emb @ frame_embeddings.T).flatten()
    top_indices = similarities.argsort()[::-1][:top_k]

    results = []
    for idx in top_indices:
        ts = keyframes[idx]["timestamp"]
        scene = get_scene_for_timestamp(ts, SCENES)
        results.append({
            "rank": len(results) + 1,
            "frame_index": int(idx),
            "timestamp": ts,
            "scene_name": scene["name"],
            "product": scene.get("product"),
            "image": keyframes[idx]["image"],
            "similarity": float(similarities[idx]),
        })
    return results


def show_retrieved_frames(query, results):
    """Display retrieved frames side by side."""
    n = len(results)
    fig, axes = plt.subplots(1, n, figsize=(6 * n, 5))
    if n == 1:
        axes = [axes]
    for ax, r in zip(axes, results):
        ax.imshow(r["image"])
        label = f"#{r['rank']} t={r['timestamp']:.1f}s\n{r['scene_name']}\nsim={r['similarity']:.3f}"
        ax.set_title(label, fontsize=10)
        ax.axis("off")
    fig.suptitle(f'Query: "{query}"', fontsize=13, fontweight="bold")
    plt.tight_layout()
    plt.show()


# Test retrieval
test_query = "red circle product"
results = retrieve_frames(test_query, frame_embeddings, keyframes, clip_model, clip_processor)

print(f"Query: '{test_query}'\n")
for r in results:
    prod = r['product']
    prod_str = f"{prod['name']} ({prod['color']} {prod['shape']})" if prod else "(no product)"
    print(f"  #{r['rank']} t={r['timestamp']:.1f}s | {r['scene_name']} | {prod_str} | sim={r['similarity']:.4f}")

show_retrieved_frames(test_query, results)

def build_frame_context(results):
    """Build a text description of retrieved frames."""
    lines = ["Retrieved video frames (ranked by visual similarity):"]
    for r in results:
        prod = r["product"]
        prod_str = (
            f"{prod['name']}: {prod['color']} {prod['shape']}, ${prod['price']:.2f}"
            if prod else "(title/closing card)"
        )
        lines.append(
            f"  #{r['rank']} timestamp={r['timestamp']:.1f}s, scene={r['scene_name']}, "
            f"product={prod_str}, similarity={r['similarity']:.3f}"
        )
    return "\n".join(lines)


VIDEO_RAG_SYSTEM_PROMPT = (
    "You are a helpful video analyst. You are given keyframes extracted from a product "
    "showcase video, along with their timestamps and metadata. Answer the user's question "
    "based on what you can SEE in the frames and the provided metadata. Reference specific "
    "timestamps when relevant (e.g., 'at t=5.0s'). If the frames do not contain enough "
    "information, say so. Do not invent products or details not shown."
)


def generate_answer_with_frames(question, results):
    """Send retrieved frames + metadata to a multimodal LLM via llm_cascade."""
    context_text = build_frame_context(results)
    prompt = f"{context_text}\n\nQUESTION: {question}"
    # Pass PIL frames directly; the cascade handles per-provider encoding.
    frames = [r["image"] for r in results]
    resp = llm.generate_multimodal(
        prompt, images=frames, system_prompt=VIDEO_RAG_SYSTEM_PROMPT)
    return resp.text, f"{resp.provider}/{resp.model}"


def generate_answer_without_rag(question):
    """Direct LLM answer with NO video frames."""
    return generate_text(question)


print("Answer generation functions ready (llm_cascade).")

def answer_with_video_rag(question, top_k=3):
    """Full Video RAG: extract frames -> CLIP retrieval -> multimodal LLM answer."""
    print(f"  [Step 1] Embedding query with CLIP ...")
    results = retrieve_frames(
        question, frame_embeddings, keyframes, clip_model, clip_processor, top_k=top_k,
    )
    print(f"  [Step 2] Retrieved {len(results)} frame(s):")
    for r in results:
        prod_name = r['product']['name'] if r['product'] else '(card)'
        print(f"           #{r['rank']} t={r['timestamp']:.1f}s — {prod_name} (sim={r['similarity']:.3f})")

    print(f"  [Step 3] Generating multimodal answer ...")
    answer, provider = generate_answer_with_frames(question, results)
    print(f"  [Step 4] Answer generated (provider: {provider})")
    return answer, results, provider


print("Video RAG pipeline ready.")

questions = [
    "What is the red product shown in the video and how much does it cost?",
    #"At what point in the video does the green triangle appear?",
    #"What is the most expensive product in this showcase video?",
    #"What shape is the orange product?",
    "What text appears at the beginning of the video?",
]


def preview(text, max_len=800):
    text = text or ""
    return text[:max_len] + ("..." if len(text) > max_len else "")


if not has_llm_provider():
    print("Error: No LLM API key configured. Set at least one API key in Colab Secrets.")
else:
    for i, q in enumerate(questions, start=1):
        print("\n" + "=" * 80)
        print(f"Q{i}. {q}")
        print("=" * 80)

        # --- WITH VIDEO RAG ---
        print("\n--- WITH VIDEO RAG ---")
        try:
            answer_rag, retrieved, prov = answer_with_video_rag(q)
            show_retrieved_frames(q, retrieved)
            print(f"\n  Answer (provider: {prov}):")
            print(preview(answer_rag))
        except Exception as e:
            print(f"  RAG error: {e}")

        # --- WITHOUT RAG ---
        print("\n--- WITHOUT RAG (LLM knowledge only) ---")
        try:
            answer_direct, prov_direct = generate_answer_without_rag(q)
            print(f"  Answer (provider: {prov_direct}):")
            print(preview(answer_direct))
        except Exception as e:
            print(f"  Direct error: {e}")

        print()

# Change this to any question about the product showcase video
my_question = "How many products are shown in total, and what are their prices?"

if has_llm_provider():
    print(f"Question: {my_question}\n")

    print("=" * 60)
    print("WITH VIDEO RAG:")
    print("=" * 60)
    try:
        answer, retrieved, provider = answer_with_video_rag(my_question, top_k=4)
        show_retrieved_frames(my_question, retrieved)
        print(f"\nAnswer ({provider}):\n{answer}")
    except Exception as e:
        print(f"Error: {e}")

    print("\n" + "=" * 60)
    print("WITHOUT RAG:")
    print("=" * 60)
    try:
        answer_direct, prov_direct = generate_answer_without_rag(my_question)
        print(f"\nAnswer ({prov_direct}):\n{answer_direct}")
    except Exception as e:
        print(f"Error: {e}")
else:
    print("No LLM API key configured. Set at least one API key in Colab Secrets.")

# Direct Gemini call on the whole video file -- Gemini-only feature that the
# cascade does not exercise for other providers.
from google import genai
from google.genai import types as genai_types
from llm_cascade.providers import _get_key

gemini_key = _get_key("GEMINI_API_KEY")
if not gemini_key:
    print("Skipping: GEMINI_API_KEY not configured.")
else:
    question = "What products appear in this video and in what order?"
    video_bytes = Path(VIDEO_PATH).read_bytes()

    client = genai.Client(api_key=gemini_key)
    parts = [
        genai_types.Part.from_text(text=question),
        genai_types.Part.from_bytes(data=video_bytes, mime_type="video/mp4"),
    ]
    resp = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=parts,
        config=genai_types.GenerateContentConfig(
            thinking_config=genai_types.ThinkingConfig(thinking_budget=0),
        ),
    )
    print(f"Question: {question}")
    print(f"Video:    {VIDEO_PATH} ({Path(VIDEO_PATH).stat().st_size // 1024} KB)")
    print(f"\nAnswer:\n{resp.text}")