TKK_E32232028/recognize.py

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)