compare view

core/detect/checkpoint/frames.py

@@ -147,6 +147,108 @@ def load_cached_frames_b64(timeline_id: str) -> list[dict]:
     return result
 
 
+# ---------------------------------------------------------------------------
+# Debug overlay storage — per job/stage/frame
+# ---------------------------------------------------------------------------
+
+def _overlay_prefix(timeline_id: str, job_id: str, stage: str) -> str:
+    return f"{CACHE_PREFIX}/{timeline_id}/overlays/{job_id}/{stage}/"
+
+
+def _overlay_key(timeline_id: str, job_id: str, stage: str, seq: int, name: str) -> str:
+    return f"{CACHE_PREFIX}/{timeline_id}/overlays/{job_id}/{stage}/{seq}_{name}.png"
+
+
+def save_overlays(
+    timeline_id: str,
+    job_id: str,
+    stage: str,
+    seq: int,
+    overlays: dict[str, str],
+):
+    """
+    Save debug overlay images (base64 PNG) to blob storage.
+
+    overlays: {overlay_key: base64_png_string}
+    e.g. {"edge_overlay_b64": "iVBOR...", "lines_overlay_b64": "iVBOR..."}
+    """
+    from core.storage.s3 import upload_file
+    import tempfile
+
+    for name, b64_data in overlays.items():
+        key = _overlay_key(timeline_id, job_id, stage, seq, name)
+        raw = base64.b64decode(b64_data)
+
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+            tmp.write(raw)
+            tmp_path = tmp.name
+
+        try:
+            upload_file(tmp_path, BUCKET, key)
+        finally:
+            os.unlink(tmp_path)
+
+    logger.info("Saved %d overlays for timeline %s job %s stage %s frame %d",
+                len(overlays), timeline_id, job_id, stage, seq)
+
+
+def load_overlays(
+    timeline_id: str,
+    job_id: str,
+    stage: str,
+    seq: int,
+) -> dict[str, str] | None:
+    """
+    Load debug overlay images from blob storage as base64 strings.
+
+    Returns {overlay_key: base64_png_string} or None if no overlays cached.
+    """
+    from core.storage.s3 import list_objects, download_to_temp
+
+    prefix = _overlay_prefix(timeline_id, job_id, stage)
+    seq_prefix = f"{seq}_"
+    objects = list_objects(BUCKET, prefix)
+
+    overlays = {}
+    for obj in objects:
+        filename = obj["key"].rsplit("/", 1)[-1]
+        if not filename.startswith(seq_prefix):
+            continue
+        name = filename[len(seq_prefix):].replace(".png", "")
+
+        tmp_path = download_to_temp(BUCKET, obj["key"])
+        try:
+            with open(tmp_path, "rb") as f:
+                overlays[name] = base64.b64encode(f.read()).decode()
+        finally:
+            os.unlink(tmp_path)
+
+    return overlays if overlays else None
+
+
+def list_overlay_frames(
+    timeline_id: str,
+    job_id: str,
+    stage: str,
+) -> list[int]:
+    """List frame sequences that have cached overlays."""
+    from core.storage.s3 import list_objects
+
+    prefix = _overlay_prefix(timeline_id, job_id, stage)
+    objects = list_objects(BUCKET, prefix)
+
+    seqs = set()
+    for obj in objects:
+        filename = obj["key"].rsplit("/", 1)[-1]
+        seq_str = filename.split("_")[0]
+        try:
+            seqs.add(int(seq_str))
+        except ValueError:
+            continue
+
+    return sorted(seqs)
+
+
 def clear_cache(timeline_id: str):
     """Delete the frame cache for a timeline."""
     from core.storage.s3 import delete_objects