192 lines
6.5 KiB
Python
192 lines
6.5 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import argparse
|
|
import numpy as np
|
|
from deepface import DeepFace
|
|
import cv2
|
|
from db import get_connection_from_env, ensure_table, load_all_embeddings
|
|
from face_match import is_recognized_match
|
|
|
|
CASCADE_PATH = "src/face.xml"
|
|
EMBED_DIR = "embeddings"
|
|
|
|
|
|
def normalize_embedding(raw):
|
|
emb = raw
|
|
if isinstance(emb, dict):
|
|
if 'embedding' in emb:
|
|
emb = emb['embedding']
|
|
else:
|
|
emb = emb
|
|
if isinstance(emb, list) and len(emb) > 0 and isinstance(emb[0], dict) and 'embedding' in emb[0]:
|
|
emb = emb[0]['embedding']
|
|
if isinstance(emb, dict) and 'embedding' in emb:
|
|
emb = emb['embedding']
|
|
try:
|
|
return np.array(emb).astype(float).reshape(-1)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def load_embeddings_from_dir():
|
|
embs = {}
|
|
if not os.path.exists(EMBED_DIR):
|
|
return embs
|
|
for fn in os.listdir(EMBED_DIR):
|
|
if fn.endswith('.json'):
|
|
name = fn[:-5]
|
|
with open(os.path.join(EMBED_DIR, fn), 'r') as f:
|
|
data = json.load(f)
|
|
vec = normalize_embedding(data)
|
|
if vec is None:
|
|
print(f"Warning: could not parse embedding for {name}, skipping")
|
|
continue
|
|
embs[name] = vec
|
|
return embs
|
|
|
|
|
|
def load_embeddings_from_db():
|
|
embs = {}
|
|
conn = get_connection_from_env()
|
|
ensure_table(conn)
|
|
rows = load_all_embeddings(conn)
|
|
conn.close()
|
|
for row in rows:
|
|
name = row.get('name')
|
|
emb_raw = row.get('embedding')
|
|
if isinstance(emb_raw, str):
|
|
try:
|
|
emb_raw = json.loads(emb_raw)
|
|
except Exception:
|
|
pass
|
|
vec = normalize_embedding(emb_raw)
|
|
if vec is None:
|
|
print(f"Warning: could not parse embedding for {name}, skipping")
|
|
continue
|
|
embs[name] = vec
|
|
return embs
|
|
|
|
|
|
def load_embeddings(use_db=True):
|
|
if use_db:
|
|
return load_embeddings_from_db()
|
|
return load_embeddings_from_dir()
|
|
|
|
|
|
def crop_first_face(image_path):
|
|
img = cv2.imread(image_path)
|
|
if img is None:
|
|
return None
|
|
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))
|
|
if len(faces)==0:
|
|
return None
|
|
x,y,w,h = faces[0]
|
|
# pad and resize similar to enrollment to help detectors
|
|
crop = img[y:y+h, x:x+w]
|
|
pad = int(0.25 * max(w, h))
|
|
padded = cv2.copyMakeBorder(crop, pad, pad, pad, pad, borderType=cv2.BORDER_CONSTANT, value=[0,0,0])
|
|
resized = cv2.resize(padded, (224,224), interpolation=cv2.INTER_AREA)
|
|
return resized
|
|
|
|
|
|
def recognize(image_path, use_db=True):
|
|
embs = load_embeddings(use_db=use_db)
|
|
if not embs:
|
|
if use_db:
|
|
print("No enrolled embeddings found in DB. Enroll first.")
|
|
else:
|
|
print("No enrolled embeddings found. Run enroll_faces.py first.")
|
|
return
|
|
crop = crop_first_face(image_path)
|
|
if crop is None:
|
|
print("No face detected in image")
|
|
return
|
|
tmp = "/tmp/query_crop.jpg"
|
|
cv2.imwrite(tmp, crop)
|
|
|
|
# Try compute embedding with enforcement, fallback to no-enforce + opencv backend
|
|
try:
|
|
query_emb = DeepFace.represent(tmp, model_name='Facenet', enforce_detection=True)
|
|
except Exception as e:
|
|
print("Primary represent failed:", e)
|
|
try:
|
|
query_emb = DeepFace.represent(tmp, model_name='Facenet', enforce_detection=False, detector_backend='opencv')
|
|
print("Fallback represent succeeded (enforce_detection=False)")
|
|
except Exception as e2:
|
|
print("Failed to compute embedding for query:", e2)
|
|
return
|
|
|
|
# normalize query embedding to numeric vector
|
|
def parse_embedding(obj):
|
|
# obj may be list of numbers, numpy array, dict with 'embedding', or list with dict
|
|
if isinstance(obj, dict):
|
|
if 'embedding' in obj:
|
|
return np.array(obj['embedding']).astype(float).reshape(-1)
|
|
# maybe dict keyed differently
|
|
# try to find embedding-like key
|
|
for k in obj:
|
|
if isinstance(obj[k], (list, tuple, np.ndarray)):
|
|
try:
|
|
return np.array(obj[k]).astype(float).reshape(-1)
|
|
except Exception:
|
|
continue
|
|
if isinstance(obj, list):
|
|
if len(obj) > 0 and isinstance(obj[0], dict) and 'embedding' in obj[0]:
|
|
return np.array(obj[0]['embedding']).astype(float).reshape(-1)
|
|
# list of numbers
|
|
try:
|
|
return np.array(obj).astype(float).reshape(-1)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
return np.array(obj).astype(float).reshape(-1)
|
|
except Exception:
|
|
raise ValueError("Cannot parse embedding")
|
|
|
|
try:
|
|
q = parse_embedding(query_emb)
|
|
except Exception as e:
|
|
print("Could not normalize query embedding:", e)
|
|
return
|
|
|
|
# Use cosine distance for comparison (1 - cosine_similarity)
|
|
def cosine_distance(a, b):
|
|
a_norm = a / (np.linalg.norm(a) + 1e-10)
|
|
b_norm = b / (np.linalg.norm(b) + 1e-10)
|
|
return 1.0 - np.dot(a_norm, b_norm)
|
|
|
|
best = (None, float('inf'))
|
|
second_best = (None, float('inf'))
|
|
for name, emb in embs.items():
|
|
try:
|
|
e = np.array(emb).astype(float).reshape(-1)
|
|
except Exception:
|
|
print(f"Skipping {name}: cannot parse gallery embedding")
|
|
continue
|
|
dist = float(cosine_distance(q, e))
|
|
if dist < best[1]:
|
|
second_best = best
|
|
best = (name, dist)
|
|
elif dist < second_best[1]:
|
|
second_best = (name, dist)
|
|
|
|
THRESHOLD = 0.25
|
|
name, score = best
|
|
if name is None:
|
|
print("No match found")
|
|
else:
|
|
match = is_recognized_match(score, THRESHOLD, second_score=second_best[1])
|
|
print(f"Best match: {name} (cosine distance={score:.4f}) -> {'MATCH' if match else 'NO MATCH'} (threshold={THRESHOLD})")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("image", help="Path to image file")
|
|
parser.add_argument("--use-db", action="store_true", default=True, help="Load embeddings from MySQL (default)")
|
|
parser.add_argument("--use-dir", action="store_false", dest="use_db", help="Load embeddings from embeddings/ directory")
|
|
args = parser.parse_args()
|
|
recognize(args.image, use_db=args.use_db)
|