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
283 lines
8.9 KiB
Python
283 lines
8.9 KiB
Python
#!/opt/homebrew/bin/python3.11
|
||
"""
|
||
Face Clustering Processor
|
||
職責：將短暫的 Face ID 聚合為持續的 Person ID，並自動綁定 Speaker。
|
||
"""
|
||
|
||
import cv2
|
||
import json
|
||
import numpy as np
|
||
import os
|
||
import sys
|
||
import psycopg2
|
||
from sklearn.cluster import AgglomerativeClustering
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
try:
|
||
from deepface import DeepFace
|
||
|
||
HAS_DEEPFACE = True
|
||
except ImportError:
|
||
print("❌ DeepFace not found. Run: pip install deepface")
|
||
sys.exit(1)
|
||
|
||
# 設定
|
||
UUID = os.getenv("UUID", "quick_preview")
|
||
OUTPUT_DIR = os.getenv("MOMENTRY_OUTPUT_DIR", "./output")
|
||
VIDEO_PATH = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.mp4")
|
||
FACE_JSON_PATH = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.face.json")
|
||
OUTPUT_JSON_PATH = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.face_clustered.json")
|
||
ASRX_JSON_PATH = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.asrx.json")
|
||
DB_URL = os.getenv("DATABASE_URL", "postgresql://accusys@localhost:5432/momentry")
|
||
|
||
|
||
def optimized_clustering(embeddings):
|
||
"""
|
||
Optimized Clustering for large datasets (e.g. 25k faces).
|
||
Strategy: Sample -> Agglomerative -> Centroid Assignment
|
||
"""
|
||
import numpy as np
|
||
from sklearn.cluster import AgglomerativeClustering
|
||
from sklearn.metrics.pairwise import cosine_distances
|
||
|
||
n_faces = len(embeddings)
|
||
print(f" 🚀 Starting optimized clustering for {n_faces} faces...")
|
||
|
||
# 1. Sampling
|
||
sample_size = min(5000, n_faces)
|
||
if n_faces > sample_size:
|
||
indices = np.random.choice(n_faces, sample_size, replace=False)
|
||
sample_embeddings = embeddings[indices]
|
||
else:
|
||
sample_embeddings = embeddings
|
||
indices = np.arange(n_faces)
|
||
|
||
print(f" 📊 Sampling {len(sample_embeddings)} faces for clustering structure...")
|
||
|
||
# 2. Agglomerative Clustering on Sample
|
||
clustering = AgglomerativeClustering(
|
||
n_clusters=None, distance_threshold=0.4, metric="cosine", linkage="average"
|
||
)
|
||
sample_labels = clustering.fit_predict(sample_embeddings)
|
||
|
||
unique_labels = set(sample_labels)
|
||
n_clusters = len(unique_labels)
|
||
print(f" 🔍 Found {n_clusters} unique clusters in sample.")
|
||
|
||
# 3. Compute Centroids for each cluster
|
||
centroids = []
|
||
for label in unique_labels:
|
||
cluster_mask = sample_labels == label
|
||
cluster_faces = sample_embeddings[cluster_mask]
|
||
# Mean embedding
|
||
centroid = np.mean(cluster_faces, axis=0)
|
||
centroids.append(centroid)
|
||
|
||
centroids = np.array(centroids) # Shape: (n_clusters, 512)
|
||
|
||
# 4. Assign all faces to nearest centroid
|
||
# Batch processing to save memory
|
||
print(f" 🏃 Assigning {n_faces} faces to {n_clusters} clusters...")
|
||
all_labels = np.zeros(n_faces, dtype=int)
|
||
|
||
batch_size = 5000
|
||
for start in range(0, n_faces, batch_size):
|
||
end = min(start + batch_size, n_faces)
|
||
batch = embeddings[start:end]
|
||
dists = cosine_distances(batch, centroids)
|
||
all_labels[start:end] = np.argmin(dists, axis=1)
|
||
|
||
return all_labels
|
||
|
||
|
||
def main():
|
||
if not os.path.exists(FACE_JSON_PATH):
|
||
print("❌ Face JSON not found.")
|
||
return
|
||
|
||
with open(FACE_JSON_PATH) as f:
|
||
face_data = json.load(f)
|
||
|
||
frames_list = face_data.get("frames", [])
|
||
if not frames_list:
|
||
print("❌ No frames in JSON.")
|
||
return
|
||
|
||
cap = cv2.VideoCapture(VIDEO_PATH)
|
||
embeddings = []
|
||
face_refs = []
|
||
|
||
print(f"🔍 Extracting face embeddings from {UUID}...")
|
||
|
||
for frame_idx, frame_obj in enumerate(frames_list):
|
||
ts = frame_obj.get("timestamp")
|
||
faces = frame_obj.get("faces", [])
|
||
if not faces:
|
||
continue
|
||
|
||
if ts is not None:
|
||
cap.set(cv2.CAP_PROP_POS_MSEC, ts * 1000)
|
||
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
continue
|
||
|
||
for face_idx, face in enumerate(faces):
|
||
x, y, w, h = face["x"], face["y"], face["width"], face["height"]
|
||
margin = 5
|
||
crop = frame[
|
||
max(0, y - margin) : y + h + margin, max(0, x - margin) : x + w + margin
|
||
]
|
||
|
||
if crop is None or crop.size == 0:
|
||
continue
|
||
|
||
try:
|
||
res = DeepFace.represent(
|
||
img_path=crop, model_name="ArcFace", enforce_detection=False
|
||
)
|
||
if res and "embedding" in res[0]:
|
||
embeddings.append(res[0]["embedding"])
|
||
face_refs.append({"frame_idx": frame_idx, "face_idx": face_idx})
|
||
except Exception:
|
||
pass
|
||
|
||
cap.release()
|
||
|
||
if not embeddings:
|
||
print("❌ No embeddings extracted.")
|
||
return
|
||
|
||
embeddings = np.array(embeddings)
|
||
print(f"✅ Extracted {len(embeddings)} face embeddings.")
|
||
|
||
# 2. 聚類
|
||
print(f"🧠 Clustering {len(embeddings)} faces...")
|
||
clustering = AgglomerativeClustering(
|
||
n_clusters=None, distance_threshold=0.4, metric="cosine", linkage="average"
|
||
)
|
||
labels = clustering.fit_predict(embeddings)
|
||
|
||
unique_labels = set(labels)
|
||
label_to_person = {l: f"Person_{i}" for i, l in enumerate(unique_labels)}
|
||
print(
|
||
f"👥 Detected {len(unique_labels)} unique persons: {[label_to_person[l] for l in unique_labels]}"
|
||
)
|
||
|
||
# 3. 更新 JSON
|
||
for ref, label in zip(face_refs, labels):
|
||
f_idx = ref["frame_idx"]
|
||
face_idx = ref["face_idx"]
|
||
person_id = label_to_person[label]
|
||
|
||
if f_idx < len(frames_list):
|
||
faces = frames_list[f_idx].get("faces", [])
|
||
if face_idx < len(faces):
|
||
frames_list[f_idx]["faces"][face_idx]["person_id"] = person_id
|
||
|
||
# 保存
|
||
with open(OUTPUT_JSON_PATH, "w", encoding="utf-8") as f:
|
||
json.dump(face_data, f, indent=2, ensure_ascii=False)
|
||
print(f"✅ Saved clustered data to {OUTPUT_JSON_PATH}")
|
||
|
||
# 4. 自動綁定 Speaker
|
||
auto_bind_speakers()
|
||
|
||
|
||
def auto_bind_speakers():
|
||
if not os.path.exists(OUTPUT_JSON_PATH) or not os.path.exists(ASRX_JSON_PATH):
|
||
print("⚠️ Missing data for speaker binding.")
|
||
return
|
||
|
||
with open(OUTPUT_JSON_PATH) as f:
|
||
face_clustered = json.load(f)
|
||
with open(ASRX_JSON_PATH) as f:
|
||
asrx_data = json.load(f)
|
||
|
||
print("🔗 Auto-binding Speakers to Persons...")
|
||
|
||
# 建立 Face 時間列表
|
||
face_spans = []
|
||
for frame_obj in face_clustered.get("frames", []):
|
||
ts = frame_obj.get("timestamp")
|
||
for face in frame_obj.get("faces", []):
|
||
person_id = face.get("person_id")
|
||
if person_id and ts is not None:
|
||
face_spans.append({"ts": ts, "person_id": person_id})
|
||
|
||
speaker_person_counts = {}
|
||
|
||
# 對於每個說話片段，找出畫面中出現的人
|
||
for seg in asrx_data.get("segments", []):
|
||
start = seg.get("start")
|
||
end = seg.get("end")
|
||
speaker = seg.get("speaker_id")
|
||
if not speaker:
|
||
continue
|
||
|
||
# 找時間重疊
|
||
candidates = [f for f in face_spans if start <= f["ts"] <= end]
|
||
if candidates:
|
||
# 投票
|
||
person_counts = {}
|
||
for c in candidates:
|
||
pid = c["person_id"]
|
||
person_counts[pid] = person_counts.get(pid, 0) + 1
|
||
|
||
if speaker not in speaker_person_counts:
|
||
speaker_person_counts[speaker] = {}
|
||
|
||
best_person = max(person_counts, key=person_counts.get)
|
||
speaker_person_counts[speaker][best_person] = (
|
||
speaker_person_counts[speaker].get(best_person, 0) + 1
|
||
)
|
||
|
||
# 寫入資料庫
|
||
try:
|
||
conn = psycopg2.connect(DB_URL)
|
||
cur = conn.cursor()
|
||
|
||
for speaker, persons in speaker_person_counts.items():
|
||
if not persons:
|
||
continue
|
||
best_person = max(persons, key=persons.get)
|
||
print(
|
||
f" 🎤 {speaker} is likely {best_person} ({persons[best_person]} votes)"
|
||
)
|
||
|
||
# 1. 找或建 Talent
|
||
cur.execute("SELECT id FROM talents WHERE real_name = %s", (best_person,))
|
||
row = cur.fetchone()
|
||
|
||
if row:
|
||
talent_id = row[0]
|
||
else:
|
||
cur.execute(
|
||
"INSERT INTO talents (real_name) VALUES (%s) RETURNING id",
|
||
(best_person,),
|
||
)
|
||
talent_id = cur.fetchone()[0]
|
||
print(f" ✨ Created Talent #{talent_id} ({best_person})")
|
||
|
||
# 2. 綁定 Speaker
|
||
cur.execute(
|
||
"""
|
||
INSERT INTO identity_bindings (talent_id, binding_type, binding_value, source, confidence)
|
||
VALUES (%s, 'speaker', %s, 'auto_cluster', 0.8)
|
||
ON CONFLICT (binding_type, binding_value) DO UPDATE SET talent_id = EXCLUDED.talent_id
|
||
""",
|
||
(talent_id, speaker),
|
||
)
|
||
print(f" ✅ Bound {speaker} -> {best_person}")
|
||
|
||
conn.commit()
|
||
cur.close()
|
||
conn.close()
|
||
except Exception as e:
|
||
print(f" ❌ DB Error: {e}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|