github-actions
Deploy to Hugging Face
c794b6b
Raw
History Blame Contribute Delete
6.27 kB
import cv2
import numpy as np
from insightface.app import FaceAnalysis
import os
import json
from datetime import datetime, timezone
REAL_FACES_DB = "faces_db"
TEMP_DB_ROOT = "temp_face_database"
TEMP_EMB_ROOT = "temp_faces_db"
os.makedirs(TEMP_DB_ROOT, exist_ok=True)
os.makedirs(TEMP_EMB_ROOT, exist_ok=True)
def load_database():
db = {}
if os.path.exists(REAL_FACES_DB):
for file in os.listdir(REAL_FACES_DB):
if file.endswith(".npy"):
name = file.replace(".npy", "")
db[name] = np.load(os.path.join(REAL_FACES_DB, file))
if os.path.exists(TEMP_EMB_ROOT):
for file in os.listdir(TEMP_EMB_ROOT):
if file.endswith(".npy"):
name = file.replace(".npy", "")
db[name] = np.load(os.path.join(TEMP_EMB_ROOT, file))
return db
def cosine_similarity(a, b):
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
def get_next_unknown_id():
existing = [d for d in os.listdir(TEMP_DB_ROOT)
if os.path.isdir(os.path.join(TEMP_DB_ROOT, d)) and d.startswith("unknown_")]
if not existing:
return 1
ids = []
for d in existing:
try:
ids.append(int(d.split("_")[1]))
except (IndexError, ValueError):
pass
return max(ids) + 1 if ids else 1
def log_interaction(interaction_graph, a, b, timestamp):
if a == b:
return
interaction_graph.setdefault(a, {}).setdefault(b, []).append({
"timestamp": timestamp,
"camera": "cam_1"
})
interaction_graph.setdefault(b, {}).setdefault(a, []).append({
"timestamp": timestamp,
"camera": "cam_1"
})
def build_levels(root, graph):
level_1 = set(graph.get(root, {}).keys())
level_2 = set()
for p in level_1:
level_2.update(graph.get(p, {}).keys())
level_2 -= level_1
level_2.discard(root)
level_3 = set()
for p in level_2:
level_3.update(graph.get(p, {}).keys())
level_3 -= level_2
level_3 -= level_1
level_3.discard(root)
return level_1, level_2, level_3
def format_level(interaction_graph, level_set, via=None):
result = []
for person in level_set:
entry = {
"name": person,
"interactions": interaction_graph.get(person, {})
}
if via:
entry["interacted_via"] = via.get(person, "")
result.append(entry)
return result
def main():
app = FaceAnalysis(name='buffalo_l', allowed_modules=['detection', 'recognition'])
app.prepare(ctx_id=-1, det_size=(640, 640))
db = load_database()
root_person = input("Enter the name of the person to track: ").strip()
interaction_graph = {}
frame_id = 0
cap = cv2.VideoCapture(0)
if not cap.isOpened():
raise SystemExit("Could not open webcam.")
try:
while True:
ret, frame = cap.read()
if not ret:
break
frame_id += 1
faces = app.get(frame)
detected_people = []
bboxes = []
for face in faces:
x1, y1, x2, y2 = face.bbox.astype(int)
emb = face.embedding
best_match = "Unknown"
best_score = 0.0
for name, db_emb in db.items():
if db_emb.ndim == 1:
score = cosine_similarity(emb, db_emb)
else:
scores = [cosine_similarity(emb, view) for view in db_emb]
score = max(scores) if scores else 0.0
if score > best_score:
best_score = score
best_match = name
threshold = 0.35 if best_match.startswith("unknown") else 0.30
if best_score > threshold:
label = best_match
color = (0, 255, 0)
else:
new_id = get_next_unknown_id()
best_match = f"unknown_{new_id}"
db[best_match] = np.array([emb])
label = best_match
color = (0, 255, 0)
detected_people.append(best_match)
bboxes.append((x1, y1, x2, y2))
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
cv2.putText(frame, label, (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
frame_width = frame.shape[1]
proximity_threshold = 0.25 * frame_width
seen_pairs = set()
timestamp = datetime.now(timezone.utc).isoformat()
for i in range(len(detected_people)):
for j in range(i + 1, len(detected_people)):
a = detected_people[i]
b = detected_people[j]
if a == b:
continue
(x1a, y1a, x2a, y2a) = bboxes[i]
(x1b, y1b, x2b, y2b) = bboxes[j]
center_a = ((x1a + x2a) / 2, (y1a + y2a) / 2)
center_b = ((x1b + x2b) / 2, (y1b + y2b) / 2)
distance = np.linalg.norm(np.array(center_a) - np.array(center_b))
if distance < proximity_threshold:
pair = tuple(sorted([a, b]))
if pair not in seen_pairs:
log_interaction(interaction_graph, a, b, timestamp)
seen_pairs.add(pair)
cv2.imshow("Live Face Recognition", frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
finally:
cap.release()
cv2.destroyAllWindows()
level_1, level_2, level_3 = build_levels(root_person, interaction_graph)
output = {
"root_person": root_person,
"contacts": {
"level_1": format_level(interaction_graph, level_1),
"level_2": format_level(interaction_graph, level_2),
"level_3": format_level(interaction_graph, level_3)
}
}
with open("interaction_output.json", "w") as f:
json.dump(output, f, indent=2)
if __name__ == '__main__':
main()