8f05a7c188
- Update ASR, face, OCR, pose processors - Add release pre-flight check script - Add synonym generation, chunk processing scripts - Add face recognition, stamp search utilities
127 lines
3.4 KiB
Python
127 lines
3.4 KiB
Python
#!/opt/homebrew/bin/python3.11
|
|
"""
|
|
Update person demographics (Age, Gender) using InsightFace.
|
|
This script scans the representative face of each person and updates the DB.
|
|
"""
|
|
|
|
import os
|
|
import cv2
|
|
import json
|
|
import psycopg2
|
|
import numpy as np
|
|
import insightface
|
|
from insightface.app import FaceAnalysis
|
|
|
|
# Configuration
|
|
DB_CONFIG = {"host": "localhost", "user": "accusys", "dbname": "momentry"}
|
|
VIDEO_PATH = "output/384b0ff44aaaa1f1/384b0ff44aaaa1f1.mp4"
|
|
|
|
|
|
def get_face_app():
|
|
app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
|
|
app.prepare(ctx_id=0, det_size=(640, 640))
|
|
return app
|
|
|
|
|
|
def get_person_frames(conn):
|
|
"""Get one frame timestamp for each person."""
|
|
cur = conn.cursor()
|
|
# Get the first appearance time for each person to save processing time
|
|
cur.execute("""
|
|
SELECT person_id, MIN(start_time) as start_time
|
|
FROM dev.person_appearances
|
|
GROUP BY person_id
|
|
""")
|
|
return cur.fetchall()
|
|
|
|
|
|
def get_frame_at_time(video_path, time_sec):
|
|
"""Extract a single frame from video."""
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
return None
|
|
|
|
# InsightFace expects BGR (OpenCV default), but we need to seek correctly
|
|
# cv2 CAP_PROP_POS_MSEC is sometimes inaccurate, better to use FPS if known
|
|
# But for a "representative" frame, seeking by ms is usually "good enough"
|
|
cap.set(cv2.CAP_PROP_POS_MSEC, time_sec * 1000)
|
|
ret, frame = cap.read()
|
|
cap.release()
|
|
|
|
if ret:
|
|
return frame
|
|
return None
|
|
|
|
|
|
def analyze_face(app, frame):
|
|
"""Run InsightFace to get age/gender."""
|
|
faces = app.get(frame)
|
|
if faces:
|
|
# Just take the first face found
|
|
face = faces[0]
|
|
age = int(face.age) if hasattr(face, "age") else None
|
|
gender = (
|
|
"female"
|
|
if (hasattr(face, "gender") and face.gender == 0)
|
|
else ("male" if (hasattr(face, "gender") and face.gender == 1) else None)
|
|
)
|
|
return age, gender
|
|
return None, None
|
|
|
|
|
|
def update_person_db(conn, person_id, age, gender):
|
|
"""Update DB."""
|
|
cur = conn.cursor()
|
|
cur.execute(
|
|
"""
|
|
UPDATE dev.person_identities
|
|
SET age = %s, gender = %s
|
|
WHERE person_id = %s
|
|
""",
|
|
(age, gender, person_id),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def main():
|
|
print("=== Person Demographics Updater ===")
|
|
|
|
# 1. Init DB
|
|
conn = psycopg2.connect(**DB_CONFIG)
|
|
|
|
# 2. Init Model
|
|
print("Loading InsightFace model...")
|
|
app = get_face_app()
|
|
|
|
# 3. Get List
|
|
persons = get_person_frames(conn)
|
|
print(f"Found {len(persons)} persons to process.")
|
|
|
|
for i, (person_id, start_time) in enumerate(persons):
|
|
print(
|
|
f"Processing {i + 1}/{len(persons)}: {person_id} (Time: {start_time:.1f}s)"
|
|
)
|
|
|
|
# 4. Get Frame
|
|
frame = get_frame_at_time(VIDEO_PATH, start_time)
|
|
|
|
if frame is not None:
|
|
# 5. Analyze
|
|
age, gender = analyze_face(app, frame)
|
|
|
|
if age and gender:
|
|
print(f" -> Detected: Age {age}, Gender {gender}")
|
|
# 6. Update
|
|
update_person_db(conn, person_id, age, gender)
|
|
else:
|
|
print(f" -> Face not found or attributes missing in this frame.")
|
|
else:
|
|
print(f" -> Failed to retrieve frame.")
|
|
|
|
print("=== Done ===")
|
|
conn.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|