Update app.py

app.py

@@ -1,199 +1,244 @@
-# SecureFace ID – FIXED VERSION
 import os
 import cv2
 import numpy as np
 import gradio as gr
+import json
 from ultralytics import YOLO
-from huggingface_hub import hf_hub_download
-import insightface
 from insightface.app import FaceAnalysis
-import faiss
-
-# ==================== PATHS ====================
-KNOWN_EMBS_PATH = "known_embeddings.npy"
-KNOWN_NAMES_PATH = "known_names.npy"
-
-# ==================== MODELS ====================
-# Load YOLO for fast detection
-print("Loading YOLOv8...")
-model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt")
-detector = YOLO(model_path)
-
-# Load InsightFace for embedding extraction
-print("Loading InsightFace...")
-recognizer = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
-recognizer.prepare(ctx_id=0, det_size=(640,640))
-
-# FAISS index setup
-index = faiss.IndexHNSWFlat(512, 32)
-index.hnsw.efSearch = 16
-known_names = []
-
-# Load database at startup
-if os.path.exists(KNOWN_EMBS_PATH) and os.path.exists(KNOWN_NAMES_PATH):
-    try:
-        embs = np.load(KNOWN_EMBS_PATH)
-        known_names = np.load(KNOWN_NAMES_PATH, allow_pickle=True).tolist()
-        if embs.shape[0] > 0:
-            index.add(embs.astype('float32'))
-            print(f"✅ Loaded {len(known_names)} identities from disk.")
-    except Exception as e:
-        print(f"⚠️ Database Error: {e}")
+from huggingface_hub import hf_hub_download
 
-# ==================== PROCESS FRAME ====================
-def process_frame(frame, blur_type="gaussian", intensity=50, expand=1.4, show_labels=True):
-    if frame is None: return None
-    
-    img = frame.copy()
-    h, w = img.shape[:2]
-    
-    # 1. Detect Faces with YOLO
-    results = detector(img, conf=0.4, verbose=False)[0]
+# ==========================================
+# CONFIGURATION & STATE MANAGEMENT
+# ==========================================
+DB_FILE = "face_db.json"
+EMBEDDINGS_FILE = "face_embeddings.npy"
 
-    for box in results.boxes:
-        x1, y1, x2, y2 = map(int, box.xyxy[0])
+class FaceSystem:
+    def __init__(self):
+        print("🚀 Initializing AI Models...")
         
-        # Calculate expanded crop for recognition context
-        ew = int((x2-x1)*(expand-1)/2)
-        eh = int((y2-y1)*(expand-1)/2)
-        cx1 = max(0, x1-ew); cy1 = max(0, y1-eh)
-        cx2 = min(w, x2+ew); cy2 = min(h, y2+eh)
-
-        # 2. Recognition Logic
-        # We crop the face and convert to BGR (InsightFace expects BGR)
-        crop = cv2.cvtColor(img[cy1:cy2, cx1:cx2], cv2.COLOR_RGB2BGR)
+        # 1. Load YOLOv8-Face (Best for detection speed/accuracy)
+        # We download a specific version trained for faces
+        model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt")
+        self.detector = YOLO(model_path)
         
-        # Run InsightFace on the crop
-        faces = recognizer.get(crop)
-
-        name = "Unknown"
-        match_found = False
+        # 2. Load InsightFace (Best for recognition accuracy)
+        # 'buffalo_l' is the large model (higher accuracy). Use 'buffalo_s' if you need more speed.
+        self.recognizer = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
+        self.recognizer.prepare(ctx_id=0, det_size=(640, 640))
         
-        if faces and index.ntotal > 0:
-            # Take the largest face in the crop (usually the correct one)
-            main_face = max(faces, key=lambda x: (x.bbox[2]-x.bbox[0]) * (x.bbox[3]-x.bbox[1]))
-            
-            emb = main_face.normed_embedding.reshape(1, -1).astype('float32')
-            D, I = index.search(emb, k=1)
-            
-            # Threshold: Lower is better for L2 distance. 
-            # 0.8 is a safe balance; 0.6 is very strict.
-            if D[0][0] < 0.8: 
-                name = known_names[I[0][0]]
-                match_found = True
-
-        # 3. Blur Logic (Privacy)
-        if blur_type != "none":
-            face_region = img[y1:y2, x1:x2] # Blur only the tight box, not expanded
-            if blur_type == "gaussian":
-                k = max(21, int(min(x2-x1, y2-y1) * intensity / 100) | 1)
-                blurred = cv2.GaussianBlur(face_region, (k,k), 0)
-                img[y1:y2, x1:x2] = blurred
-            elif blur_type == "pixelate":
-                # Map intensity 10-100 to pixel block size 20-3
-                block_size = max(3, int(20 * (1 - intensity/120))) 
-                small = cv2.resize(face_region, (max(1, (x2-x1)//block_size), max(1, (y2-y1)//block_size)))
-                blurred = cv2.resize(small, (x2-x1, y2-y1), interpolation=cv2.INTER_NEAREST)
-                img[y1:y2, x1:x2] = blurred
-            elif blur_type == "solid":
-                 cv2.rectangle(img, (x1,y1), (x2,y2), (0,0,0), -1)
-
-        # 4. Draw Labels (Identity)
-        if show_labels:
-            color = (0, 255, 0) if match_found else (0, 0, 255) # Green for known, Red for unknown
-            
-            # Draw box
-            cv2.rectangle(img, (x1,y1), (x2,y2), color, 2)
+        # 3. Load Database
+        self.known_names = []
+        self.known_embeddings = np.empty((0, 512))
+        self.load_db()
+        print("✅ System Ready.")
+
+    def load_db(self):
+        """Load names and vector embeddings from disk."""
+        if os.path.exists(DB_FILE) and os.path.exists(EMBEDDINGS_FILE):
+            with open(DB_FILE, 'r') as f:
+                self.known_names = json.load(f)
+            self.known_embeddings = np.load(EMBEDDINGS_FILE)
+            print(f"📂 Loaded {len(self.known_names)} identities.")
+        else:
+            print("📂 Database empty. Starting fresh.")
+
+    def save_db(self):
+        """Save current memory to disk."""
+        with open(DB_FILE, 'w') as f:
+            json.dump(self.known_names, f)
+        np.save(EMBEDDINGS_FILE, self.known_embeddings)
+
+    def enroll_user(self, name, image):
+        """Analyzes an image and adds the person to the database."""
+        if image is None or name.strip() == "":
+            return "⚠️ Error: Missing name or photo."
+        
+        # InsightFace expects BGR format (OpenCV standard)
+        img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        
+        # Detect face for enrollment
+        faces = self.recognizer.get(img_bgr)
+        
+        if len(faces) == 0:
+            return "⚠️ Error: No face detected. Please use a clear, front-facing photo."
+        
+        # Get the largest face (in case there are others in background)
+        # We sort by area (width * height)
+        face = sorted(faces, key=lambda x: (x.bbox[2]-x.bbox[0]) * (x.bbox[3]-x.bbox[1]))[-1]
+        
+        # Extract embedding (512-dimensional vector)
+        embedding = face.normed_embedding.reshape(1, -1)
+        
+        # Add to memory
+        if self.known_embeddings.shape[0] == 0:
+            self.known_embeddings = embedding
+        else:
+            self.known_embeddings = np.vstack([self.known_embeddings, embedding])
             
-            # Draw label background and text
-            label_str = name
-            (tw, th), _ = cv2.getTextSize(label_str, cv2.FONT_HERSHEY_DUPLEX, 0.8, 2)
-            cv2.rectangle(img, (x1, y1-30), (x1+tw+10, y1), color, -1)
-            cv2.putText(img, label_str, (x1+5, y1-8), cv2.FONT_HERSHEY_DUPLEX, 0.8, (255,255,255), 2)
-
-    return img
-
-# ==================== ENROLL FUNCTION ====================
-def enroll_person(name, face_image):
-    global index, known_names
-
-    # FIX 1: Proper None check for numpy array
-    if face_image is None or not name or not name.strip():
-        return "⚠️ Error: Please provide both a name and a photo."
+        self.known_names.append(name)
+        
+        # Save to disk
+        self.save_db()
+        
+        return f"✅ Success: '{name}' added to database."
 
-    # Convert RGB (Gradio) to BGR (OpenCV/InsightFace)
-    bgr = cv2.cvtColor(face_image, cv2.COLOR_RGB2BGR)
+    def recognize_and_process(self, frame, blur_intensity=20, threshold=0.5):
+        """
+        The Core Loop: Detect -> Identify -> Anonymize
+        """
+        if frame is None: return None
+        
+        img_vis = frame.copy()
+        h, w = img_vis.shape[:2]
+        
+        # 1. DETECT (YOLOv8)
+        # conf=0.5 reduces false positives
+        results = self.detector(img_vis, conf=0.5, verbose=False)
+        
+        for result in results:
+            boxes = result.boxes
+            for box in boxes:
+                x1, y1, x2, y2 = map(int, box.xyxy[0])
+                
+                # Add margin to the face crop for better recognition
+                margin = 0
+                cx1 = max(0, x1 - margin)
+                cy1 = max(0, y1 - margin)
+                cx2 = min(w, x2 + margin)
+                cy2 = min(h, y2 + margin)
+                
+                face_crop = img_vis[cy1:cy2, cx1:cx2]
+                
+                # 2. IDENTIFY (InsightFace + Vector Math)
+                name = "Unknown"
+                color = (200, 0, 0) # Red for unknown
+                
+                if self.known_embeddings.shape[0] > 0 and face_crop.size > 0:
+                    # Convert crop to BGR for InsightFace
+                    face_crop_bgr = cv2.cvtColor(face_crop, cv2.COLOR_RGB2BGR)
+                    
+                    # Extract embedding
+                    analysis = self.recognizer.get(face_crop_bgr)
+                    
+                    if len(analysis) > 0:
+                        # Get embedding of the main face in the crop
+                        target_emb = analysis[0].normed_embedding
+                        
+                        # Calculate Cosine Similarity against ALL known faces at once
+                        # (Dot product of normalized vectors = Cosine Similarity)
+                        similarities = np.dot(self.known_embeddings, target_emb)
+                        
+                        # Find best match
+                        best_idx = np.argmax(similarities)
+                        best_score = similarities[best_idx]
+                        
+                        if best_score > threshold:
+                            name = self.known_names[best_idx]
+                            color = (0, 255, 0) # Green for known
+                            # Optional: Show confidence score
+                            # name += f" ({int(best_score*100)}%)"
+
+                # 3. PRIVACY (Pixelation)
+                # We extract the ROI again (strictly inside the box)
+                roi = img_vis[y1:y2, x1:x2]
+                if roi.size > 0:
+                    # Pixelation logic: Downscale -> Upscale
+                    # Map intensity (10-100) to a factor. Lower factor = bigger blocks.
+                    # Intensity 10 = Block size 20px
+                    # Intensity 100 = Block size 3px (barely blurred)
+                    block_size = max(3, int(30 - (blur_intensity / 4)))
+                    
+                    h_roi, w_roi = roi.shape[:2]
+                    
+                    # Downscale
+                    small = cv2.resize(roi, (max(1, w_roi // block_size), max(1, h_roi // block_size)), interpolation=cv2.INTER_LINEAR)
+                    # Upscale (Nearest Neighbor creates the pixel effect)
+                    pixelated = cv2.resize(small, (w_roi, h_roi), interpolation=cv2.INTER_NEAREST)
+                    
+                    # Apply back to image
+                    img_vis[y1:y2, x1:x2] = pixelated
+
+                # 4. OVERLAY (ID Label)
+                # We draw the label ON TOP of the blurred face
+                
+                # Box
+                cv2.rectangle(img_vis, (x1, y1), (x2, y2), color, 2)
+                
+                # Text Background
+                label_size, baseline = cv2.getTextSize(name, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
+                cv2.rectangle(img_vis, (x1, y1 - label_size[1] - 10), (x1 + label_size[0] + 10, y1), color, -1)
+                
+                # Text
+                cv2.putText(img_vis, name, (x1 + 5, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+                
+        return img_vis
+
+# Initialize System
+system = FaceSystem()
+
+# ==========================================
+# GRADIO INTERFACE
+# ==========================================
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald", secondary_hue="zinc")) as demo:
     
-    faces = recognizer.get(bgr)
+    gr.Markdown(
+        """
+        # 👁️ SecureVision Pro
+        **Enterprise-Grade Identity Protection & Recognition System**
+        """
+    )
     
-    if not faces:
-        return "⚠️ Error: No face detected. Please use a clear frontal photo."
-
-    # Pick the largest face if multiple are found
-    main_face = max(faces, key=lambda x: (x.bbox[2]-x.bbox[0]) * (x.bbox[3]-x.bbox[1]))
-    new_emb = main_face.normed_embedding.reshape(1, 512)
-
-    # Load existing data to ensure sync
-    if os.path.exists(KNOWN_EMBS_PATH):
-        embs = np.load(KNOWN_EMBS_PATH)
-        # Ensure 2D array
-        if len(embs.shape) == 1: embs = embs.reshape(1, -1)
-    else:
-        embs = np.empty((0,512))
-
-    # Append new data
-    embs = np.vstack([embs, new_emb])
-    known_names.append(name)
-
-    # Save to disk
-    np.save(KNOWN_EMBS_PATH, embs)
-    np.save(KNOWN_NAMES_PATH, np.array(known_names))
-    
-    # Rebuild Index
-    index.reset()
-    index.add(embs.astype('float32'))
-
-    return f"✅ Success: **{name}** has been enrolled!"
-
-# ==================== GRADIO UI ====================
-with gr.Blocks(title="SecureFace ID", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🛡️ SecureFace ID")
-
-    with gr.Tab("📹 Live Surveillance"):
-        with gr.Row():
-            with gr.Column():
-                cam = gr.Image(sources=["webcam"], streaming=True, label="Live Feed", height=450)
-            with gr.Column():
-                output = gr.Image(label="Protected Stream", height=450)
-        
-        with gr.Row():
-            blur_type = gr.Radio(["gaussian", "pixelate", "solid", "none"], value="pixelate", label="Privacy Filter")
-            intensity = gr.Slider(1, 100, 80, label="Blur Intensity")
-            expand = gr.Slider(1.0, 2.0, 1.3, label="Context Area")
-            show_names = gr.Checkbox(True, label="Show IDs Overlay")
+    with gr.Tabs():
+        # TAB 1: MONITOR
+        with gr.Tab("📹 Live Monitor"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    input_feed = gr.Image(sources=["webcam"], streaming=True, label="Camera Feed")
+                with gr.Column(scale=2):
+                    output_feed = gr.Image(label="Processed Stream (Privacy + ID)")
             
-        # Connect the stream
-        cam.stream(process_frame, [cam, blur_type, intensity, expand, show_names], output)
-
-    with gr.Tab("👤 Enroll Person"):
-        with gr.Row():
-            with gr.Column():
-                name_input = gr.Textbox(label="Full Name / ID", placeholder="e.g. Agent Smith")
-                photo_input = gr.Image(label="Reference Photo", sources=["upload", "webcam"], height=300)
-                enroll_btn = gr.Button("Add to Database", variant="primary")
-            with gr.Column():
-                enroll_status = gr.Markdown("### Status: Waiting...")
-                db_view = gr.Markdown()
-
-    # Database viewer updater
-    def get_db_status():
-        if not known_names: return "Database is empty."
-        return f"### 📂 Registered Users ({len(known_names)}):\n" + "\n".join([f"- {n}" for n in list(set(known_names))])
-
-    # Event wiring
-    enroll_btn.click(enroll_person, inputs=[name_input, photo_input], outputs=enroll_status)
-    enroll_btn.click(get_db_status, outputs=db_view)
-    demo.load(get_db_status, outputs=db_view)
+            with gr.Accordion("⚙️ System Settings", open=True):
+                blur_slider = gr.Slider(1, 100, value=50, label="Privacy Level (Pixelation Strength)")
+                conf_slider = gr.Slider(0.1, 0.9, value=0.5, label="Recognition Strictness (Threshold)")
+
+            # Connect the stream
+            input_feed.stream(
+                fn=system.recognize_and_process, 
+                inputs=[input_feed, blur_slider, conf_slider], 
+                outputs=output_feed
+            )
+
+        # TAB 2: DATABASE
+        with gr.Tab("👤 Database Management"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### Enroll New Personnel")
+                    new_name = gr.Textbox(label="Full Name / ID")
+                    new_photo = gr.Image(sources=["upload", "webcam"], label="Reference Photo")
+                    add_btn = gr.Button("Enroll User", variant="primary")
+                    status_msg = gr.Markdown("")
+                
+                with gr.Column():
+                    gr.Markdown("### Database Status")
+                    # A function to list current users
+                    def get_user_list():
+                        if not system.known_names: return "No users enrolled."
+                        return "\n".join([f"• {n}" for n in system.known_names])
+                    
+                    user_list = gr.Markdown(get_user_list)
+                    refresh_btn = gr.Button("Refresh List")
+            
+            # Enroll Logic
+            add_btn.click(
+                fn=system.enroll_user,
+                inputs=[new_name, new_photo],
+                outputs=status_msg
+            )
+            # Refresh Logic
+            refresh_btn.click(fn=get_user_list, outputs=user_list)
+            # Auto-refresh list on enroll
+            add_btn.click(fn=get_user_list, outputs=user_list)
 
 if __name__ == "__main__":
     demo.launch()