197 lines
6.4 KiB
Python
197 lines
6.4 KiB
Python
import os
|
|
import argparse
|
|
import time
|
|
import json
|
|
import cv2
|
|
import numpy as np
|
|
from deepface import DeepFace
|
|
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_gallery_from_dir(embed_dir=EMBED_DIR):
|
|
gallery = {}
|
|
if not os.path.exists(embed_dir):
|
|
return gallery
|
|
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
|
|
gallery[name] = vec
|
|
return gallery
|
|
|
|
|
|
def load_gallery_from_db():
|
|
gallery = {}
|
|
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
|
|
gallery[name] = vec
|
|
return gallery
|
|
|
|
|
|
def load_gallery(embed_dir=EMBED_DIR, use_db=True):
|
|
if use_db:
|
|
return load_gallery_from_db()
|
|
return load_gallery_from_dir(embed_dir)
|
|
|
|
|
|
def parse_embedding(obj):
|
|
if isinstance(obj, dict):
|
|
if 'embedding' in obj:
|
|
return np.array(obj['embedding']).astype(float).reshape(-1)
|
|
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)
|
|
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')
|
|
|
|
|
|
def compute_embedding(crop_img):
|
|
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)
|
|
tmp = os.path.join('/tmp', f'gui_query_{int(time.time())}.jpg')
|
|
cv2.imwrite(tmp, resized)
|
|
try:
|
|
emb = DeepFace.represent(tmp, model_name='Facenet', enforce_detection=True)
|
|
except Exception as e:
|
|
try:
|
|
emb = DeepFace.represent(tmp, model_name='Facenet', enforce_detection=False, detector_backend='opencv')
|
|
except Exception:
|
|
return None
|
|
try:
|
|
vec = parse_embedding(emb)
|
|
except Exception:
|
|
return None
|
|
try:
|
|
os.remove(tmp)
|
|
except Exception:
|
|
pass
|
|
return vec
|
|
|
|
|
|
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 - float(np.dot(a_norm, b_norm))
|
|
|
|
|
|
def process_image(path, out_path=None, threshold=0.25, show_window=True, use_db=True):
|
|
img = cv2.imread(path)
|
|
if img is None:
|
|
print('Could not open image', path)
|
|
return False
|
|
gallery = load_gallery(use_db=use_db)
|
|
print(f'Loaded {len(gallery)} enrolled identities')
|
|
face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30,30))
|
|
print(f'Found {len(faces)} face(s) in image')
|
|
results = []
|
|
for (x,y,w,h) in faces:
|
|
crop = img[y:y+h, x:x+w]
|
|
emb = compute_embedding(crop)
|
|
name = 'Unknown'
|
|
score = None
|
|
if emb is not None and len(gallery) > 0:
|
|
best = (None, float('inf'))
|
|
for gname, gvec in gallery.items():
|
|
d = cosine_distance(emb, gvec)
|
|
if d < best[1]:
|
|
best = (gname, d)
|
|
if best[0] is not None and is_recognized_match(best[1], threshold):
|
|
name = best[0]
|
|
score = best[1]
|
|
else:
|
|
name = 'Unknown'
|
|
score = best[1]
|
|
results.append((x,y,w,h,name,score))
|
|
|
|
# draw
|
|
for (x,y,w,h,name,score) in results:
|
|
cv2.rectangle(img, (x,y), (x+w, y+h), (0,255,0), 2)
|
|
label = name
|
|
if score is not None:
|
|
label = f"{name} {score:.2f}"
|
|
cv2.putText(img, label, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2)
|
|
|
|
if out_path:
|
|
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
|
cv2.imwrite(out_path, img)
|
|
print('Saved annotated image to', out_path)
|
|
if show_window:
|
|
cv2.imshow('Recognition', img)
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|
|
return True
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('image', help='Path to image')
|
|
parser.add_argument('--out', help='Path to save annotated image')
|
|
parser.add_argument('--threshold', type=float, default=0.25)
|
|
parser.add_argument('--no-window', action='store_true')
|
|
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()
|
|
process_image(
|
|
args.image,
|
|
out_path=args.out,
|
|
threshold=args.threshold,
|
|
show_window=(not args.no_window),
|
|
use_db=args.use_db,
|
|
)
|