Update app.py

app.py

@@ -1,5 +1,6 @@
 """
-Face Privacy Tool (Gradio UI with YOLOv8 Segmentation & Video Support)
+Face Recognition Tool (YOLO + DeepFace)
+Extended from Face Privacy Tool to add face recognition capabilities
 """
 
 # --- Standard Libraries ---
@@ -7,9 +8,12 @@ import logging
 import atexit
 import tempfile
 import os
+import json
+import pickle
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Tuple, Optional
+from pathlib import Path
 
 # --- Computer Vision & UI Libraries ---
 import cv2
@@ -17,29 +21,183 @@ import numpy as np
 import gradio as gr
 from ultralytics import YOLO
 
+# --- Face Recognition Libraries ---
+try:
+    from deepface import DeepFace
+    DEEPFACE_AVAILABLE = True
+except ImportError:
+    DEEPFACE_AVAILABLE = False
+    logging.warning("DeepFace not installed. Install with: pip install deepface")
+
 # --- Configure Logging ---
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
+# ====================================================
+# FACE DATABASE MANAGER
+# ====================================================
+class FaceDatabase:
+    """Manages known faces and their embeddings."""
+    
+    def __init__(self, db_path: str = "face_database"):
+        self.db_path = Path(db_path)
+        self.db_path.mkdir(exist_ok=True)
+        self.embeddings_file = self.db_path / "embeddings.pkl"
+        self.known_faces = {}  # {person_id: {"name": str, "embedding": np.array, "image_path": str}}
+        self.load_database()
+    
+    def load_database(self):
+        """Load all known faces and their embeddings."""
+        if self.embeddings_file.exists():
+            try:
+                with open(self.embeddings_file, 'rb') as f:
+                    self.known_faces = pickle.load(f)
+                logger.info(f"✅ Loaded {len(self.known_faces)} known faces from database")
+            except Exception as e:
+                logger.error(f"Failed to load embeddings: {e}")
+                self.known_faces = {}
+        else:
+            logger.info("No existing database found. Starting fresh.")
+    
+    def save_database(self):
+        """Save embeddings to disk."""
+        try:
+            with open(self.embeddings_file, 'wb') as f:
+                pickle.dump(self.known_faces, f)
+            logger.info("✅ Database saved successfully")
+        except Exception as e:
+            logger.error(f"Failed to save database: {e}")
+    
+    def add_person(self, person_id: str, name: str, image: np.ndarray) -> Tuple[bool, str]:
+        """Add a new person to the database."""
+        if not DEEPFACE_AVAILABLE:
+            return False, "❌ DeepFace not installed"
+        
+        try:
+            # Generate embedding
+            embedding_objs = DeepFace.represent(
+                img_path=image,
+                model_name="Facenet512",
+                enforce_detection=True
+            )
+            
+            if not embedding_objs:
+                return False, "❌ No face detected in image"
+            
+            embedding = np.array(embedding_objs[0]["embedding"])
+            
+            # Save image
+            person_dir = self.db_path / person_id
+            person_dir.mkdir(exist_ok=True)
+            image_path = person_dir / "face.jpg"
+            cv2.imwrite(str(image_path), cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+            
+            # Store in database
+            self.known_faces[person_id] = {
+                "name": name,
+                "embedding": embedding,
+                "image_path": str(image_path)
+            }
+            
+            self.save_database()
+            return True, f"✅ Added {name} (ID: {person_id}) to database"
+            
+        except Exception as e:
+            logger.error(f"Error adding person: {e}")
+            return False, f"❌ Error: {str(e)}"
+    
+    def remove_person(self, person_id: str) -> Tuple[bool, str]:
+        """Remove a person from the database."""
+        if person_id in self.known_faces:
+            del self.known_faces[person_id]
+            self.save_database()
+            return True, f"✅ Removed person with ID: {person_id}"
+        return False, f"❌ Person ID not found: {person_id}"
+    
+    def recognize_face(self, face_image: np.ndarray, threshold: float = 0.6) -> Dict[str, Any]:
+        """
+        Recognize a face by comparing with known embeddings.
+        Returns: {"match": bool, "person_id": str, "name": str, "distance": float}
+        """
+        if not DEEPFACE_AVAILABLE or len(self.known_faces) == 0:
+            return {"match": False, "person_id": "unknown", "name": "Unknown", "distance": 999}
+        
+        try:
+            # Generate embedding for input face
+            embedding_objs = DeepFace.represent(
+                img_path=face_image,
+                model_name="Facenet512",
+                enforce_detection=False  # We already detected it with YOLO
+            )
+            
+            if not embedding_objs:
+                return {"match": False, "person_id": "unknown", "name": "Unknown", "distance": 999}
+            
+            query_embedding = np.array(embedding_objs[0]["embedding"])
+            
+            # Find best match
+            best_match = None
+            best_distance = float('inf')
+            
+            for person_id, data in self.known_faces.items():
+                # Cosine distance
+                distance = np.linalg.norm(query_embedding - data["embedding"])
+                
+                if distance < best_distance:
+                    best_distance = distance
+                    best_match = {
+                        "match": distance < threshold,
+                        "person_id": person_id,
+                        "name": data["name"],
+                        "distance": distance
+                    }
+            
+            if best_match and best_match["match"]:
+                return best_match
+            else:
+                return {"match": False, "person_id": "unknown", "name": "Unknown", "distance": best_distance}
+                
+        except Exception as e:
+            logger.error(f"Recognition error: {e}")
+            return {"match": False, "person_id": "unknown", "name": "Unknown", "distance": 999}
+    
+    def list_all_people(self) -> str:
+        """Return a formatted list of all known people."""
+        if not self.known_faces:
+            return "📭 Database is empty"
+        
+        result = "### 👥 Known People:\n\n"
+        for person_id, data in self.known_faces.items():
+            result += f"- **{data['name']}** (ID: `{person_id}`)\n"
+        return result
+
+# Global database instance
+FACE_DB: Optional[FaceDatabase] = None
+
+def get_face_database() -> FaceDatabase:
+    """Get or create the global face database instance."""
+    global FACE_DB
+    if FACE_DB is None:
+        FACE_DB = FaceDatabase()
+    return FACE_DB
+
 # ====================================================
 # TEMPORARY FILE CLEANUP
 # ====================================================
 TEMP_FILES = []
 
 def cleanup_temp_files():
-    """Clean up any temporary files created during the session on exit."""
     for f in TEMP_FILES:
         try:
             if os.path.exists(f):
                 os.remove(f)
-                logger.info(f"🗑️ Cleaned up temporary file: {f}")
+                logger.info(f"🗑️ Cleaned up: {f}")
         except Exception as e:
-            logger.warning(f"⚠️ Failed to delete temporary file {f}: {e}")
+            logger.warning(f"⚠️ Failed to delete {f}: {e}")
 
 atexit.register(cleanup_temp_files)
 
 def create_temp_file(suffix=".mp4") -> str:
-    """Creates a temporary file and registers it for cleanup."""
     path = tempfile.mktemp(suffix=suffix)
     TEMP_FILES.append(path)
     return path
@@ -54,564 +212,313 @@ SENSITIVITY_MAP = {
 }
 
 def get_confidence_from_sensitivity(sensitivity: str) -> float:
-    """Converts user-friendly sensitivity text to numerical confidence threshold."""
     return SENSITIVITY_MAP.get(sensitivity, 0.5)
 
 # ====================================================
-# CONFIGURATION DATA CLASSES
+# CONFIGURATION
 # ====================================================
-@dataclass
-class BlurConfig:
-    """Configuration for blur effects."""
-    type: str = "pixelate"
-    intensity: float = 25.0
-    pixel_size: int = 25
-    solid_color: Tuple[int, int, int] = (0, 0, 0)
-    adaptive_blur: bool = True
-    min_kernel: int = 15
-    max_kernel: int = 95
-
 @dataclass
 class DetectionConfig:
-    """Configuration for the face detector."""
     min_confidence: float = 0.5
     model_path: str = "yolov8n-face.pt"
 
-@dataclass
-class AppConfig:
-    """Main application configuration."""
-    blur: BlurConfig = field(default_factory=BlurConfig)
-    detection: DetectionConfig = field(default_factory=DetectionConfig)
-    scaling_factor: float = 1.2
-    forehead_margin: int = 20
-    face_margin: int = 15
-
-# ====================================================
-# BLUR EFFECTS (STRATEGY PATTERN)
-# ====================================================
-class BlurEffect(ABC):
-    """Abstract base class for blur effects."""
-    def __init__(self, config: BlurConfig):
-        self.config = config
-
-    @abstractmethod
-    def apply(self, image: np.ndarray, roi: Tuple[int, int, int, int]) -> np.ndarray:
-        """Apply the blur effect to the region of interest (ROI)."""
-        pass
-
-class GaussianBlur(BlurEffect):
-    """Gaussian blur with adaptive kernel sizing for a natural look."""
-    def apply(self, image: np.ndarray, roi: Tuple[int, int, int, int]) -> np.ndarray:
-        x, y, w, h = roi
-        face_roi = image[y:y+h, x:x+w]
-        if face_roi.size == 0: 
-            return image
-
-        if self.config.adaptive_blur:
-            min_dim = min(w, h)
-            kernel_val = int(min_dim * (self.config.intensity / 100.0))
-            kernel_val = max(self.config.min_kernel, min(kernel_val, self.config.max_kernel))
-        else:
-            kernel_val = int(self.config.intensity)
-        
-        kernel_val = kernel_val | 1  # Ensure kernel size is odd
-        blurred_roi = cv2.GaussianBlur(face_roi, (kernel_val, kernel_val), 0)
-        image[y:y+h, x:x+w] = blurred_roi
-        return image
-
-class PixelateBlur(BlurEffect):
-    """Pixelation effect for a retro/digital privacy look."""
-    def apply(self, image: np.ndarray, roi: Tuple[int, int, int, int]) -> np.ndarray:
-        x, y, w, h = roi
-        face_roi = image[y:y+h, x:x+w]
-        if face_roi.size == 0: 
-            return image
-
-        h_roi, w_roi = face_roi.shape[:2]
-        pixel_size = self.config.pixel_size
-        if pixel_size <= 0: 
-            return image
-
-        small = cv2.resize(face_roi, (max(1, w_roi // pixel_size), max(1, h_roi // pixel_size)), interpolation=cv2.INTER_LINEAR)
-        pixelated = cv2.resize(small, (w_roi, h_roi), interpolation=cv2.INTER_NEAREST)
-        image[y:y+h, x:x+w] = pixelated
-        return image
-
-class SolidColorBlur(BlurEffect):
-    """Solid color rectangle overlay for complete redaction."""
-    def apply(self, image: np.ndarray, roi: Tuple[int, int, int, int]) -> np.ndarray:
-        x, y, w, h = roi
-        cv2.rectangle(image, (x, y), (x+w, y+h), self.config.solid_color, -1)
-        return image
-
-def get_blur_effect(config: BlurConfig) -> BlurEffect:
-    """Factory function to create a blur effect instance."""
-    blur_effects = {"gaussian": GaussianBlur, "pixelate": PixelateBlur, "solid": SolidColorBlur}
-    blur_class = blur_effects.get(config.type)
-    if not blur_class: 
-        raise ValueError(f"Unknown blur type: {config.type}")
-    return blur_class(config)
-
 # ====================================================
-# YOLOv8 FACE DETECTOR (SINGLETON)
+# YOLO FACE DETECTOR
 # ====================================================
 class YOLOv8FaceDetector:
-    """Unified face detector using YOLOv8-Face model."""
     def __init__(self, config: DetectionConfig):
         try:
-            logger.info(f"Attempting to load YOLOv8-Face model: {config.model_path}")
+            logger.info(f"Loading model: {config.model_path}")
             self.model = YOLO(config.model_path)
             self.min_conf = config.min_confidence
-            logger.info("✅ Model loaded successfully.")
+            logger.info("✅ Model loaded")
         except Exception as e:
-            logger.error(f"❌ Failed to load model: {e}")
-            raise RuntimeError(f"Model loading failed. Ensure '{config.model_path}' is available.") from e
-
-    def detect_faces(self, image: np.ndarray, conf_threshold: float, return_annotated: bool = False) -> Tuple[List[Dict[str, Any]], Optional[np.ndarray]]:
-        """Detects faces in an image with a given confidence threshold."""
+            logger.error(f"❌ Model loading failed: {e}")
+            raise RuntimeError(f"Cannot load model '{config.model_path}'") from e
+
+    def detect_faces(self, image: np.ndarray, conf_threshold: float, 
+                     recognize: bool = False) -> Tuple[List[Dict[str, Any]], np.ndarray]:
+        """
+        Detect faces and optionally recognize them.
+        Returns: (faces_list, annotated_image)
+        """
         results = self.model(image, conf=conf_threshold, verbose=False)
         faces = []
-        annotated_image = image.copy() if return_annotated else None
+        annotated_image = image.copy()
+        
+        face_db = get_face_database() if recognize else None
 
         for r in results:
-            if r.boxes is None: 
+            if r.boxes is None:
                 continue
             for box in r.boxes:
                 x1, y1, x2, y2 = map(int, box.xyxy[0])
                 confidence = float(box.conf[0])
-                faces.append({"x": x1, "y": y1, "width": x2 - x1, "height": y2 - y1, "confidence": confidence})
                 
-                if return_annotated:
-                    # Draw bounding box
+                face_info = {
+                    "x": x1, "y": y1, 
+                    "width": x2 - x1, 
+                    "height": y2 - y1,
+                    "confidence": confidence
+                }
+                
+                # Face recognition
+                if recognize and face_db and DEEPFACE_AVAILABLE:
+                    # Crop face region
+                    face_crop = image[y1:y2, x1:x2]
+                    if face_crop.size > 0:
+                        recognition_result = face_db.recognize_face(face_crop)
+                        face_info.update(recognition_result)
+                        
+                        # Draw bounding box (green if known, red if unknown)
+                        color = (0, 255, 0) if recognition_result["match"] else (0, 0, 255)
+                        cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 3)
+                        
+                        # Draw label with ID and name
+                        if recognition_result["match"]:
+                            label = f"{recognition_result['name']} ({recognition_result['person_id']})"
+                        else:
+                            label = "Unknown"
+                        
+                        # Text background
+                        (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
+                        cv2.rectangle(annotated_image, (x1, y1 - h - 10), (x1 + w, y1), color, -1)
+                        cv2.putText(annotated_image, label, (x1, y1 - 5), 
+                                  cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
+                else:
+                    # Simple detection without recognition
                     cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 3)
-                    
-                    # Simplified label - just "Face" without percentage
                     label = "Face"
                     (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
                     cv2.rectangle(annotated_image, (x1, y1 - h - 10), (x1 + w, y1), (0, 255, 0), -1)
-                    cv2.putText(annotated_image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)
+                    cv2.putText(annotated_image, label, (x1, y1 - 5), 
+                              cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)
+                
+                faces.append(face_info)
         
         return faces, annotated_image
 
 GLOBAL_DETECTOR: Optional[YOLOv8FaceDetector] = None
 
 def get_global_detector() -> YOLOv8FaceDetector:
-    """Initializes and returns the global singleton detector instance."""
     global GLOBAL_DETECTOR
     if GLOBAL_DETECTOR is None:
         GLOBAL_DETECTOR = YOLOv8FaceDetector(DetectionConfig())
     return GLOBAL_DETECTOR
 
 # ====================================================
-# CORE APPLICATION LOGIC
-# ====================================================
-class FacePrivacyApp:
-    def __init__(self, config: AppConfig, detector: YOLOv8FaceDetector):
-        self.config = config
-        self.blur_effect = get_blur_effect(config.blur)
-        self.detector = detector
-
-    def _expand_bbox(self, bbox: Dict[str, Any], img_shape: Tuple[int, int]) -> Tuple[int, int, int, int]:
-        """Expands a bounding box to include margins for better coverage."""
-        h_img, w_img = img_shape
-        new_w = int(bbox["width"] * self.config.scaling_factor)
-        new_h = int(bbox["height"] * self.config.scaling_factor)
-        x_offset = (new_w - bbox["width"]) // 2
-        y_offset = (new_h - bbox["height"]) // 2
-        x = max(0, bbox["x"] - x_offset - self.config.face_margin)
-        y = max(0, bbox["y"] - y_offset - self.config.forehead_margin)
-        w = min(w_img - x, new_w + 2 * self.config.face_margin)
-        h = min(h_img - y, new_h + self.config.forehead_margin)
-        return x, y, w, h
-
-    def process_image(self, image: np.ndarray, conf_threshold: float) -> np.ndarray:
-        """Applies blur to all detected faces in an image."""
-        writable_image = image.copy()
-        faces, _ = self.detector.detect_faces(writable_image, conf_threshold, return_annotated=False)
-        for face in faces:
-            expanded_roi = self._expand_bbox(face, writable_image.shape[:2])
-            writable_image = self.blur_effect.apply(writable_image, expanded_roi)
-        return writable_image
-
-# ====================================================
-# GRADIO HANDLER FUNCTIONS
+# GRADIO HANDLERS - FACE RECOGNITION
 # ====================================================
-def get_app_instance(blur_type: str, blur_amount: float, blur_size: float) -> FacePrivacyApp:
-    """Creates a FacePrivacyApp instance from UI settings."""
-    detector = get_global_detector()
-    app_config = AppConfig(
-        scaling_factor=blur_size,
-        blur=BlurConfig(type=blur_type, intensity=blur_amount, pixel_size=int(blur_amount))
-    )
-    return FacePrivacyApp(app_config, detector)
-
-def process_media(media, blur_type, blur_amount, blur_size, sensitivity):
-    """Process single image with blur effect."""
-    if media is None: 
-        return None
-    try:
-        # Convert sensitivity to confidence threshold
-        confidence = get_confidence_from_sensitivity(sensitivity)
-        app = get_app_instance(blur_type, blur_amount, blur_size)
-        return app.process_image(media, confidence)
-    except Exception as e:
-        logger.error(f"Image processing error: {e}")
-        gr.Warning(f"An error occurred: {e}")
-        return media
-
-def process_video(video_file, blur_type, blur_amount, blur_size, sensitivity, progress=gr.Progress()):
-    """Process video with blur effect."""
-    if video_file is None: 
-        return None, "⚠️ No video provided."
-    try:
-        # Convert sensitivity to confidence threshold
-        confidence = get_confidence_from_sensitivity(sensitivity)
-        app = get_app_instance(blur_type, blur_amount, blur_size)
-        cap = cv2.VideoCapture(video_file.name)
-        if not cap.isOpened(): 
-            return None, "❌ Cannot open video file."
-
-        out_path = create_temp_file()
-        fourcc = cv2.VideoWriter_fourcc(*'avc1')
-        fps, w, h = cap.get(cv2.CAP_PROP_FPS), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        out_vid = cv2.VideoWriter(out_path, fourcc, fps, (w, h))
-        
-        if not out_vid.isOpened():
-            logger.warning("H.264 codec failed, falling back to mp4v.")
-            out_vid = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
-
-        frame_num = 0
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret: 
-                break
-            frame_num += 1
-            progress(frame_num / max(total_frames, 1), desc=f"Processing frame {frame_num}/{total_frames}")
-            processed_frame = app.process_image(frame, confidence)
-            out_vid.write(processed_frame)
-        
-        cap.release()
-        out_vid.release()
-        return out_path, f"✅ Processed {frame_num} frames."
-    except Exception as e:
-        logger.error(f"Video processing error: {e}")
-        gr.Error(f"Video processing failed: {e}")
-        return None, f"❌ Error: {e}"
-
-def detect_faces_image(image, sensitivity):
-    """Detect faces in single image."""
-    if image is None: 
-        return None, "⚠️ No image provided."
+def recognize_faces_image(image, sensitivity):
+    """Recognize faces in a single image."""
+    if image is None:
+        return None, "⚠️ No image provided"
+    
     try:
-        # Convert sensitivity to confidence threshold
         confidence = get_confidence_from_sensitivity(sensitivity)
         detector = get_global_detector()
-        faces, annotated_image = detector.detect_faces(image, confidence, return_annotated=True)
+        faces, annotated = detector.detect_faces(image, confidence, recognize=True)
         
-        # Simplified result - just show count
         if faces:
-            result = f"✅ **{len(faces)} face(s) detected!**"
+            known = sum(1 for f in faces if f.get("match", False))
+            unknown = len(faces) - known
+            result = f"✅ **Detected {len(faces)} face(s):** {known} known, {unknown} unknown"
         else:
-            result = "❌ **No faces detected.**"
+            result = "❌ **No faces detected**"
         
-        return annotated_image, result
+        return annotated, result
     except Exception as e:
-        logger.error(f"Image detection error: {e}")
-        gr.Warning(f"An error occurred: {e}")
+        logger.error(f"Recognition error: {e}")
         return image, f"❌ Error: {e}"
 
-def detect_faces_video(video_file, sensitivity, progress=gr.Progress()):
-    """Detect faces in video."""
-    if video_file is None: 
-        return None, "⚠️ No video provided."
-    try:
-        # Convert sensitivity to confidence threshold
-        confidence = get_confidence_from_sensitivity(sensitivity)
-        detector = get_global_detector()
-        cap = cv2.VideoCapture(video_file.name)
-        if not cap.isOpened(): 
-            return None, "❌ Cannot open video file."
-
-        out_path = create_temp_file()
-        fourcc = cv2.VideoWriter_fourcc(*'avc1')
-        fps, w, h = cap.get(cv2.CAP_PROP_FPS), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        out_vid = cv2.VideoWriter(out_path, fourcc, fps, (w, h))
-        
-        if not out_vid.isOpened(): 
-            out_vid = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
-
-        frame_num, frames_with_faces = 0, 0
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret: 
-                break
-            frame_num += 1
-            progress(frame_num / max(total_frames, 1), desc=f"Analyzing frame {frame_num}/{total_frames}")
-            faces, annotated_frame = detector.detect_faces(frame, confidence, return_annotated=True)
-            if faces:
-                frames_with_faces += 1
-            out_vid.write(annotated_frame)
-        
-        cap.release()
-        out_vid.release()
-        
-        # Simplified result - just show frame count
-        if frames_with_faces > 0:
-            result = f"✅ **Faces detected in {frames_with_faces}/{frame_num} frames!**"
-        else:
-            result = f"❌ **No faces detected in {frame_num} frames.**"
-        
-        return out_path, result
-    except Exception as e:
-        logger.error(f"Video detection error: {e}")
-        gr.Error(f"Video detection failed: {e}")
-        return None, f"❌ Error: {e}"
-
-def detect_faces_webcam(image, sensitivity):
-    """Detect faces in webcam stream."""
-    if image is None: 
+def recognize_faces_webcam(image, sensitivity):
+    """Recognize faces in webcam stream."""
+    if image is None:
         return None
+    
     try:
-        # Convert sensitivity to confidence threshold
         confidence = get_confidence_from_sensitivity(sensitivity)
         detector = get_global_detector()
-        _, annotated_image = detector.detect_faces(image, confidence, return_annotated=True)
-        return annotated_image
+        _, annotated = detector.detect_faces(image, confidence, recognize=True)
+        return annotated
     except Exception as e:
-        logger.error(f"Webcam detection error: {e}")
+        logger.error(f"Webcam recognition error: {e}")
         return image
 
+# ====================================================
+# DATABASE MANAGEMENT HANDLERS
+# ====================================================
+def add_person_to_db(person_id, name, image):
+    """Add a new person to the face database."""
+    if not person_id or not name:
+        return None, "⚠️ Please provide both ID and Name"
+    
+    if image is None:
+        return None, "⚠️ Please upload a face image"
+    
+    db = get_face_database()
+    success, message = db.add_person(person_id, name, image)
+    
+    if success:
+        return db.list_all_people(), message
+    else:
+        return db.list_all_people(), message
+
+def remove_person_from_db(person_id):
+    """Remove a person from the database."""
+    if not person_id:
+        return None, "⚠️ Please provide a person ID"
+    
+    db = get_face_database()
+    success, message = db.remove_person(person_id)
+    return db.list_all_people(), message
+
+def refresh_database_list():
+    """Refresh the list of known people."""
+    db = get_face_database()
+    return db.list_all_people()
+
 # ====================================================
 # GRADIO UI
 # ====================================================
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Face Privacy Tool") as demo:
-    gr.Markdown("# 🔐 Face Privacy Tool")
-    gr.Markdown("AI-powered face detection and privacy protection using YOLOv8-Face. Obscure faces in images, videos, and live webcam feeds, or use the detection mode to verify their presence.")
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Face Recognition Tool") as demo:
+    gr.Markdown("# 🎭 Face Recognition Tool")
+    gr.Markdown("AI-powered face detection and recognition using YOLOv8 + DeepFace. Detect faces and identify known individuals.")
+    
+    if not DEEPFACE_AVAILABLE:
+        gr.Markdown("⚠️ **Warning:** DeepFace not installed. Install with: `pip install deepface`")
 
     with gr.Row():
-        # ========== SETTINGS SIDEBAR (WRAPPED IN FRAME) ==========
+        # ========== SETTINGS SIDEBAR ==========
         with gr.Column(scale=1):
-            with gr.Column(scale=1, variant="panel"):
-                gr.Markdown("### ⚙️ Global Settings")
-                
-                with gr.Accordion("Privacy Settings", open=True):
-                    blur_type = gr.Radio(
-                        ["gaussian", "pixelate", "solid"], 
-                        value="pixelate", 
-                        label="Blur Type", 
-                        info="Choose how to obscure faces."
-                    )
-                    blur_amount = gr.Slider(
-                        1, 100, 
-                        step=1, 
-                        value=15, 
-                        label="Blur Intensity/Size", 
-                        info="Higher = more obscured."
-                    )
-                    blur_size = gr.Slider(
-                        1.0, 2.0, 
-                        step=0.05, 
-                        value=1.1, 
-                        label="Coverage Area", 
-                        info="Expand blur beyond face boundary."
-                    )
+            with gr.Column(variant="panel"):
+                gr.Markdown("### ⚙️ Settings")
                 
                 with gr.Accordion("Detection Settings", open=True):
-                    # Changed from Slider to Radio for sensitivity
                     detection_sensitivity = gr.Radio(
                         choices=list(SENSITIVITY_MAP.keys()),
                         value="Balanced (Default)",
-                        label="Detection Sensitivity",
-                        info="How strict the face detection should be"
+                        label="Detection Sensitivity"
                     )
 
-        # ========== MAIN CONTENT AREA ==========
+        # ========== MAIN CONTENT ==========
         with gr.Column(scale=2):
             with gr.Tabs():
-                # ========== PRIVACY MODE ==========
-                with gr.TabItem("🔒 Privacy Mode (Blur Faces)"):
-                    gr.Markdown("### Apply privacy protection to your media.")
+                # ========== FACE RECOGNITION TAB ==========
+                with gr.TabItem("🎭 Face Recognition"):
+                    gr.Markdown("### Identify known faces in your images")
                     
                     with gr.Tabs():
-                        # Image Tab
+                        # Image Recognition
                         with gr.TabItem("📷 Image"):
                             with gr.Row():
-                                img_in_blur = gr.Image(
-                                    sources=["upload", "clipboard"], 
-                                    type="numpy", 
+                                img_in_recog = gr.Image(
+                                    sources=["upload", "clipboard"],
+                                    type="numpy",
                                     label="Input Image",
-                                    height=500,
-                                )
-                                img_out_blur = gr.Image(
-                                    type="numpy", 
-                                    label="Protected Image",
-                                    height=500,
-                                )
-                            with gr.Row():
-                                blur_img_btn = gr.Button("Apply Privacy Blur", variant="primary", scale=3)
-                                gr.ClearButton([img_in_blur, img_out_blur], scale=1)
-                            gr.Examples(
-                                examples=[
-                                    ["./examples/single_face.jpg"],
-                                    ["./examples/two_faces.png"],
-                                    ["./examples/group_photo.png"],
-                                    ["./examples/group2.webp"],
-                                ],
-                                inputs=img_in_blur,
-                                label="Click an example to try"
-                            )
-                        
-                        # Video Tab
-                        with gr.TabItem("🎥 Video"):
-                            with gr.Row():
-                                vid_in_blur = gr.File(
-                                    file_types=[".mp4", ".mov", ".avi"], 
-                                    label="Input Video"
+                                    height=500
                                 )
                                 with gr.Column():
-                                    vid_out_blur = gr.Video(
-                                        label="Protected Video",
+                                    img_out_recog = gr.Image(
+                                        type="numpy",
+                                        label="Recognition Result",
                                         height=500
                                     )
-                                    vid_status_blur = gr.Markdown("")
+                                    img_status_recog = gr.Markdown("_Upload an image to start._")
+                            
                             with gr.Row():
-                                blur_vid_btn = gr.Button("Process Video", variant="primary", scale=3)
-                                gr.ClearButton([vid_in_blur, vid_out_blur, vid_status_blur], scale=1)
-
-                        # Webcam Tab
+                                recog_img_btn = gr.Button("Recognize Faces", variant="primary", scale=3)
+                                gr.ClearButton([img_in_recog, img_out_recog, img_status_recog], scale=1)
+                        
+                        # Webcam Recognition
                         with gr.TabItem("📹 Webcam"):
-                            gr.Markdown("**Real-time privacy protection from your webcam feed.**")
+                            gr.Markdown("**Real-time face recognition from webcam**")
                             with gr.Row():
-                                web_in_blur = gr.Image(
-                                    sources=["webcam"], 
-                                    type="numpy", 
-                                    streaming=True, 
+                                web_in_recog = gr.Image(
+                                    sources=["webcam"],
+                                    type="numpy",
+                                    streaming=True,
                                     label="Live Webcam",
-                                    height=500,
+                                    height=500
                                 )
-                                web_out_blur = gr.Image(
-                                    type="numpy", 
-                                    label="Protected Feed",
-                                    height=500,
+                                web_out_recog = gr.Image(
+                                    type="numpy",
+                                    label="Recognition Feed",
+                                    height=500
                                 )
-
-                # ========== DETECTION MODE ==========
-                with gr.TabItem("🔍 Detection Mode (Check for Faces)"):
-                    gr.Markdown("### Verify if your media contains human faces.")
+                
+                # ========== DATABASE MANAGEMENT TAB ==========
+                with gr.TabItem("🗄️ Manage Database"):
+                    gr.Markdown("### Add or remove known people")
                     
-                    with gr.Tabs():
-                        # Image Detection Tab
-                        with gr.TabItem("📷 Image"):
-                            with gr.Row():
-                                img_in_detect = gr.Image(
-                                    sources=["upload", "clipboard"], 
-                                    type="numpy", 
-                                    label="Input Image", 
-                                    height=500,
-                                )
-                                with gr.Column():
-                                    img_out_detect = gr.Image(
-                                        type="numpy", 
-                                        label="Detection Result", 
-                                        height=500,
-                                    )
-                                    img_status_detect = gr.Markdown("_Upload an image to start._")
-                            
-                            with gr.Row():
-                                detect_img_btn = gr.Button("Detect Faces", variant="primary", scale=3)
-                                gr.ClearButton([img_in_detect, img_out_detect, img_status_detect], scale=1)
-                            gr.Examples(
-                                examples=[
-                                    ["./examples/single_face.jpg"],
-                                    ["./examples/two_faces.png"],
-                                    ["./examples/group_photo.png"],
-                                    ["./examples/group2.webp"]
-                                ],
-                                inputs=img_in_detect,
-                                label="Click an example to try"
+                    with gr.Row():
+                        with gr.Column():
+                            gr.Markdown("#### ➕ Add New Person")
+                            add_person_id = gr.Textbox(label="Person ID", placeholder="e.g., EMP001")
+                            add_person_name = gr.Textbox(label="Name", placeholder="e.g., John Doe")
+                            add_person_image = gr.Image(
+                                sources=["upload", "webcam"],
+                                type="numpy",
+                                label="Face Image"
                             )
-
-                        # Video Detection Tab
-                        with gr.TabItem("🎥 Video"):
-                            with gr.Row():
-                                vid_in_detect = gr.File(
-                                    file_types=[".mp4", ".mov", ".avi"], 
-                                    label="Input Video"
-                                )
-                                with gr.Column():
-                                    vid_out_detect = gr.Video(
-                                        label="Annotated Video",
-                                        height=500
-                                    )
-                                    vid_status_detect = gr.Markdown("_Upload a video to start._")
-                            with gr.Row():
-                                detect_vid_btn = gr.Button("Analyze Video for Faces", variant="primary", scale=3)
-                                gr.ClearButton([vid_in_detect, vid_out_detect, vid_status_detect], scale=1)
-
-                        # Webcam Detection Tab
-                        with gr.TabItem("📹 Webcam"):
-                            gr.Markdown("**Live face detection from your webcam feed.**")
-                            with gr.Row():
-                                web_in_detect = gr.Image(
-                                    sources=["webcam"], 
-                                    type="numpy", 
-                                    streaming=True, 
-                                    label="Live Feed",
-                                    height=500,
-                                )
-                                web_out_detect = gr.Image(
-                                    type="numpy", 
-                                    label="Detection Result",
-                                    height=500,
-                                )
+                            add_person_btn = gr.Button("Add Person", variant="primary")
+                            add_status = gr.Markdown("")
+                        
+                        with gr.Column():
+                            gr.Markdown("#### ➖ Remove Person")
+                            remove_person_id = gr.Textbox(label="Person ID to Remove", placeholder="e.g., EMP001")
+                            remove_person_btn = gr.Button("Remove Person", variant="stop")
+                            remove_status = gr.Markdown("")
+                    
+                    gr.Markdown("---")
+                    with gr.Row():
+                        refresh_btn = gr.Button("🔄 Refresh List", scale=1)
+                        database_list = gr.Markdown(value="Loading...")
 
     # ========== EVENT HANDLERS ==========
-    # Privacy Mode (updated to use detection_sensitivity instead of detection_confidence)
-    blur_img_btn.click(
-        process_media, 
-        inputs=[img_in_blur, blur_type, blur_amount, blur_size, detection_sensitivity], 
-        outputs=img_out_blur
-    )
-    blur_vid_btn.click(
-        process_video, 
-        inputs=[vid_in_blur, blur_type, blur_amount, blur_size, detection_sensitivity], 
-        outputs=[vid_out_blur, vid_status_blur]
+    # Recognition
+    recog_img_btn.click(
+        recognize_faces_image,
+        inputs=[img_in_recog, detection_sensitivity],
+        outputs=[img_out_recog, img_status_recog]
     )
-    web_in_blur.stream(
-        process_media, 
-        inputs=[web_in_blur, blur_type, blur_amount, blur_size, detection_sensitivity], 
-        outputs=web_out_blur
+    
+    web_in_recog.stream(
+        recognize_faces_webcam,
+        inputs=[web_in_recog, detection_sensitivity],
+        outputs=web_out_recog
     )
-
-    # Detection Mode (updated to use detection_sensitivity instead of detection_confidence)
-    detect_img_btn.click(
-        detect_faces_image, 
-        inputs=[img_in_detect, detection_sensitivity], 
-        outputs=[img_out_detect, img_status_detect]
+    
+    # Database Management
+    add_person_btn.click(
+        add_person_to_db,
+        inputs=[add_person_id, add_person_name, add_person_image],
+        outputs=[database_list, add_status]
     )
-    detect_vid_btn.click(
-        detect_faces_video, 
-        inputs=[vid_in_detect, detection_sensitivity], 
-        outputs=[vid_out_detect, vid_status_detect]
+    
+    remove_person_btn.click(
+        remove_person_from_db,
+        inputs=[remove_person_id],
+        outputs=[database_list, remove_status]
     )
-    web_in_detect.stream(
-        detect_faces_webcam, 
-        inputs=[web_in_detect, detection_sensitivity], 
-        outputs=web_out_detect
+    
+    refresh_btn.click(
+        refresh_database_list,
+        outputs=database_list
     )
+    
+    # Load database list on startup
+    demo.load(refresh_database_list, outputs=database_list)
 
 # ====================================================
-# MAIN ENTRY POINT
+# MAIN
 # ====================================================
 if __name__ == "__main__":
-    logger.info("🚀 Initializing Face Privacy Tool...")
+    logger.info("🚀 Starting Face Recognition Tool...")
     try:
         get_global_detector()
-        logger.info("✅ Systems ready. Launching Gradio interface...")
+        get_face_database()
+        logger.info("✅ Systems ready. Launching...")
         demo.launch()
     except Exception as e:
-        logger.error(f"❌ Startup failed: {e}")
-        logger.info("💡 Make sure 'yolov8n-face.pt' is available in the current directory or will be downloaded automatically by ultralytics.")
-        
+        logger.error(f"❌ Startup failed: {e}")