phase 10

detect/checkpoint/serializer.py

@@ -0,0 +1,133 @@
+"""State serialization — DetectState ↔ JSON-compatible dict."""
+
+from __future__ import annotations
+
+import dataclasses
+
+from detect.models import (
+    BoundingBox,
+    BrandDetection,
+    Frame,
+    PipelineStats,
+    TextCandidate,
+)
+
+
+# ---------------------------------------------------------------------------
+# Serialize helpers
+# ---------------------------------------------------------------------------
+
+def serialize_frame_meta(frame: Frame) -> dict:
+    meta = {
+        "sequence": frame.sequence,
+        "chunk_id": frame.chunk_id,
+        "timestamp": frame.timestamp,
+        "perceptual_hash": frame.perceptual_hash,
+    }
+    return meta
+
+
+def serialize_text_candidate(tc: TextCandidate) -> dict:
+    bbox_dict = dataclasses.asdict(tc.bbox)
+    candidate = {
+        "frame_sequence": tc.frame.sequence,
+        "bbox": bbox_dict,
+        "text": tc.text,
+        "ocr_confidence": tc.ocr_confidence,
+    }
+    return candidate
+
+
+def serialize_state(state: dict, frames_manifest: dict[int, str]) -> dict:
+    """
+    Serialize DetectState to a JSON-compatible dict.
+
+    Frame images are replaced with S3 key references.
+    TextCandidate.frame references become frame_sequence integers.
+    """
+    frames = state.get("frames", [])
+    filtered = state.get("filtered_frames", [])
+
+    manifest_strs = {str(k): v for k, v in frames_manifest.items()}
+    frames_meta = [serialize_frame_meta(f) for f in frames]
+    filtered_seqs = [f.sequence for f in filtered]
+
+    boxes_serialized = {}
+    for seq, boxes in state.get("boxes_by_frame", {}).items():
+        boxes_serialized[str(seq)] = [dataclasses.asdict(b) for b in boxes]
+
+    text_candidates = [serialize_text_candidate(tc) for tc in state.get("text_candidates", [])]
+    unresolved = [serialize_text_candidate(tc) for tc in state.get("unresolved_candidates", [])]
+    detections = [dataclasses.asdict(d) for d in state.get("detections", [])]
+    stats = dataclasses.asdict(state.get("stats", PipelineStats()))
+
+    checkpoint = {
+        "job_id": state.get("job_id", ""),
+        "video_path": state.get("video_path", ""),
+        "profile_name": state.get("profile_name", ""),
+        "config_overrides": state.get("config_overrides", {}),
+        "frames_manifest": manifest_strs,
+        "frames_meta": frames_meta,
+        "filtered_frame_sequences": filtered_seqs,
+        "boxes_by_frame": boxes_serialized,
+        "text_candidates": text_candidates,
+        "unresolved_candidates": unresolved,
+        "detections": detections,
+        "stats": stats,
+    }
+    return checkpoint
+
+
+# ---------------------------------------------------------------------------
+# Deserialize helpers
+# ---------------------------------------------------------------------------
+
+def deserialize_text_candidate(d: dict, frame_map: dict[int, Frame]) -> TextCandidate:
+    frame = frame_map[d["frame_sequence"]]
+    bbox = BoundingBox(**d["bbox"])
+    candidate = TextCandidate(
+        frame=frame,
+        bbox=bbox,
+        text=d["text"],
+        ocr_confidence=d["ocr_confidence"],
+    )
+    return candidate
+
+
+def deserialize_state(checkpoint: dict, frames: list[Frame]) -> dict:
+    """Reconstitute DetectState from a checkpoint dict + loaded frames."""
+    frame_map = {f.sequence: f for f in frames}
+
+    filtered_seqs = set(checkpoint.get("filtered_frame_sequences", []))
+    filtered_frames = [f for f in frames if f.sequence in filtered_seqs]
+
+    boxes_by_frame = {}
+    for seq_str, box_dicts in checkpoint.get("boxes_by_frame", {}).items():
+        seq = int(seq_str)
+        boxes_by_frame[seq] = [BoundingBox(**b) for b in box_dicts]
+
+    text_candidates = [
+        deserialize_text_candidate(d, frame_map)
+        for d in checkpoint.get("text_candidates", [])
+    ]
+    unresolved_candidates = [
+        deserialize_text_candidate(d, frame_map)
+        for d in checkpoint.get("unresolved_candidates", [])
+    ]
+    detections = [BrandDetection(**d) for d in checkpoint.get("detections", [])]
+    stats = PipelineStats(**checkpoint.get("stats", {}))
+
+    state = {
+        "job_id": checkpoint.get("job_id", ""),
+        "video_path": checkpoint.get("video_path", ""),
+        "profile_name": checkpoint.get("profile_name", ""),
+        "config_overrides": checkpoint.get("config_overrides", {}),
+        "frames": frames,
+        "filtered_frames": filtered_frames,
+        "boxes_by_frame": boxes_by_frame,
+        "text_candidates": text_candidates,
+        "unresolved_candidates": unresolved_candidates,
+        "detections": detections,
+        "stats": stats,
+    }
+    return state