Spaces:
Build error
Build error
File size: 4,683 Bytes
f10370d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
# run.py
# run.py (Add this to the very top, before imports)
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # Forces TensorFlow to use CPU only
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" # Suppress warnings
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" # For GPU, but harmless on CPU
os.environ["NUMEXPR_NUM_THREADS"] = "1" # Limit threads for NumPy/NumExpr
import argparse
import cv2
import sys
import numpy as np
import tensorflow as tf
# Import your core logic modules
from src.utils import load_embeddings
from src.detect import detect_face
from src.embed import get_embedding
from src.recognize import recognize_face_by_embedding
from src.register import register_new_user
# Define the logic to process a captured frame for registration or recognition
def process_frame(frame, mode, detector_key, data_store):
# 1. Detection
detected_faces = detect_face(frame, detector_type=detector_key)
if detected_faces:
# Focus on the largest face for processing
main_face = max(detected_faces, key=lambda x: x['box'][2] * x['box'][3])
x, y, w, h = main_face['box']
face_img = frame[y:y+h, x:x+w]
# Draw bounding box on the displayed frame
cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 255), 2)
if mode == 'register':
# Ask for name when face is captured
name = input("Enter name to register: ")
# 2. Registration Logic
if register_new_user(face_img, name):
return f"✅ Registration successful for {name}! Please run recognize mode next."
else:
return "❌ Registration failed. Could not generate embedding."
elif mode in ['recognize', 'identify']:
# 3. Recognition Logic
embedding = get_embedding(face_img)
if embedding is None:
return "⚠️ Could not generate embedding for recognition."
name, distance = recognize_face_by_embedding(
embedding,
data_store['embeddings'],
data_store['names']
)
if name != "Unknown":
cv2.putText(frame, name, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
return f"✅ Recognized: {name} (Dist: {distance:.2f})"
else:
cv2.putText(frame, "Unknown", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
return f"⚠️ Unknown Person (Dist: {distance:.2f})"
return "No face detected in the frame."
# --- Main CLI Function ---
def main():
parser = argparse.ArgumentParser(description="Smart Office Face Recognition System CLI.")
parser.add_argument('--mode', required=True, choices=['register', 'recognize', 'identify'], help="Operation mode (register or identify/recognize).")
parser.add_argument('--detector', required=True, choices=['cnn', 'classical'], help="Detector backend: cnn (MTCNN) or classical (Haar Cascade).")
args = parser.parse_args()
data_store = load_embeddings() # Initial load of persistent data
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("Error: Could not open webcam.")
sys.exit(1)
print(f"\n--- Running in {args.mode.upper()} mode with {args.detector.upper()} detector. ---")
print("Press 'c' to CAPTURE/PROCESS a face or 'q' to QUIT.\n")
while True:
ret, frame = cap.read()
if not ret:
break
# Display instructions on the frame
cv2.putText(frame, f"MODE: {args.mode.upper()} | DETECTOR: {args.detector.upper()}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
cv2.putText(frame, "Press 'c' to capture.", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
cv2.imshow('Smart Office Face Recognition CLI', frame)
key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
break
elif key == ord('c'):
print(f"Capture command received. Processing frame in {args.mode} mode...")
# Process the captured frame
result = process_frame(frame, args.mode, args.detector, data_store)
print(result)
# FIX: Pylance Scope Fix - If registered, reload the local data_store variable for immediate recognition.
if "Registration successful" in result:
data_store = load_embeddings()
print("Data store reloaded for new user recognition.")
cap.release()
cv2.destroyAllWindows()
if __name__ == '__main__':
main() |