| import os |
| import shutil |
| import cv2 |
| import numpy as np |
| from insightface.app import FaceAnalysis |
|
|
| |
| _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 |
|
|
| |
| DEFAULT_DB_ROOT = "face_database" |
| DEFAULT_EMB_ROOT = "faces_db" |
|
|
| 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.") |
| |
| face = faces[0] |
| |
| if not hasattr(face, "kps") or face.kps is None: |
| raise ValueError("Landmarks not available for augmentation.") |
| landmarks = face.kps |
| |
| left_eye = landmarks[0] |
| right_eye = landmarks[1] |
| |
| 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/<name>.npy` with shape (N, 512). |
| Returns the generated embeddings. |
| """ |
| detector = app if app else get_app() |
| if known_embedding is not None: |
| |
| |
| |
| embeddings = [known_embedding] |
| |
| |
| 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. |
| """ |
| |
| dest_original = copy_image_to_folder(name, image_path, db_root) |
| print(f"Original image copied to {dest_original}") |
|
|
| |
| 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}") |
|
|
| |
| 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}") |
|
|
