222 lines
7.3 KiB
Python
222 lines
7.3 KiB
Python
import os
|
|
import time
|
|
import argparse
|
|
import json
|
|
from datetime import datetime
|
|
|
|
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 ensure_dir(d):
|
|
if not os.path.exists(d):
|
|
os.makedirs(d, exist_ok=True)
|
|
|
|
|
|
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(embed_dir=EMBED_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(embed_dir=EMBED_DIR, use_db=True):
|
|
if use_db:
|
|
return load_embeddings_from_db()
|
|
return load_embeddings_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_from_crop(crop_img):
|
|
# pad + resize
|
|
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'live_query_{int(time.time())}.jpg')
|
|
cv2.imwrite(tmp, resized)
|
|
try:
|
|
emb = DeepFace.represent(tmp, model_name='Facenet', enforce_detection=True)
|
|
except Exception as e:
|
|
# fallback
|
|
try:
|
|
emb = DeepFace.represent(tmp, model_name='Facenet', enforce_detection=False, detector_backend='opencv')
|
|
except Exception as e2:
|
|
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 run_live(embed_dir=EMBED_DIR, threshold=0.25, camera_id=0, show_window=True, use_db=True):
|
|
gallery = load_embeddings(embed_dir, use_db=use_db)
|
|
print(f"Loaded {len(gallery)} enrolled embeddings")
|
|
cap = cv2.VideoCapture(camera_id)
|
|
if not cap.isOpened():
|
|
print("Cannot open camera")
|
|
return
|
|
face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
|
|
print("Live recognition started. Press 'q' to quit. Press 's' to save snapshot.")
|
|
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))
|
|
results = []
|
|
for (x,y,w,h) in faces:
|
|
crop = frame[y:y+h, x:x+w]
|
|
emb = compute_embedding_from_crop(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(frame, (x,y), (x+w, y+h), (0,255,0), 2)
|
|
label = name
|
|
if score is not None:
|
|
label = f"{name} {score:.2f}"
|
|
cv2.putText(frame, label, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2)
|
|
|
|
if show_window:
|
|
cv2.imshow('Live Recognize', frame)
|
|
key = cv2.waitKey(1) & 0xFF
|
|
if key == ord('q'):
|
|
break
|
|
if key == ord('s'):
|
|
ts = int(time.time())
|
|
fn = f'snapshot_{ts}.jpg'
|
|
cv2.imwrite(fn, frame)
|
|
print('Saved', fn)
|
|
else:
|
|
# headless: print results to console
|
|
if results:
|
|
for (_,_,_,_,name,score) in results:
|
|
print(f"Detected: {name} (score={score})")
|
|
# small sleep to reduce CPU
|
|
time.sleep(0.01)
|
|
cap.release()
|
|
if show_window:
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--embed-dir', default=EMBED_DIR)
|
|
parser.add_argument('--threshold', type=float, default=0.25)
|
|
parser.add_argument('--camera', type=int, default=0)
|
|
parser.add_argument('--no-window', action='store_true', help='Run headless (no GUI)')
|
|
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()
|
|
run_live(
|
|
embed_dir=args.embed_dir,
|
|
threshold=args.threshold,
|
|
camera_id=args.camera,
|
|
show_window=(not args.no_window),
|
|
use_db=args.use_db,
|
|
)
|