app.py

import gradio as gr
import numpy as np
from PIL import Image
import os
import json
from datetime import datetime

# Try to import InsightFace, fallback gracefully if not available
INSIGHTFACE_AVAILABLE = False
try:
    from insightface.app.face_analysis import FaceAnalysis
    import onnxruntime as ort
    INSIGHTFACE_AVAILABLE = True
    print("✓ InsightFace available")
except Exception as e:
    print(f"InsightFace not available: {e}")
    print("Will use demo mode")

class FaceMatchingSystem:
    def __init__(self):
        """Initialize the face matching system"""
        self.app = None
        self.face_database = {}
        self.model_status = "Initializing..."
        self.setup_models()
        
    def setup_models(self):
        """Setup the face recognition models"""
        try:
            if INSIGHTFACE_AVAILABLE:
                print("Attempting to load InsightFace models...")
                try:
                    self.app = FaceAnalysis(
                        name='buffalo_l',
                        providers=['CPUExecutionProvider']
                    )
                    self.app.prepare(ctx_id=0, det_thresh=0.5, det_size=(640, 640))
                    self.model_status = "✓ InsightFace models loaded successfully"
                    print(self.model_status)
                    return
                except Exception as e:
                    print(f"Failed to load InsightFace models: {e}")
            
            # Fallback to demo mode
            self.app = MockFaceApp()
            self.model_status = "Demo mode active (InsightFace not available)"
            print(self.model_status)
            
        except Exception as e:
            print(f"Error in model setup: {e}")
            self.app = MockFaceApp()
            self.model_status = f"Demo mode (Error: {str(e)})"
    
    def extract_face_embedding(self, image):
        """Extract face embedding from image"""
        try:
            if image is None:
                return None, "No image provided"
                
            # Convert PIL to numpy array (RGB format)
            if isinstance(image, Image.Image):
                image_array = np.array(image.convert('RGB'))
            else:
                image_array = image
            
            # Use the face analysis app
            if hasattr(self.app, 'get'):
                faces = self.app.get(image_array)
            else:
                return np.random.rand(512), "Demo mode: mock embedding generated"
            
            if len(faces) == 0:
                return None, "No face detected in the image"
            
            # Use the largest face if multiple detected
            if len(faces) > 1:
                faces = sorted(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse=True)
            
            face = faces[0]
            embedding = face.embedding
            confidence = getattr(face, 'det_score', 0.95)
            
            return embedding, f"Face detected (confidence: {confidence:.3f})"
            
        except Exception as e:
            print(f"Error extracting face embedding: {e}")
            return None, f"Error processing image: {str(e)}"
    
    def add_face_to_database(self, image, person_name):
        """Add a face to the database"""
        if not person_name or not person_name.strip():
            return "Please provide a valid person name", ""
        
        person_name = person_name.strip()
        
        embedding, message = self.extract_face_embedding(image)
        if embedding is None:
            return f"Failed to add {person_name}: {message}", ""
        
        # Store embedding in database
        self.face_database[person_name] = {
            'embedding': embedding.tolist() if hasattr(embedding, 'tolist') else embedding,
            'added_at': datetime.now().isoformat()
        }
        
        # Save database
        self.save_database()
        
        return f"✓ Successfully added {person_name} to database ({message})", self.get_database_info()
    
    def match_face(self, image, threshold=0.6):
        """Match a face against the database"""
        if not self.face_database:
            return "Database is empty. Please add faces first.", "", 0.0
        
        embedding, message = self.extract_face_embedding(image)
        if embedding is None:
            return f"Face matching failed: {message}", "", 0.0
        
        best_match = None
        best_similarity = 0.0
        
        for person_name, data in self.face_database.items():
            stored_embedding = np.array(data['embedding'])
            
            # Calculate cosine similarity
            similarity = np.dot(embedding, stored_embedding) / (
                np.linalg.norm(embedding) * np.linalg.norm(stored_embedding)
            )
            
            if similarity > best_similarity:
                best_similarity = similarity
                best_match = person_name
        
        if best_similarity >= threshold:
            confidence_percentage = best_similarity * 100
            return (
                f"✓ Match Found: {best_match}",
                f"Confidence: {confidence_percentage:.1f}%",
                confidence_percentage
            )
        else:
            return (
                "❌ No match found",
                f"Best similarity: {best_similarity*100:.1f}% (below threshold {threshold*100:.1f}%)",
                best_similarity * 100
            )
    
    def save_database(self):
        """Save the face database"""
        try:
            with open('face_database.json', 'w') as f:
                json.dump(self.face_database, f, indent=2)
        except Exception as e:
            print(f"Failed to save database: {e}")
    
    def load_database(self):
        """Load the face database"""
        try:
            if os.path.exists('face_database.json'):
                with open('face_database.json', 'r') as f:
                    self.face_database = json.load(f)
                print(f"Loaded {len(self.face_database)} faces from database")
        except Exception as e:
            print(f"Failed to load database: {e}")
            self.face_database = {}
    
    def get_database_info(self):
        """Get information about the current database"""
        if not self.face_database:
            return "Database is empty"
        
        info = f"Database contains {len(self.face_database)} faces:\\n"
        for name, data in self.face_database.items():
            added_date = data.get('added_at', 'Unknown')[:10]
            info += f"• {name} (added: {added_date})\\n"
        
        return info
    
    def clear_database(self):
        """Clear the entire database"""
        self.face_database = {}
        self.save_database()
        return "Database cleared successfully", ""

class MockFaceApp:
    """Mock face app for demo purposes when InsightFace is not available"""
    def __init__(self):
        self.face_counter = 0
        
    def get(self, image):
        if image is None:
            return []
            
        # Create deterministic embedding based on image hash
        image_hash = hash(str(np.array(image).mean())) % 1000
        
        class MockFace:
            def __init__(self, image_hash):
                np.random.seed(image_hash)
                self.embedding = np.random.rand(512)
                self.embedding = self.embedding / np.linalg.norm(self.embedding)
                self.det_score = 0.85 + (image_hash % 15) / 100
                self.bbox = [50, 50, 200, 200]
        
        return [MockFace(image_hash)]

# Initialize the system
print("Initializing Face Matching System...")
face_system = FaceMatchingSystem()
face_system.load_database()

# Create a simple, robust Gradio interface
with gr.Blocks(title="FaceMatch Pro") as demo:
    
    gr.Markdown("# 🎯 FaceMatch Pro")
    gr.Markdown("### Professional Face Recognition System")
    
    # Status display
    status_display = gr.Textbox(
        label="System Status", 
        value=face_system.model_status,
        interactive=False
    )
    
    with gr.Tabs():
        # Tab 1: Add Face
        with gr.Tab("Add Face"):
            gr.Markdown("### Add a face to the database")
            
            with gr.Row():
                with gr.Column():
                    add_image = gr.Image(label="Upload Photo", type="pil")
                    person_name = gr.Textbox(label="Person Name", placeholder="Enter name...")
                    add_btn = gr.Button("Add to Database", variant="primary")
                
                with gr.Column():
                    add_result = gr.Textbox(label="Result", lines=3)
                    database_info = gr.Textbox(
                        label="Database Info", 
                        lines=6, 
                        value=face_system.get_database_info()
                    )
            
            add_btn.click(
                face_system.add_face_to_database,
                inputs=[add_image, person_name],
                outputs=[add_result, database_info]
            )
        
        # Tab 2: Match Face
        with gr.Tab("Match Face"):
            gr.Markdown("### Find face matches")
            
            with gr.Row():
                with gr.Column():
                    match_image = gr.Image(label="Upload Photo to Match", type="pil")
                    threshold = gr.Slider(
                        minimum=0.3,
                        maximum=0.9,
                        value=0.6,
                        step=0.05,
                        label="Matching Threshold"
                    )
                    match_btn = gr.Button("Find Matches", variant="primary")
                
                with gr.Column():
                    match_result = gr.Textbox(label="Match Result", lines=2)
                    confidence_text = gr.Textbox(label="Confidence Details", lines=2)
                    confidence_score = gr.Number(label="Confidence Score (%)", precision=1)
            
            match_btn.click(
                face_system.match_face,
                inputs=[match_image, threshold],
                outputs=[match_result, confidence_text, confidence_score]
            )
        
        # Tab 3: Database Management
        with gr.Tab("Database"):
            gr.Markdown("### Database Management")
            
            db_stats = gr.Textbox(
                label="Database Contents",
                lines=8,
                value=face_system.get_database_info()
            )
            
            with gr.Row():
                refresh_btn = gr.Button("Refresh Info", variant="secondary")
                clear_btn = gr.Button("Clear Database", variant="stop")
            
            clear_result = gr.Textbox(label="Action Result", lines=2)
            
            refresh_btn.click(
                lambda: face_system.get_database_info(),
                outputs=[db_stats]
            )
            
            clear_btn.click(
                face_system.clear_database,
                outputs=[clear_result, db_stats]
            )

if __name__ == "__main__":
    demo.launch()