81 lines
2.1 KiB
Python
81 lines
2.1 KiB
Python
"""Frame image storage — save/load to S3/MinIO as JPEGs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from detect.models import Frame
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BUCKET = os.environ.get("S3_BUCKET_OUT", "out")
|
|
CHECKPOINT_PREFIX = "checkpoints"
|
|
|
|
|
|
def save_frames(job_id: str, frames: list[Frame]) -> dict[int, str]:
|
|
"""
|
|
Save frame images to S3 as JPEGs.
|
|
|
|
Returns manifest: {sequence: s3_key}
|
|
"""
|
|
from core.storage.s3 import upload_file
|
|
|
|
manifest = {}
|
|
|
|
for frame in frames:
|
|
key = f"{CHECKPOINT_PREFIX}/{job_id}/frames/{frame.sequence}.jpg"
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
|
|
img = Image.fromarray(frame.image)
|
|
img.save(tmp, format="JPEG", quality=85)
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
upload_file(tmp_path, BUCKET, key)
|
|
finally:
|
|
os.unlink(tmp_path)
|
|
|
|
manifest[frame.sequence] = key
|
|
|
|
logger.info("Saved %d frames to s3://%s/%s/%s/frames/",
|
|
len(frames), BUCKET, CHECKPOINT_PREFIX, job_id)
|
|
return manifest
|
|
|
|
|
|
def load_frames(manifest: dict[int, str], frame_metadata: list[dict]) -> list[Frame]:
|
|
"""
|
|
Load frame images from S3 and reconstitute Frame objects.
|
|
|
|
frame_metadata: list of dicts with sequence, chunk_id, timestamp, perceptual_hash.
|
|
"""
|
|
from core.storage.s3 import download_to_temp
|
|
|
|
meta_map = {m["sequence"]: m for m in frame_metadata}
|
|
frames = []
|
|
|
|
for seq, key in manifest.items():
|
|
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)
|
|
|
|
meta = meta_map.get(seq, {})
|
|
frame = Frame(
|
|
sequence=seq,
|
|
chunk_id=meta.get("chunk_id", 0),
|
|
timestamp=meta.get("timestamp", 0.0),
|
|
image=image_array,
|
|
perceptual_hash=meta.get("perceptual_hash", ""),
|
|
)
|
|
frames.append(frame)
|
|
|
|
frames.sort(key=lambda f: f.sequence)
|
|
return frames
|