Update app.py

app.py

@@ -3,6 +3,7 @@ import cv2
 import numpy as np
 import gradio as gr
 import json
+import tempfile
 from ultralytics import YOLO
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
@@ -17,13 +18,11 @@ class FaceSystem:
     def __init__(self):
         print("🚀 Initializing AI Models...")
         
-        # 1. Load YOLOv8-Face (Best for detection speed/accuracy)
-        # We download a specific version trained for faces
+        # 1. Load YOLOv8-Face
         model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt")
         self.detector = YOLO(model_path)
         
-        # 2. Load InsightFace (Best for recognition accuracy)
-        # 'buffalo_l' is the large model (higher accuracy). Use 'buffalo_s' if you need more speed.
+        # 2. Load InsightFace
         self.recognizer = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
         self.recognizer.prepare(ctx_id=0, det_size=(640, 640))
         
@@ -34,7 +33,6 @@ class FaceSystem:
         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)
@@ -44,56 +42,41 @@ class FaceSystem:
             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."
+            return "⚠️ Error: No face detected."
         
-        # Get the largest face (in case there are others in background)
-        # We sort by area (width * height)
+        # Get largest face
         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])
             
         self.known_names.append(name)
-        
-        # Save to disk
         self.save_db()
-        
         return f"✅ Success: '{name}' added to database."
 
     def recognize_and_process(self, frame, blur_intensity=20, threshold=0.5):
-        """
-        The Core Loop: Detect -> Identify -> Anonymize
-        """
+        """Core processing logic for a single frame"""
         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
+        # Detect
         results = self.detector(img_vis, conf=0.5, verbose=False)
         
         for result in results:
@@ -101,120 +84,125 @@ class FaceSystem:
             for box in boxes:
                 x1, y1, x2, y2 = map(int, box.xyxy[0])
                 
-                # Add margin to the face crop for better recognition
+                # Add context margin
                 margin = 0
-                cx1 = max(0, x1 - margin)
-                cy1 = max(0, y1 - margin)
-                cx2 = min(w, x2 + margin)
-                cy2 = min(h, y2 + margin)
-                
+                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)
+                # Identify
                 name = "Unknown"
-                color = (200, 0, 0) # Red for unknown
+                color = (200, 0, 0) # Red
                 
                 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:
+                        if similarities[best_idx] > threshold:
                             name = self.known_names[best_idx]
-                            color = (0, 255, 0) # Green for known
-                            # Optional: Show confidence score
-                            # name += f" ({int(best_score*100)}%)"
+                            color = (0, 255, 0) # Green
 
-                # 3. PRIVACY (Pixelation)
-                # We extract the ROI again (strictly inside the box)
+                # Privacy Blur (Pixelate)
                 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
+                # Overlay
                 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)
+                label_size, _ = 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
+# Initialize
 system = FaceSystem()
 
 # ==========================================
-# GRADIO INTERFACE
+# VIDEO PROCESSING HELPER
 # ==========================================
+def process_video_file(video_path, blur_intensity, threshold):
+    """Reads a video file, processes every frame, saves it, and returns path."""
+    if video_path is None: return None
+    
+    cap = cv2.VideoCapture(video_path)
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    
+    # Create temp output file
+    temp_out = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
+    output_path = temp_out.name
+    
+    # Setup writer (mp4v is usually safe for CPU)
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret: break
+        
+        # Process frame using our existing core function
+        # We convert BGR (OpenCV) to RGB (needed for our function) then back to BGR
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        processed_rgb = system.recognize_and_process(frame_rgb, blur_intensity, threshold)
+        processed_bgr = cv2.cvtColor(processed_rgb, cv2.COLOR_RGB2BGR)
+        
+        writer.write(processed_bgr)
+        
+    cap.release()
+    writer.release()
+    return output_path
+
 # ==========================================
-# GRADIO INTERFACE
+# GRADIO INTERFACE (Fixed UI)
 # ==========================================
-# We removed the 'theme=' argument to prevent version errors
 with gr.Blocks() as demo:
     
-    gr.Markdown(
-        """
-        # 👁️ SecureVision Pro
-        **Enterprise-Grade Identity Protection & Recognition System**
-        """
-    )
+    gr.Markdown("# 👁️ SecureVision Pro")
     
     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)")
+        # --- SURVEILLANCE TAB ---
+        with gr.Tab("📹 Surveillance"):
             
-            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"):
+            # Global Settings for this tab
+            with gr.Accordion("⚙️ Settings", open=False):
+                blur_slider = gr.Slider(1, 100, value=50, label="Privacy Level")
+                conf_slider = gr.Slider(0.1, 0.9, value=0.5, label="Recognition Threshold")
+
+            with gr.Tabs():
+                # Sub-Tab 1: Live Webcam
+                with gr.TabItem("Live Webcam"):
+                    with gr.Row():
+                        web_in = gr.Image(sources=["webcam"], streaming=True, label="Live Feed")
+                        web_out = gr.Image(label="Live Output")
+                    web_in.stream(system.recognize_and_process, [web_in, blur_slider, conf_slider], web_out)
+
+                # Sub-Tab 2: Upload Image
+                with gr.TabItem("Upload Image"):
+                    with gr.Row():
+                        img_in = gr.Image(sources=["upload", "clipboard"], label="Upload Image")
+                        img_out = gr.Image(label="Processed Image")
+                    img_btn = gr.Button("Analyze Image", variant="primary")
+                    img_btn.click(system.recognize_and_process, [img_in, blur_slider, conf_slider], img_out)
+
+                # Sub-Tab 3: Upload Video
+                with gr.TabItem("Upload Video"):
+                    with gr.Row():
+                        vid_in = gr.Video(label="Upload Video")
+                        vid_out = gr.Video(label="Processed Output")
+                    vid_btn = gr.Button("Process Video", variant="primary")
+                    vid_btn.click(process_video_file, [vid_in, blur_slider, conf_slider], vid_out)
+
+        # --- DATABASE TAB ---
+        with gr.Tab("👤 Database"):
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("### Enroll New Personnel")
@@ -225,7 +213,6 @@ with gr.Blocks() as demo:
                 
                 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])
@@ -233,16 +220,9 @@ with gr.Blocks() as demo:
                     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)
+            add_btn.click(system.enroll_user, [new_name, new_photo], status_msg)
+            refresh_btn.click(get_user_list, outputs=user_list)
+            add_btn.click(get_user_list, outputs=user_list)
 
 if __name__ == "__main__":
     demo.launch()