TKK_E32232028/detect_and_enroll.py

import os
import time
import argparse
import json
import tempfile
from datetime import datetime

import cv2
from deepface import DeepFace
from db import get_connection_from_env, ensure_table, save_embedding

CASCADE_PATH = "src/face.xml"
DEFAULT_EMBED_DIR = "embeddings"


def ensure_dir(d):
    if not os.path.exists(d):
        os.makedirs(d, exist_ok=True)


def compute_and_save_embedding(crop_img, name, embed_dir=DEFAULT_EMBED_DIR, save_crop=True, save_db=True):
    ensure_dir(embed_dir)
    # save temporary crop
    ts = int(time.time())
    crop_path = os.path.join(tempfile.gettempdir(), f"{name}_{ts}_crop.jpg")

    # Add padding around crop to help DeepFace detectors (and avoid tiny crops)
    h, w = crop_img.shape[:2]
    pad = int(0.25 * max(w, h))
    # create padded canvas
    padded = cv2.copyMakeBorder(crop_img, pad, pad, pad, pad, borderType=cv2.BORDER_CONSTANT, value=[0, 0, 0])

    # Resize to a reasonable size for the embedding model (helps detection)
    resize_side = 224
    resized = cv2.resize(padded, (resize_side, resize_side), interpolation=cv2.INTER_AREA)
    cv2.imwrite(crop_path, resized)

    # compute embedding with DeepFace
    print("Computing embedding (this may take a moment on first run)...")
    try:
        emb = DeepFace.represent(crop_path, model_name="Facenet", enforce_detection=True)
    except Exception as first_err:
        # fallback: try again without enforcement and using opencv backend
        print("First attempt failed (enforce_detection=True):", first_err)
        print("Retrying with enforce_detection=False and detector_backend='opencv'...")
        try:
            emb = DeepFace.represent(crop_path, model_name="Facenet", enforce_detection=False, detector_backend='opencv')
        except Exception as second_err:
            print("Second attempt failed:", second_err)
            return False
    # Guarantee embedding is a plain list (sometimes it's returned as numpy array)
    try:
        # if emb is dict with 'embedding' key, handle that case
        if isinstance(emb, dict) and 'embedding' in emb:
            vec = emb['embedding']
        else:
            vec = list(emb)
    except Exception:
        # fallback: try to convert to list directly
        vec = list(emb)
    # save embedding JSON
    out_json = os.path.join(embed_dir, f"{name}.json")
    meta = {
        "name": name,
        "timestamp": datetime.utcnow().isoformat() + 'Z',
        "embedding": vec,
    }
    with open(out_json, "w") as f:
        json.dump(meta, f)
    print("Saved embedding to", out_json)
    if save_db:
        try:
            conn = get_connection_from_env()
            ensure_table(conn)
            save_embedding(conn, name, json.dumps(vec))
            conn.close()
            print("Saved embedding to database for", name)
        except Exception as e:
            print("DB save failed:", e)
    # option: save crop permanently
    if save_crop:
        out_img = os.path.join(embed_dir, f"{name}.jpg")
        cv2.imwrite(out_img, crop_img)
        print("Saved face crop to", out_img)
    # cleanup temp file
    try:
        os.remove(crop_path)
    except Exception:
        pass
    return True


def run_camera_loop(show_window=False, embed_dir=DEFAULT_EMBED_DIR, save_crop=True, camera_id=0, save_db=True):
    cap = cv2.VideoCapture(camera_id)
    if not cap.isOpened():
        print("ERROR: Cannot open camera (id=%s)" % camera_id)
        return
    face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
    print("Camera opened. Waiting for faces... (press Ctrl+C to stop)")
    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("Failed to read frame from camera")
                break
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(60, 60))
            if len(faces) > 0:
                print(f"Detected {len(faces)} face(s)")
                # draw boxes for optional display
                for (x, y, w, h) in faces:
                    cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                if show_window:
                    cv2.imshow('Detect & Enroll', frame)
                    # small wait so window refreshes; but continue to enrollment prompt
                    cv2.waitKey(1)
                # take first face crop for enrollment
                x, y, w, h = faces[0]
                crop = frame[y:y+h, x:x+w]
                # Ask user for name
                print("Face ready to enroll.")
                name = input("Enter name for enrollment (empty to skip): ").strip()
                if name == "":
                    print("Skipped enrollment for this detection.")
                else:
                    ok = compute_and_save_embedding(
                        crop,
                        name,
                        embed_dir=embed_dir,
                        save_crop=save_crop,
                        save_db=save_db,
                    )
                    if ok:
                        print("Enrollment finished for", name)
                    else:
                        print("Enrollment FAILED for", name)
                # after enrollment (or skip) wait a little to avoid immediate re-detection of same frame
                print("Resuming detection in 2 seconds...")
                time.sleep(2)
            else:
                # nothing detected
                if show_window:
                    cv2.imshow('Detect & Enroll', frame)
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        break
                # tiny sleep to reduce CPU
                time.sleep(0.05)
    except KeyboardInterrupt:
        print("Stopped by user")
    finally:
        cap.release()
        if show_window:
            cv2.destroyAllWindows()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Detect faces from camera and enroll embeddings')
    parser.add_argument('--show', action='store_true', help='Show camera window (requires GUI/Qt)')
    parser.add_argument('--embed-dir', default=DEFAULT_EMBED_DIR, help='Directory to save embeddings and crops')
    parser.add_argument('--no-save-crop', action='store_true', help='Do not save face crop images')
    parser.add_argument('--camera', type=int, default=0, help='Camera device id')
    parser.add_argument('--no-db', action='store_true', help='Do not save embeddings to MySQL')
    args = parser.parse_args()

    run_camera_loop(
        show_window=args.show,
        embed_dir=args.embed_dir,
        save_crop=(not args.no_save_crop),
        camera_id=args.camera,
        save_db=(not args.no_db),
    )