mediaproc/core/detect/checkpoint/frames.py

"""
Frame cache — per-timeline frame storage in blob storage (S3/MinIO).

Frames are extracted from chunks once, cached as JPEGs at
cache/timelines/{timeline_id}/frames/{seq}.jpg in the app's
blob storage. Any job on the timeline reads from the cache.
Cache is clearable and rebuildable from chunks.

Uses the same storage backend as the rest of the app, so it
works across lambdas, GPU boxes, and local dev.
"""

from __future__ import annotations

import base64
import io
import logging
import os
import tempfile

import numpy as np
from PIL import Image

from core.detect.models import Frame

logger = logging.getLogger(__name__)

BUCKET = os.environ.get("S3_BUCKET", "mpr")
CACHE_PREFIX = "cache/timelines"


def _frame_key(timeline_id: str, seq: int) -> str:
    return f"{CACHE_PREFIX}/{timeline_id}/frames/{seq}.jpg"


def _list_prefix(timeline_id: str) -> str:
    return f"{CACHE_PREFIX}/{timeline_id}/frames/"


def cache_exists(timeline_id: str) -> bool:
    """Check if frame cache exists for a timeline."""
    from core.storage.s3 import list_objects

    objects = list_objects(BUCKET, _list_prefix(timeline_id))
    return len(objects) > 0


def cache_frames(timeline_id: str, frames: list[Frame], quality: int = 85) -> int:
    """
    Write frames to blob storage as JPEGs.

    Returns number of frames cached.
    """
    from core.storage.s3 import upload_file

    for frame in frames:
        key = _frame_key(timeline_id, frame.sequence)

        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
            img = Image.fromarray(frame.image)
            img.save(tmp, format="JPEG", quality=quality)
            tmp_path = tmp.name

        try:
            upload_file(tmp_path, BUCKET, key)
        finally:
            os.unlink(tmp_path)

    logger.info("Cached %d frames for timeline %s", len(frames), timeline_id)
    return len(frames)


def load_cached_frames(timeline_id: str) -> list[Frame]:
    """
    Load all cached frames as Frame objects with numpy arrays.

    Returns empty list if cache doesn't exist.
    """
    from core.storage.s3 import list_objects, download_to_temp

    objects = list_objects(BUCKET, _list_prefix(timeline_id))
    if not objects:
        return []

    frames = []
    for obj in objects:
        key = obj["key"]
        filename = key.rsplit("/", 1)[-1]
        if not filename.endswith(".jpg"):
            continue
        seq = int(filename.replace(".jpg", ""))

        tmp_path = download_to_temp(BUCKET, key)
        try:
            img = Image.open(tmp_path).convert("RGB")
            image_array = np.array(img)
        finally:
            os.unlink(tmp_path)

        frame = Frame(
            sequence=seq,
            chunk_id=0,
            timestamp=0.0,
            image=image_array,
            perceptual_hash="",
        )
        frames.append(frame)

    frames.sort(key=lambda f: f.sequence)
    return frames


def load_cached_frames_b64(timeline_id: str) -> list[dict]:
    """
    Load cached frames as base64 JPEGs for the UI.

    Returns list of {seq, timestamp, jpeg_b64}.
    """
    from core.storage.s3 import list_objects, download_to_temp

    objects = list_objects(BUCKET, _list_prefix(timeline_id))
    if not objects:
        return []

    result = []
    for obj in objects:
        key = obj["key"]
        filename = key.rsplit("/", 1)[-1]
        if not filename.endswith(".jpg"):
            continue
        seq = int(filename.replace(".jpg", ""))

        tmp_path = download_to_temp(BUCKET, key)
        try:
            with open(tmp_path, "rb") as f:
                jpeg_b64 = base64.b64encode(f.read()).decode()
        finally:
            os.unlink(tmp_path)

        result.append({
            "seq": seq,
            "timestamp": 0.0,
            "jpeg_b64": jpeg_b64,
        })

    result.sort(key=lambda f: f["seq"])
    return result


def clear_cache(timeline_id: str):
    """Delete the frame cache for a timeline."""
    from core.storage.s3 import delete_objects

    prefix = _list_prefix(timeline_id)
    delete_objects(BUCKET, prefix)
    logger.info("Cleared frame cache for timeline %s", timeline_id)


def frames_to_b64(frames: list[Frame], quality: int = 75) -> list[dict]:
    """
    Convert in-memory Frame objects to base64 JPEG dicts.

    For API responses when frames are already in memory.
    """
    result = []
    for frame in frames:
        buf = io.BytesIO()
        img = Image.fromarray(frame.image)
        img.save(buf, format="JPEG", quality=quality)
        jpeg_b64 = base64.b64encode(buf.getvalue()).decode()

        result.append({
            "seq": frame.sequence,
            "timestamp": frame.timestamp,
            "jpeg_b64": jpeg_b64,
        })

    result.sort(key=lambda f: f["seq"])
    return result