65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
#!/opt/homebrew/bin/python3.11
|
|
"""Insert face detections from traced JSON into DB."""
|
|
import json, os, sys
|
|
import psycopg2
|
|
import psycopg2.extras
|
|
|
|
DB_URL = os.environ.get("DATABASE_URL", "postgresql://accusys@localhost:5432/momentry")
|
|
|
|
def insert_faces(file_uuid, traced_json_path, schema):
|
|
conn = psycopg2.connect(DB_URL)
|
|
cur = conn.cursor()
|
|
|
|
with open(traced_json_path) as f:
|
|
data = json.load(f)
|
|
|
|
frames = data.get("frames", {})
|
|
metadata = data.get("metadata", {})
|
|
fps = metadata.get("fps", 24.0)
|
|
|
|
total = 0
|
|
for frame_num_str, frame_data in sorted(frames.items(), key=lambda x: int(x[0])):
|
|
frame_num = int(frame_num_str)
|
|
ts = frame_num / fps
|
|
faces = frame_data.get("faces", [])
|
|
|
|
for face in faces:
|
|
x = int(face.get("x", 0))
|
|
y = int(face.get("y", 0))
|
|
w = int(face.get("width", 0))
|
|
h = int(face.get("height", 0))
|
|
confidence = face.get("confidence", 0.0)
|
|
trace_id = face.get("trace_id")
|
|
embedding = face.get("embedding")
|
|
|
|
try:
|
|
cur.execute(
|
|
f"""
|
|
INSERT INTO {schema}.face_detections
|
|
(file_uuid, frame_number, timestamp_secs, x, y, width, height, confidence, trace_id, embedding)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT DO NOTHING
|
|
""",
|
|
(file_uuid, frame_num, ts, x, y, w, h, confidence, trace_id, embedding),
|
|
)
|
|
if cur.rowcount > 0:
|
|
total += 1
|
|
except Exception as e:
|
|
print(f"[INSERT] Error at frame {frame_num}: {e}")
|
|
conn.rollback()
|
|
|
|
conn.commit()
|
|
cur.close()
|
|
conn.close()
|
|
print(f"[INSERT] Inserted {total} face detections into {schema}.face_detections")
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description="Insert face detections")
|
|
parser.add_argument("--file-uuid", required=True)
|
|
parser.add_argument("--face-json", required=True)
|
|
parser.add_argument("--schema", default="public")
|
|
args = parser.parse_args()
|
|
|
|
insert_faces(args.file_uuid, args.face_json, args.schema)
|