import cv2 import insightface from insightface.app import FaceAnalysis import numpy as np import os import time from register_face import register_face REAL_FACES_DB = "faces_db" TEMP_DB_ROOT = "temp_face_database" TEMP_EMB_ROOT = "temp_faces_db" # Ensure temp directories exist os.makedirs(TEMP_DB_ROOT, exist_ok=True) os.makedirs(TEMP_EMB_ROOT, exist_ok=True) # Global database db = {} def load_database(): global db db = {} # Load real faces 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)) # Load temp faces 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)) print(f"Loaded {len(db)} faces from database") def cosine_similarity(a, b): return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)) def recognize_face(face_embedding): best_match = "Unknown" best_score = 0.0 for name, db_emb in db.items(): if db_emb.ndim == 1: score = cosine_similarity(face_embedding, db_emb) else: scores = [cosine_similarity(face_embedding, 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 # Dynamic Thresholding is handled in the main loop now, but we return the best match here # We can just return the match and let the loop decide based on the name return best_match, best_score def get_next_unknown_id(): # Find all folders starting with "unknown_" 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 # Extract numbers 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 main(input_path: str = "input_video.mp4", output_path: str = "output_recognized.mp4"): app = FaceAnalysis(name="buffalo_l") app.prepare(ctx_id=0, det_size=(640, 640)) load_database() cap = cv2.VideoCapture(input_path) if not cap.isOpened(): raise Exception(f"Error opening video file {input_path}") fourcc = cv2.VideoWriter_fourcc(*'mp4v') fps = cap.get(cv2.CAP_PROP_FPS) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) out = cv2.VideoWriter(output_path, fourcc, fps, (w, h)) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) print(f"Processing {frame_count} frames...") try: while True: ret, frame = cap.read() if not ret: break faces = app.get(frame) for face in faces: bbox = face.bbox.astype(int) x1, y1, x2, y2 = bbox best_match, score = recognize_face(face.normed_embedding) if best_match.startswith("unknown"): threshold = 0.35 else: threshold = 0.30 if score > threshold: name = best_match else: name = "Unknown" if name != "Unknown": if name in db: current_db_emb = db[name] if current_db_emb.ndim == 1: current_db_emb = np.expand_dims(current_db_emb, axis=0) updated_emb = np.vstack([current_db_emb, face.normed_embedding]) if len(updated_emb) > 50: updated_emb = updated_emb[-50:] db[name] = updated_emb if name.startswith("unknown"): try: npy_path = os.path.join(TEMP_EMB_ROOT, f"{name}.npy") np.save(npy_path, updated_emb) except OSError as e: print(f"Failed to update persistent DB for {name}: {e}") if name == "Unknown": h, w, _ = frame.shape pad_w = int((x2 - x1) * 0.25) pad_h = int((y2 - y1) * 0.25) crop_x1 = max(0, x1 - pad_w) crop_y1 = max(0, y1 - pad_h) crop_x2 = min(w, x2 + pad_w) crop_y2 = min(h, y2 + pad_h) face_crop = frame[crop_y1:crop_y2, crop_x1:crop_x2] crop_faces = app.get(face_crop) if len(crop_faces) == 0: print("Skipping unknown registration: No face detected in crop (False Positive).") continue crop_emb = crop_faces[0].embedding check_match, check_score = recognize_face(crop_emb) check_threshold = 0.35 if check_match.startswith("unknown") else 0.30 if check_score > check_threshold: print(f"Crop matched {check_match} ({check_score:.2f})! Updating instead of registering new.") if check_match in db: current_db_emb = db[check_match] if current_db_emb.ndim == 1: current_db_emb = np.expand_dims(current_db_emb, axis=0) updated_emb = np.vstack([current_db_emb, crop_emb]) if len(updated_emb) > 50: updated_emb = updated_emb[-50:] db[check_match] = updated_emb if check_match.startswith("unknown"): try: np.save(os.path.join(TEMP_EMB_ROOT, f"{check_match}.npy"), updated_emb) except OSError: pass cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) text_y = y2 + 25 cv2.putText(frame, f"{check_match} ({check_score:.2f})", (x1, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2) continue print(f"New unknown face detected (score: {score:.2f})") new_id = get_next_unknown_id() new_name = f"unknown_{new_id}" temp_img_path = f"{new_name}.jpg" cv2.imwrite(temp_img_path, face_crop) try: print(f"Registering new person: {new_name}") new_embeddings = register_face(new_name, temp_img_path, TEMP_DB_ROOT, TEMP_EMB_ROOT, known_embedding=face.normed_embedding) db[new_name] = new_embeddings name = new_name except Exception as e: print(f"Failed to register unknown face: {e}") try: import shutil failed_folder = os.path.join(TEMP_DB_ROOT, new_name) if os.path.exists(failed_folder): shutil.rmtree(failed_folder) if new_name in db: del db[new_name] except Exception as cleanup_e: print(f"Cleanup failed: {cleanup_e}") if os.path.exists(temp_img_path): os.remove(temp_img_path) cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) text_y = y2 + 25 cv2.putText(frame, f"{name} ({score:.2f})", (x1, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2) out.write(frame) finally: cap.release() out.release() print(f"Face recognition video saved as {output_path}") if __name__ == '__main__': main()