momentry_core/scripts/store_traced_faces.py

#!/opt/homebrew/bin/python3.11
"""
Store Traced Faces - Pipeline integration for face trace + position data

Flow:
1. Reads face.json output from face_processor.py
2. Runs face_tracker.py to assign trace_id per face (IoU + embedding)
3. Inserts traced faces into face_detections table with trace_id and position (x,y,w,h)

Usage:
    python store_traced_faces.py --file-uuid <uuid> [--face-json <path>]

TKG Export:
    trace_id + position (x,y,w,h) per frame enables spatial-temporal graph construction.
    Each trace is a temporal entity; position tracks movement across frames.
"""

import sys
import os
import json
import argparse
import psycopg2
import psycopg2.extras
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "utils"))

# Config
DB_URL = os.environ.get("DATABASE_URL", "postgresql://accusys@localhost:5432/momentry")
SCHEMA = os.environ.get("MOMENTRY_DB_SCHEMA", "dev")
OUTPUT_DIR = os.environ.get("MOMENTRY_OUTPUT_DIR", "/Users/accusys/momentry/output_dev")


def get_conn():
    return psycopg2.connect(DB_URL)


def run_face_tracker(face_json_path: str, traced_json_path: str) -> str:
    """Run face_tracker.py on face.json, returns path to face_traced.json"""
    from face_tracker import track_faces

    with open(face_json_path) as f:
        face_data = json.load(f)

    # V2.0 uses list format (FaceResult), convert to dict for face_tracker
    if isinstance(face_data.get("frames"), list):
        frames_dict = {}
        for frame in face_data["frames"]:
            fnum = str(frame["frame"])
            frames_dict[fnum] = {
                "frame_number": frame["frame"],
                "time_seconds": frame.get("timestamp", 0),
                "faces": frame.get("faces", []),
            }
        face_data["frames"] = frames_dict
        # Preserve metadata (fps needed by face_tracker)
        if "metadata" not in face_data:
            face_data["metadata"] = {
                "fps": face_data.get("fps", 30.0),
                "total_frames": face_data.get("frame_count", 0),
            }

    print(f"[TRACE] Processing {len(face_data.get('frames', {}))} frames")

    face_data = track_faces(face_data, use_embedding=True)
    metadata = face_data.get("metadata", {})
    metadata["tracking_method"] = "iou_embedding"
    metadata["tracked_at"] = datetime.now().isoformat()
    face_data["metadata"] = metadata

    with open(traced_json_path, "w") as f:
        json.dump(face_data, f, indent=2, ensure_ascii=False)

    trace_count = len(face_data.get("traces", {}))
    print(f"[TRACE] Completed: {trace_count} traces -> {traced_json_path}")
    return traced_json_path


def store_traced_faces(file_uuid: str, traced_json_path: str, schema: str = SCHEMA):
    """Insert traced face detections into face_detections table with trace_id"""
    conn = get_conn()
    cur = conn.cursor()

    with open(traced_json_path) as f:
        data = json.load(f)

    frames = data.get("frames", {})
    total_stored = 0

    for frame_num_str, frame_data in sorted(frames.items(), key=lambda x: int(x[0])):
        frame_num = int(frame_num_str)
        faces = frame_data.get("faces", [])

        for face in faces:
            trace_id = face.get("trace_id")
            if trace_id is None:
                continue

            x = face.get("x", 0)
            y = face.get("y", 0)
            w = face.get("width", 0)
            h = face.get("height", 0)
            confidence = face.get("confidence", 0.0)
            face_id = face.get("face_id")
            attributes = face.get("attributes")
            embedding = face.get("embedding")

            bbox = json.dumps({"x": x, "y": y, "width": w, "height": h})
            embed_vec = embedding if embedding and len(embedding) > 0 else None

            try:
                cur.execute(
                    f"""
                    INSERT INTO {schema}.face_detections
                        (file_uuid, frame_number, face_id, trace_id,
                         x, y, width, height, confidence, embedding)
                    VALUES (%s, %s, %s, %s,
                            %s, %s, %s, %s, %s, %s)
                    ON CONFLICT DO NOTHING
                    """,
                    (
                        file_uuid, frame_num, face_id, trace_id,
                        x, y, w, h, confidence,
                        embed_vec,
                    ),
                )
                total_stored += 1
            except Exception as e:
                print(f"[TRACE] Error storing face at frame {frame_num}: {e}")
                conn.rollback()
                continue

    conn.commit()

    # Log trace summary
    cur.execute(
        f"SELECT COUNT(DISTINCT trace_id) FROM {schema}.face_detections WHERE file_uuid = %s AND trace_id IS NOT NULL",
        (file_uuid,),
    )
    db_trace_count = cur.fetchone()[0]

    cur.close()
    conn.close()

    print(f"[TRACE] Stored {total_stored} face detections, {db_trace_count} unique traces in DB")
    return total_stored, db_trace_count


def main():
    parser = argparse.ArgumentParser(description="Store traced faces in DB")
    parser.add_argument("--file-uuid", required=True, help="Video file UUID")
    parser.add_argument("--face-json", help="Path to face.json (default: auto-detect)")
    parser.add_argument("--schema", default=SCHEMA, help="DB schema name")
    args = parser.parse_args()

    face_json = args.face_json or os.path.join(
        OUTPUT_DIR, f"{args.file_uuid}.face.json"
    )
    traced_json = os.path.join(OUTPUT_DIR, f"{args.file_uuid}.face_traced.json")

    if not os.path.exists(face_json):
        print(f"[TRACE] face.json not found: {face_json}", file=sys.stderr)
        sys.exit(1)

    # Step 1: Run face tracker
    run_face_tracker(face_json, traced_json)

    # Step 2: Store in DB with trace_id
    total, traces = store_traced_faces(args.file_uuid, traced_json, args.schema)
    print(f"[TRACE] Done: {total} detections, {traces} traces")


if __name__ == "__main__":
    main()