105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
import os
|
|
import time
|
|
import json
|
|
import tempfile
|
|
from datetime import datetime
|
|
|
|
import cv2
|
|
from deepface import DeepFace
|
|
from db import get_connection_from_env, ensure_table, save_embedding
|
|
|
|
CASCADE_PATH = "src/face.xml"
|
|
EMBED_DIR = "embeddings"
|
|
|
|
|
|
def ensure_dir(d):
|
|
if not os.path.exists(d):
|
|
os.makedirs(d, exist_ok=True)
|
|
|
|
|
|
def compute_and_save(crop_img, name, save_db=True):
|
|
ensure_dir(EMBED_DIR)
|
|
ts = int(time.time())
|
|
tmp = os.path.join(tempfile.gettempdir(), f"{name}_{ts}_crop.jpg")
|
|
# add padding and resize as in detect_and_enroll
|
|
h, w = crop_img.shape[:2]
|
|
pad = int(0.25 * max(w, h))
|
|
padded = cv2.copyMakeBorder(crop_img, pad, pad, pad, pad, borderType=cv2.BORDER_CONSTANT, value=[0,0,0])
|
|
resized = cv2.resize(padded, (224,224), interpolation=cv2.INTER_AREA)
|
|
cv2.imwrite(tmp, resized)
|
|
try:
|
|
emb = DeepFace.represent(tmp, model_name='Facenet', enforce_detection=True)
|
|
except Exception as e:
|
|
print('First attempt failed:', e)
|
|
try:
|
|
emb = DeepFace.represent(tmp, model_name='Facenet', enforce_detection=False, detector_backend='opencv')
|
|
except Exception as e2:
|
|
print('Second attempt failed:', e2)
|
|
return False
|
|
# normalize embedding to list
|
|
if isinstance(emb, dict) and 'embedding' in emb:
|
|
vec = emb['embedding']
|
|
elif isinstance(emb, list) and len(emb) > 0 and isinstance(emb[0], dict) and 'embedding' in emb[0]:
|
|
vec = emb[0]['embedding']
|
|
else:
|
|
vec = list(emb)
|
|
out_json = os.path.join(EMBED_DIR, f"{name}.json")
|
|
meta = {"name":name, "timestamp":datetime.utcnow().isoformat()+"Z", "embedding":vec}
|
|
with open(out_json, 'w') as f:
|
|
json.dump(meta, f)
|
|
out_img = os.path.join(EMBED_DIR, f"{name}.jpg")
|
|
cv2.imwrite(out_img, resized)
|
|
print('Saved embedding and crop for', name)
|
|
if save_db:
|
|
try:
|
|
conn = get_connection_from_env()
|
|
ensure_table(conn)
|
|
save_embedding(conn, name, json.dumps(vec))
|
|
conn.close()
|
|
print('Saved embedding to database for', name)
|
|
except Exception as e:
|
|
print('DB save failed:', e)
|
|
try:
|
|
os.remove(tmp)
|
|
except Exception:
|
|
pass
|
|
return True
|
|
|
|
|
|
def main(camera_id=0):
|
|
cap = cv2.VideoCapture(camera_id)
|
|
if not cap.isOpened():
|
|
print('Cannot open camera')
|
|
return
|
|
face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
|
|
print("Press 'c' to capture current detected face and enroll; 'q' to quit")
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
print('Failed to read frame')
|
|
break
|
|
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(60,60))
|
|
for (x,y,w,h) in faces:
|
|
cv2.rectangle(frame, (x,y), (x+w, y+h), (0,255,0), 2)
|
|
cv2.imshow('Manual Detect & Enroll', frame)
|
|
key = cv2.waitKey(1) & 0xFF
|
|
if key == ord('q'):
|
|
break
|
|
if key == ord('c'):
|
|
if len(faces)==0:
|
|
print('No faces detected to capture')
|
|
continue
|
|
x,y,w,h = faces[0]
|
|
crop = frame[y:y+h, x:x+w]
|
|
name = input('Enter name for enrollment: ').strip()
|
|
if name == '':
|
|
print('Empty name, skipping')
|
|
else:
|
|
compute_and_save(crop, name)
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|