99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
import os
|
|
import sys
|
|
import argparse
|
|
from deepface import DeepFace
|
|
import cv2
|
|
import json
|
|
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)
|
|
|
|
|
|
def crop_faces(image_path):
|
|
img = cv2.imread(image_path)
|
|
if img is None:
|
|
return []
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
|
|
faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
|
|
crops = []
|
|
for i, (x, y, w, h) in enumerate(faces):
|
|
crop = img[y:y+h, x:x+w]
|
|
crops.append(crop)
|
|
return crops
|
|
|
|
|
|
def enroll(name, image_path, save_db=True):
|
|
ensure_dir(EMBED_DIR)
|
|
crops = crop_faces(image_path)
|
|
if len(crops) == 0:
|
|
print("No faces found to enroll")
|
|
return
|
|
# Use first face crop for embedding
|
|
crop = crops[0]
|
|
tmp_path = f"/tmp/{name}_crop.jpg"
|
|
cv2.imwrite(tmp_path, crop)
|
|
print("Computing embedding for", name)
|
|
# DeepFace.represent returns embedding vector or dict depending on params
|
|
embedding = DeepFace.represent(tmp_path, model_name='Facenet', enforce_detection=True)
|
|
# Save embedding to JSON file
|
|
out_file = os.path.join(EMBED_DIR, f"{name}.json")
|
|
with open(out_file, "w") as f:
|
|
json.dump(embedding, f)
|
|
print("Saved embedding to", out_file)
|
|
|
|
# Optionally save to DB
|
|
if save_db:
|
|
try:
|
|
conn = get_connection_from_env()
|
|
ensure_table(conn)
|
|
# make embedding JSON-serializable: try to extract vector
|
|
vec = None
|
|
if isinstance(embedding, dict) and 'embedding' in embedding:
|
|
vec = embedding['embedding']
|
|
elif isinstance(embedding, list) and len(embedding) > 0 and isinstance(embedding[0], dict) and 'embedding' in embedding[0]:
|
|
vec = embedding[0]['embedding']
|
|
else:
|
|
try:
|
|
# try convert to list
|
|
vec = list(embedding)
|
|
except Exception:
|
|
vec = None
|
|
if vec is not None:
|
|
save_embedding(conn, name, json.dumps(vec))
|
|
print("Saved embedding to database for", name)
|
|
else:
|
|
print("Could not parse embedding to save to DB")
|
|
except Exception as e:
|
|
print("DB save failed:", e)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Enroll a face and save embedding")
|
|
parser.add_argument("name", help="Person name")
|
|
parser.add_argument("image", help="Path to image file")
|
|
parser.add_argument(
|
|
"--save-db",
|
|
action="store_true",
|
|
default=True,
|
|
help="Save embedding to MySQL (default: true)",
|
|
)
|
|
parser.add_argument(
|
|
"--no-db",
|
|
action="store_false",
|
|
dest="save_db",
|
|
help="Do not save embedding to MySQL",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
enroll(args.name, args.image, save_db=args.save_db)
|