import os import shutil import cv2 import numpy as np from insightface.app import FaceAnalysis # Lazy-load app to avoid multiple allocations when imported _app = None def get_app(): """Return FaceAnalysis using the same model pack as FaceMatcher (buffalo_sc by default).""" global _app if _app is None: model_pack = os.getenv("FACE_MODEL_PACK", "buffalo_sc") model_root = os.getenv("FACE_MODEL_ROOT", "/app/model_cache") _app = FaceAnalysis( name=model_pack, root=model_root, providers=["CPUExecutionProvider"], allowed_modules=["detection", "recognition"], ) _app.prepare(ctx_id=-1, det_size=(320, 320)) return _app # Base directories (Default) DEFAULT_DB_ROOT = "face_database" # per‑person image folders DEFAULT_EMB_ROOT = "faces_db" # embeddings stored here def ensure_dir(path: str) -> None: """Create a directory if it does not exist.""" os.makedirs(path, exist_ok=True) def ensure_person_folder(name: str, db_root: str) -> str: """Return the absolute path to the folder for *name*, creating it if needed.""" folder = os.path.join(db_root, name) ensure_dir(folder) return folder def copy_image_to_folder(name: str, image_path: str, db_root: str) -> str: """Copy *image_path* into the person's folder. If a file with the same name already exists, a numeric suffix is added. Returns the final destination path. """ if not os.path.isfile(image_path): raise FileNotFoundError(f"Image not found: {image_path}") dest_folder = ensure_person_folder(name, db_root) base_name = os.path.basename(image_path) dest_path = os.path.join(dest_folder, base_name) if os.path.exists(dest_path): name_root, ext = os.path.splitext(base_name) counter = 1 while True: new_name = f"{name_root}_{counter}{ext}" dest_path = os.path.join(dest_folder, new_name) if not os.path.exists(dest_path): break counter += 1 shutil.copy2(image_path, dest_path) return dest_path def augment_image(src_path: str, dest_path: str, app=None) -> None: """Create a synthetic occlusion (black rectangle over the eyes) and save it. The function reads *src_path*, detects the face, draws a rectangle covering the eye region, and writes the result to *dest_path*. """ img = cv2.imread(src_path) if img is None: raise ValueError(f"Unable to read image for augmentation: {src_path}") detector = app if app else get_app() faces = detector.get(img) if len(faces) == 0: raise ValueError("No face detected for augmentation.") # Use the first detected face face = faces[0] # InsightFace returns 5 landmarks: left eye, right eye, nose, left mouth, right mouth if not hasattr(face, "kps") or face.kps is None: raise ValueError("Landmarks not available for augmentation.") landmarks = face.kps # shape (5, 2) # Compute bounding box that covers both eyes left_eye = landmarks[0] right_eye = landmarks[1] # Expand a little to cover the whole eye region eye_center = (left_eye + right_eye) / 2 eye_width = np.linalg.norm(right_eye - left_eye) * 1.5 eye_height = eye_width * 0.6 x1 = int(eye_center[0] - eye_width / 2) y1 = int(eye_center[1] - eye_height / 2) x2 = int(eye_center[0] + eye_width / 2) y2 = int(eye_center[1] + eye_height / 2) cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 0), -1) cv2.imwrite(dest_path, img) def generate_embeddings(name: str, db_root: str, emb_root: str, known_embedding: np.ndarray = None, app=None) -> np.ndarray: """Compute embeddings for *all* images in the person's folder and store them. The resulting .npy file is saved to `emb_root/.npy` with shape (N, 512). Returns the generated embeddings. """ detector = app if app else get_app() if known_embedding is not None: # Use provided embedding directly if available # We might still want to try to get more embeddings from images if possible, # but for robustness, we ensure at least this one is saved. embeddings = [known_embedding] # Try to add embeddings from folder images too person_folder = os.path.join(db_root, name) if os.path.isdir(person_folder): for fname in os.listdir(person_folder): img_path = os.path.join(person_folder, fname) if not os.path.isfile(img_path): continue img = cv2.imread(img_path) if img is None: continue faces = detector.get(img) if len(faces) > 0: embeddings.append(faces[0].embedding) ensure_dir(emb_root) emb_array = np.stack(embeddings) from embedding_store import save_embeddings save_embeddings(name, emb_root, emb_array) print(f"Saved {len(embeddings)} embeddings for '{name}' to {emb_root}/{name}.f32emb") return emb_array person_folder = os.path.join(db_root, name) if not os.path.isdir(person_folder): raise FileNotFoundError(f"Person folder not found: {person_folder}") embeddings = [] for fname in os.listdir(person_folder): img_path = os.path.join(person_folder, fname) if not os.path.isfile(img_path): continue img = cv2.imread(img_path) if img is None: continue faces = detector.get(img) if len(faces) == 0: continue embeddings.append(faces[0].embedding) if not embeddings: raise RuntimeError(f"No valid faces found for person '{name}'.") ensure_dir(emb_root) emb_array = np.stack(embeddings) from embedding_store import save_embeddings save_embeddings(name, emb_root, emb_array) print(f"Saved {len(embeddings)} embeddings for '{name}' to {emb_root}/{name}.f32emb") return emb_array def register_face(name: str, image_path: str, db_root: str = DEFAULT_DB_ROOT, emb_root: str = DEFAULT_EMB_ROOT, known_embedding: np.ndarray = None, app=None) -> np.ndarray: """Validate the image, copy it, create an occluded version, and update embeddings. The occluded image is saved with the suffix `_occluded` before the file extension. Returns the generated embeddings. """ # Validate and copy original image dest_original = copy_image_to_folder(name, image_path, db_root) print(f"Original image copied to {dest_original}") # Create occluded version base, ext = os.path.splitext(dest_original) occluded_path = f"{base}_occluded{ext}" try: augment_image(dest_original, occluded_path, app=app) print(f"Occluded image created at {occluded_path}") except Exception as e: print(f"Warning: could not create occluded image – {e}") # Regenerate embeddings for this person (includes original + occluded) return generate_embeddings(name, db_root, emb_root, known_embedding, app=app) if __name__ == "__main__": person_name = input("Enter name for this person (folder will be created/used): ").strip() img_path = input("Enter path to image: ").strip() try: register_face(person_name, img_path) except Exception as e: print(f"Error: {e}")