ICEYOU/face_enrollment.py

"""ICEYOU - Authorized User Face Enrollment

Captures multiple webcam photos of the authorized user from various angles,
saves them under faces/owner/, and trains an OpenCV LBPH face-recognition model.

Usage:
    python face_enrollment.py
        (optionally append --name <label> to enroll secondary authorized users)

Controls in the live preview window:
    SPACE  - Capture the current frame (only if a face is detected)
    R      - Restart enrollment
    Q/ESC  - Quit
"""

import sys
import os
import json
import argparse
from pathlib import Path
from typing import List, Optional

import cv2


REQUIRED_POSES = [
    "Look straight at the camera",
    "Slowly turn your head to the LEFT",
    "Slowly turn your head to the RIGHT",
    "Tilt your head slightly UP",
    "Tilt your head slightly DOWN",
]

MIN_CAPTURES = 3
TARGET_CAPTURES = len(REQUIRED_POSES)

FACES_DIR = Path("faces")
MODEL_FILE = FACES_DIR / "model.yml"
LABELS_FILE = FACES_DIR / "labels.json"


def load_face_detectors():
    """Load both frontal and profile Haar cascades for better coverage."""
    base = Path(cv2.data.haarcascades)
    frontal_path = base / "haarcascade_frontalface_default.xml"
    alt_path = base / "haarcascade_frontalface_alt2.xml"
    profile_path = base / "haarcascade_profileface.xml"
    if not frontal_path.exists():
        print(f"ERROR: Haar cascade not found at {frontal_path}")
        sys.exit(1)
    detectors = {
        "frontal": cv2.CascadeClassifier(str(frontal_path)),
        "alt2": cv2.CascadeClassifier(str(alt_path)) if alt_path.exists() else None,
        "profile": cv2.CascadeClassifier(str(profile_path)) if profile_path.exists() else None,
    }
    return detectors


def detect_face(detectors: dict, gray):
    """Run multiple cascades and return list of face rects.
    Tries frontal first, then alt2, then profile (also mirrored for the other side).
    """
    # Improve contrast for low-light conditions
    eq = cv2.equalizeHist(gray)

    params = dict(scaleFactor=1.1, minNeighbors=4, minSize=(60, 60))

    faces = detectors["frontal"].detectMultiScale(eq, **params)
    if len(faces) == 0 and detectors["alt2"] is not None:
        faces = detectors["alt2"].detectMultiScale(eq, **params)
    if len(faces) == 0 and detectors["profile"] is not None:
        faces = detectors["profile"].detectMultiScale(eq, **params)
        if len(faces) == 0:
            # Mirror to detect the opposite-facing profile
            flipped = cv2.flip(eq, 1)
            faces = detectors["profile"].detectMultiScale(flipped, **params)
            if len(faces) > 0:
                w = eq.shape[1]
                faces = [(w - x - fw, y, fw, fh) for (x, y, fw, fh) in faces]
    return faces


def run_capture_session(user_label: str, user_dir: Path, camera_index: int = 0) -> List[Path]:
    """Open camera, prompt for poses, save captures. Returns list of saved paths."""
    user_dir.mkdir(parents=True, exist_ok=True)

    print(f"Opening camera index {camera_index}...")
    cap = cv2.VideoCapture(camera_index, cv2.CAP_DSHOW)
    if not cap.isOpened():
        print(f"ERROR: Could not open webcam at index {camera_index}.")
        print("Try a different --camera value (0, 1, 2, ...).")
        sys.exit(1)

    detectors = load_face_detectors()
    captures: List[Path] = []
    pose_idx = 0

    # Capture trigger flag (set by mouse click, auto-capture, or SPACE key)
    capture_requested = [False]

    def on_mouse(event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            capture_requested[0] = True

    print("\n=== ICEYOU Face Enrollment ===")
    print(f"User label: {user_label}")
    print(f"Target captures: {TARGET_CAPTURES} (minimum {MIN_CAPTURES})")
    print("Controls:")
    print("  - CLICK the preview window  = capture (most reliable)")
    print("  - SPACE / ENTER / C         = capture (window must have focus)")
    print("  - Auto-capture: hold a steady face for ~2s = capture automatically")
    print("  - R = restart, Q/ESC = quit\n")

    win_name = f"ICEYOU Enrollment - {user_label}"
    cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
    cv2.setMouseCallback(win_name, on_mouse)

    import time
    steady_face_start = None  # timestamp when a face first became visible & still

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("Camera read failed.")
                break

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = detect_face(detectors, gray)
            face_detected = len(faces) > 0

            display = frame.copy()
            for (x, y, w, h) in faces:
                cv2.rectangle(display, (x, y), (x + w, y + h), (0, 255, 0), 2)

            done = pose_idx >= TARGET_CAPTURES
            instruction = "ENROLLMENT COMPLETE - press Q to finish" if done else REQUIRED_POSES[pose_idx]

            # Auto-capture: 2-second steady face countdown
            auto_countdown_remaining = None
            if not done and face_detected:
                if steady_face_start is None:
                    steady_face_start = time.time()
                elapsed = time.time() - steady_face_start
                auto_countdown_remaining = max(0.0, 2.0 - elapsed)
                if elapsed >= 2.0:
                    capture_requested[0] = True
                    steady_face_start = None
            else:
                steady_face_start = None

            overlay = display.copy()
            cv2.rectangle(overlay, (0, 0), (display.shape[1], 110), (0, 0, 0), -1)
            cv2.addWeighted(overlay, 0.6, display, 0.4, 0, display)

            cv2.putText(display, f"[{len(captures)}/{TARGET_CAPTURES}]  {instruction}",
                        (10, 35), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
            if face_detected:
                status = f"Face detected ({len(faces)}) - CLICK or SPACE"
                color = (0, 255, 0)
            else:
                status = "No face auto-detected - CLICK / SPACE = manual center crop"
                color = (0, 200, 255)
            cv2.putText(display, status, (10, 70),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 2)

            if auto_countdown_remaining is not None and auto_countdown_remaining > 0:
                cv2.putText(display, f"Auto-capture in {auto_countdown_remaining:.1f}s",
                            (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 0), 2)

            # Draw center-crop guide when no face detected (helps user frame themselves)
            if not face_detected:
                h, w = display.shape[:2]
                side = min(h, w) * 6 // 10
                cx0, cy0 = (w - side) // 2, (h - side) // 2
                cv2.rectangle(display, (cx0, cy0), (cx0 + side, cy0 + side), (0, 200, 255), 2)

            cv2.imshow(win_name, display)
            key = cv2.waitKey(30) & 0xFF

            # Capture keys: SPACE (32), ENTER (13), C (99)
            capture_keys = {32, 13, ord("c"), ord("C")}
            quit_keys = {ord("q"), ord("Q"), 27}
            restart_keys = {ord("r"), ord("R")}

            if key in quit_keys:
                break
            elif key in restart_keys:
                for p in captures:
                    try:
                        p.unlink()
                    except Exception:
                        pass
                captures.clear()
                pose_idx = 0
                print("Restarting enrollment...")
            elif key in capture_keys:
                capture_requested[0] = True
            elif key not in (255, 0):
                print(f"  (unrecognized key code: {key})")

            # Handle the capture request from any source (mouse, key, auto-capture)
            if capture_requested[0]:
                capture_requested[0] = False
                if done:
                    print("Enrollment already complete - press Q to finish.")
                    continue
                if face_detected:
                    (x, y, w, h) = max(faces, key=lambda f: f[2] * f[3])
                    m = int(0.1 * max(w, h))
                    x0 = max(x - m, 0)
                    y0 = max(y - m, 0)
                    x1 = min(x + w + m, gray.shape[1])
                    y1 = min(y + h + m, gray.shape[0])
                    face_crop = gray[y0:y1, x0:x1]
                    crop_mode = "detected"
                else:
                    h, w = gray.shape
                    side = min(h, w) * 6 // 10
                    cy0 = (h - side) // 2
                    cx0 = (w - side) // 2
                    face_crop = gray[cy0:cy0 + side, cx0:cx0 + side]
                    crop_mode = "manual-center"
                face_resized = cv2.resize(face_crop, (200, 200))

                filename = user_dir / f"capture_{pose_idx + 1:02d}.png"
                cv2.imwrite(str(filename), face_resized)
                captures.append(filename)
                print(f"  Captured ({crop_mode}): {filename.name} ({REQUIRED_POSES[pose_idx]})")
                pose_idx += 1
                steady_face_start = None  # reset auto-capture so it waits again
    finally:
        cap.release()
        cv2.destroyAllWindows()

    return captures


def train_model(faces_root: Path) -> bool:
    """Train an LBPH model on all images under faces_root/<label>/*.png.
    Saves to faces/model.yml and labels.json (label_id -> label_name).
    Requires opencv-contrib-python (cv2.face module).
    """
    try:
        recognizer = cv2.face.LBPHFaceRecognizer_create()
    except AttributeError:
        print("\nWARNING: cv2.face is unavailable. Install opencv-contrib-python:")
        print("    pip install --upgrade --force-reinstall opencv-contrib-python")
        print("Photos are saved but recognition model NOT trained.")
        return False

    samples = []
    sample_labels = []
    label_map = {}

    user_dirs = sorted([d for d in faces_root.iterdir() if d.is_dir()])
    if not user_dirs:
        print("No user directories to train from.")
        return False

    for label_id, user_dir in enumerate(user_dirs):
        label_map[label_id] = user_dir.name
        images = sorted(user_dir.glob("*.png")) + sorted(user_dir.glob("*.jpg"))
        for img_path in images:
            img = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE)
            if img is None:
                continue
            samples.append(img)
            sample_labels.append(label_id)

    if not samples:
        print("No samples found to train.")
        return False

    import numpy as np
    recognizer.train(samples, np.array(sample_labels))
    MODEL_FILE.parent.mkdir(parents=True, exist_ok=True)
    recognizer.save(str(MODEL_FILE))
    with open(LABELS_FILE, "w", encoding="utf-8") as f:
        json.dump({str(k): v for k, v in label_map.items()}, f, indent=2)

    print(f"\nTrained on {len(samples)} samples across {len(label_map)} user(s).")
    print(f"Model saved: {MODEL_FILE}")
    print(f"Labels saved: {LABELS_FILE}")
    return True


def _default_camera_index() -> int:
    """Read camera_device_index from config.json if present, else 0."""
    cfg_path = Path("config.json")
    if cfg_path.exists():
        try:
            with open(cfg_path, "r", encoding="utf-8") as f:
                return int(json.load(f).get("camera_device_index", 0))
        except Exception:
            pass
    return 0


def list_available_cameras(max_index: int = 5) -> List[int]:
    """Probe camera indices 0..max_index-1 and return those that open successfully."""
    available = []
    for i in range(max_index):
        cap = cv2.VideoCapture(i, cv2.CAP_DSHOW)
        if cap.isOpened():
            available.append(i)
            cap.release()
    return available


def main():
    parser = argparse.ArgumentParser(description="ICEYOU face enrollment")
    parser.add_argument("--name", default="owner",
                        help="Label for this user (default: owner)")
    parser.add_argument("--camera", "-c", type=int, default=None,
                        help="Camera device index (default: from config.json or 0)")
    parser.add_argument("--list-cameras", action="store_true",
                        help="Probe and list available camera indices, then exit")
    args = parser.parse_args()

    if args.list_cameras:
        cams = list_available_cameras()
        if cams:
            print(f"Available camera indices: {cams}")
        else:
            print("No cameras detected.")
        return

    camera_index = args.camera if args.camera is not None else _default_camera_index()
    user_label = args.name.strip().replace(" ", "_") or "owner"
    user_dir = FACES_DIR / user_label

    if user_dir.exists() and any(user_dir.iterdir()):
        resp = input(f"Existing enrollment for '{user_label}' found. Overwrite? [y/N]: ").strip().lower()
        if resp != "y":
            print("Aborted.")
            return
        for p in user_dir.glob("*"):
            try:
                p.unlink()
            except Exception:
                pass

    captures = run_capture_session(user_label, user_dir, camera_index=camera_index)

    if len(captures) < MIN_CAPTURES:
        print(f"\nOnly {len(captures)} photo(s) captured. Need at least {MIN_CAPTURES}.")
        print("Enrollment not finalized.")
        return

    print(f"\n{len(captures)} photo(s) saved to {user_dir}")
    print("Training recognition model...")
    train_model(FACES_DIR)

    print("\nEnrollment complete.")
    print("To enable recognition, set in config.json:")
    print('    "face_recognition_enabled": true')


if __name__ == "__main__":
    main()