074cdcdbed
- Delete FaceEmbeddingDb module (face_embedding_db.rs) - Stub match_faces_iterative, generate_seed_embeddings, tmdb_match_handler - Remove sync_trace_embeddings, populate_face_embeddings_to_qdrant - Remove embedding from face.json output (face_processor.py) - Remove embedding from PG UPDATE (store_traced_faces.py) - Remove workspace traces staging (checkin.rs, qdrant_workspace.rs) - Fix tests: add pose_angle to Face, hand_nodes to TkgResult Disabled functions (need reimplement with _faces): - match_faces_iterative (identity agent) - generate_seed_embeddings (TMDb seeds) - tmdb_match_handler (TMDb matching) - cluster_face_embeddings, search_similar_faces - merge_traces_within_cuts
258 lines
8.8 KiB
Python
258 lines
8.8 KiB
Python
#!/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
|
|
from collections import defaultdict
|
|
import numpy as np
|
|
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 merge_traces_within_cuts(face_data: dict, cut_scenes: list) -> dict:
|
|
"""Merge traces within the same cut - DISABLED (no embeddings)."""
|
|
# TODO: Reimplement with Qdrant _faces collection
|
|
return face_data
|
|
|
|
|
|
def run_face_tracker(
|
|
face_json_path: str, traced_json_path: str, filter_eyes: bool = False
|
|
) -> 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"])
|
|
faces = []
|
|
for f in frame.get("faces", []):
|
|
bbox = f.get("bbox", f)
|
|
face = {
|
|
"x": bbox.get("x", f.get("x", 0)),
|
|
"y": bbox.get("y", f.get("y", 0)),
|
|
"width": bbox.get("width", f.get("width", 0)),
|
|
"height": bbox.get("height", f.get("height", 0)),
|
|
"confidence": f.get("confidence", 0.0),
|
|
}
|
|
if "landmarks" in f:
|
|
face["landmarks"] = f["landmarks"]
|
|
if "embedding" in f:
|
|
face["embedding"] = f["embedding"]
|
|
faces.append(face)
|
|
frames_dict[fnum] = {
|
|
"frame_number": frame["frame"],
|
|
"time_seconds": frame.get("timestamp", 0),
|
|
"faces": 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),
|
|
}
|
|
|
|
# Eye filter: remove faces without at least one eye landmark
|
|
if filter_eyes:
|
|
removed = 0
|
|
for fnum_str, frm_data in face_data.get("frames", {}).items():
|
|
faces = frm_data.get("faces", [])
|
|
kept = []
|
|
for face in faces:
|
|
lm = face.get("landmarks", {})
|
|
if len(lm.get("left_eye", [])) > 0 or len(lm.get("right_eye", [])) > 0:
|
|
kept.append(face)
|
|
else:
|
|
removed += 1
|
|
frm_data["faces"] = kept
|
|
print(f"[TRACE] Eye filter: {removed} faces without eyes removed")
|
|
|
|
print(f"[TRACE] Processing {len(face_data.get('frames', {}))} frames")
|
|
|
|
# Embeddings no longer loaded from DB - use IoU-only tracking
|
|
file_uuid = (
|
|
face_json_path.split("/")[-1]
|
|
.replace(".face.json", "")
|
|
.replace("_traced.json", "")
|
|
)
|
|
|
|
# Load cut boundaries from cut.json (same directory as face.json)
|
|
cut_boundaries = None
|
|
cut_scenes = None
|
|
cuts_path = face_json_path.replace("_traced.json", ".cut.json").replace(
|
|
".face.json", ".cut.json"
|
|
)
|
|
if os.path.exists(cuts_path):
|
|
with open(cuts_path) as f:
|
|
cuts = json.load(f)
|
|
cut_scenes = cuts.get("scenes", [])
|
|
cut_boundaries = {s["start_frame"] for s in cut_scenes if s["start_frame"] > 0}
|
|
print(f"[TRACE] Loaded {len(cut_boundaries)} cut boundaries")
|
|
|
|
face_data = track_faces(
|
|
face_data, use_embedding=False, cut_boundaries=cut_boundaries
|
|
)
|
|
|
|
# Merge traces within same cut (same person re-appearing after occlusion/pose change)
|
|
if cut_scenes and len(cut_scenes) > 0:
|
|
face_data = merge_traces_within_cuts(face_data, cut_scenes)
|
|
|
|
metadata = face_data.get("metadata", {})
|
|
metadata["tracking_method"] = "iou_only"
|
|
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")
|
|
if face_id is None:
|
|
face_id = f"face_{trace_id}"
|
|
attributes = face.get("attributes")
|
|
|
|
bbox = json.dumps({"x": x, "y": y, "width": w, "height": h})
|
|
|
|
try:
|
|
cur.execute(
|
|
f"""
|
|
UPDATE {schema}.face_detections
|
|
SET trace_id = %s, face_id = %s
|
|
WHERE file_uuid = %s AND frame_number = %s
|
|
AND x = %s AND y = %s AND width = %s AND height = %s
|
|
""",
|
|
(
|
|
trace_id,
|
|
face_id,
|
|
file_uuid,
|
|
frame_num,
|
|
x,
|
|
y,
|
|
w,
|
|
h,
|
|
),
|
|
)
|
|
if cur.rowcount > 0:
|
|
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")
|
|
|
|
parser.add_argument("--uuid", help="UUID for Redis tracking (accepted by executor)")
|
|
parser.add_argument(
|
|
"--filter-eyes",
|
|
action="store_true",
|
|
help="Remove faces without eye landmarks before tracking",
|
|
)
|
|
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, filter_eyes=args.filter_eyes)
|
|
|
|
# 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()
|