Sandeep2004/Cerevyn_Face_Microservice

1import os2import uuid3from threading import Lock, Thread4import cv25import faiss6import numpy as np7import imgaug.augmenters as iaa8 9from dotenv import load_dotenv10from sqlalchemy import create_engine11from sqlalchemy.orm import sessionmaker, Session12from supabase import create_client13from insightface.app import FaceAnalysis14from models import EmployeeFace15 16# ==================================================17# ENV + SUPABASE18# ==================================================19# For local dev: uses .env file via load_dotenv()20# For HuggingFace Spaces: reads from Settings > Secrets (auto-exposed as env vars)21load_dotenv()22 23SUPABASE_PROJECT_URL = os.getenv("SUPABASE_PROJECT_URL")24SUPABASE_ANON_KEY = os.getenv("ANON_KEY")25SUPABASE_SERVICE_ROLE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")26SUPABASE_DB_URL = os.getenv("SUPABASE_DB_URL")27 28if not SUPABASE_PROJECT_URL:29    raise ValueError("SUPABASE_PROJECT_URL environment variable must be set")30if not SUPABASE_ANON_KEY:31    raise ValueError("ANON_KEY environment variable must be set (set in HF Spaces Secrets)")32if not SUPABASE_DB_URL:33    raise ValueError("SUPABASE_DB_URL environment variable must be set (set in HF Spaces Secrets)")34 35# Use service-role key for storage if provided, fall back to anon36_storage_key = SUPABASE_SERVICE_ROLE_KEY or SUPABASE_ANON_KEY37supabase_storage = create_client(SUPABASE_PROJECT_URL, _storage_key)38# Keep anon client available if needed elsewhere39supabase = create_client(SUPABASE_PROJECT_URL, SUPABASE_ANON_KEY)40 41# ==================================================42# DATABASE (UNCHANGED)43# ==================================================44engine = create_engine(45    SUPABASE_DB_URL,46    pool_pre_ping=True,47    pool_recycle=1800,48)49 50SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False)51 52def get_db():53    db = SessionLocal()54    try:55        yield db56    finally:57        db.close()58 59# ==================================================60# IMAGE UPLOAD (UNCHANGED)61# ==================================================62def upload_to_bucket(file, bucket_name: str):63    file_bytes = file.file.read()64    ext = file.filename.split(".")[-1]65    file_name = f"{uuid.uuid4()}.{ext}"66 67    supabase_storage.storage.from_(bucket_name).upload(file_name, file_bytes)68    return supabase_storage.storage.from_(bucket_name).get_public_url(file_name)69 70def upload_selfie(file):71    return upload_to_bucket(file, "selfies")72 73# ==================================================74# FACE RECOGNITION CONFIG (test_enroll style)75# ==================================================76EMBEDDING_DIM = 51277FAISS_INDEX_PATH = "employee_faces.faiss"78FAISS_MAP_PATH = "employee_faces_map.npy"79 80MODEL_ROOT = "/app/models"81 82face_app = FaceAnalysis(83    name="antelopev2",84    root=MODEL_ROOT,85    providers=["CPUExecutionProvider"]86)87 88face_app.prepare(ctx_id=0)89 90print("InsightFace loaded:", face_app.models.keys())91 92 93# Serialize embedding computation to avoid race conditions94face_lock = Lock()95 96# ==================================================97# AUGMENTATION (SAME PHILOSOPHY AS test_enroll)98# ==================================================99augmenter = iaa.SomeOf((2, 4), [100    iaa.Fliplr(0.5),101    iaa.Affine(rotate=(-45, 45)),102    iaa.Multiply((0.8, 1.2)),103    iaa.GaussianBlur(sigma=(0, 1.0)),104    iaa.AdditiveGaussianNoise(scale=(10, 30)),105    iaa.Sharpen(alpha=(0.2, 0.5), lightness=(0.8, 1.2)),106    iaa.Crop(percent=(0, 0.1)),107    iaa.LinearContrast((0.75, 1.5)),108    iaa.SomeOf((0, 1), [iaa.Grayscale(alpha=1.0)])109])110 111def augment_image(image_rgb, count=50):112    return augmenter(images=[image_rgb] * count)113 114# ==================================================115# EMBEDDING EXTRACTION116# ==================================================117def extract_embedding(image_rgb):118    faces = face_app.get(image_rgb)119    if not faces:120        return None121 122    face = max(123        faces,124        key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1])125    )126    return face.normed_embedding.astype("float32")127 128def mean_embedding(images):129    embeddings = []130 131    for img in images:132        emb = extract_embedding(img)133        if emb is not None:134            embeddings.append(emb)135 136    if not embeddings:137        return None138 139    mean_emb = np.mean(embeddings, axis=0)140    mean_emb /= np.linalg.norm(mean_emb)141    return mean_emb.astype("float32")142 143# ==================================================144# FAISS + MAPPING (IMPORTANT)145# ==================================================146def rebuild_faiss(db: Session):147    """148    Builds FAISS index AND employee_id mapping149    """150    faces = db.query(EmployeeFace).order_by(EmployeeFace.employee_id).all()151    if not faces:152        return153 154    vectors = []155    id_map = []156 157    for f in faces:158        vectors.append(f.embedding)159        id_map.append(f.employee_id)160 161    vectors = np.array(vectors, dtype="float32")162 163    index = faiss.IndexFlatIP(EMBEDDING_DIM)164    index.add(vectors)165 166    faiss.write_index(index, FAISS_INDEX_PATH)167    np.save(FAISS_MAP_PATH, np.array(id_map))168 169def load_faiss():170    if not os.path.exists(FAISS_INDEX_PATH):171        return None, None172 173    index = faiss.read_index(FAISS_INDEX_PATH)174    id_map = np.load(FAISS_MAP_PATH)175 176    return index, id_map177 178# ==================================================179# FACE ENROLLMENT (DB + FAISS)180# ==================================================181def enroll_employee_face(db: Session, employee_id: int, image_bgr, image_url: str = None):182    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)183 184    augmented = augment_image(image_rgb, 12)185    with face_lock:186        mean_emb = mean_embedding(augmented)187 188    if mean_emb is None:189        return False190 191    face = (192        db.query(EmployeeFace)193        .filter(EmployeeFace.employee_id == employee_id)194        .first()195    )196 197    if face:198        face.embedding = mean_emb.tolist()199        if image_url:200            face.reference_image_url = image_url201    else:202        face = EmployeeFace(203            employee_id=employee_id,204            embedding=mean_emb.tolist(),205            reference_image_url=image_url206        )207        db.add(face)208 209    db.commit()210 211    # Rebuild FAISS in background using a fresh session to avoid closed-session issues212    def _rebuild():213        _db = SessionLocal()214        try:215            rebuild_faiss(_db)216        finally:217            _db.close()218 219    Thread(target=_rebuild, daemon=True).start()220    return True221 222# ==================================================223# FACE VERIFICATION (WITH MAPPING)224# ==================================================225def verify_employee_face(image_bgr, threshold=0.35):226    index, id_map = load_faiss()227    if index is None:228        return None229 230    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)231    emb = extract_embedding(image_rgb)232 233    if emb is None:234        return None235 236    emb = emb.reshape(1, -1).astype("float32")237    D, I = index.search(emb, 1)238 239    score = float(D[0][0])240    if score >= threshold:241        return int(id_map[I[0][0]])242 243    return None244