app.py

import gradio as gr
import sqlite3
import numpy as np
import cv2
from PIL import Image
import io
import os

# Global variable for FaceNet availability
FACE_NET_AVAILABLE = False

# Try to import FaceNet with error handling
try:
    from facenet_pytorch import MTCNN, InceptionResnetV1
    FACE_NET_AVAILABLE = True
    print("✅ FaceNet models imported successfully")
except ImportError as e:
    print(f"❌ FaceNet import error: {e}")
    FACE_NET_AVAILABLE = False

class FaceRecognitionSystem:
    def __init__(self):
        self.setup_database()
        global FACE_NET_AVAILABLE
        
        if FACE_NET_AVAILABLE:
            try:
                # Initialize MTCNN with simpler settings
                self.mtcnn = MTCNN(
                    keep_all=False,  # Only detect one face for simplicity
                    min_face_size=40,
                    thresholds=[0.6, 0.7, 0.7],
                    device='cpu'
                )
                self.resnet = InceptionResnetV1(pretrained='vggface2', classify=False, device='cpu').eval()
                print("✅ FaceNet models loaded successfully")
            except Exception as e:
                print(f"❌ Error loading FaceNet models: {e}")
                FACE_NET_AVAILABLE = False
    
    def setup_database(self):
        """Initialize SQLite database for students and attendance"""
        self.conn = sqlite3.connect('face_attendance.db', check_same_thread=False)
        self.cursor = self.conn.cursor()
        
        # Students table with face embeddings
        self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS students (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                student_id TEXT UNIQUE NOT NULL,
                class TEXT,
                face_embedding BLOB,
                registration_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # Attendance table
        self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS attendance (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                student_id TEXT,
                name TEXT,
                class TEXT,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                status TEXT DEFAULT 'Present'
            )
        ''')
        self.conn.commit()
        print("✅ Database initialized successfully")
    
    def extract_face_embedding(self, image):
        """
        Extract face embedding from image using FaceNet
        """
        global FACE_NET_AVAILABLE
        
        if not FACE_NET_AVAILABLE:
            return None, "Face recognition not available"
        
        try:
            # Convert image to PIL format
            if isinstance(image, str):
                # File path
                pil_image = Image.open(image).convert('RGB')
            elif isinstance(image, np.ndarray):
                # numpy array
                if len(image.shape) == 3:
                    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
                else:
                    return None, "Invalid image format"
            else:
                # Handle Gradio file object
                pil_image = Image.open(image).convert('RGB')
            
            # Detect face and get cropped face
            face_tensor, prob = self.mtcnn(pil_image, return_prob=True)
            
            if face_tensor is None:
                return None, "No face detected in the image"
            
            if prob < 0.9:  # Confidence threshold
                return None, "Low confidence face detection"
            
            # Ensure proper tensor dimensions [3, 160, 160]
            if face_tensor.dim() == 4:
                face_tensor = face_tensor.squeeze(0)  # Remove batch dimension if present
            
            # Get embedding
            with torch.no_grad():
                embedding = self.resnet(face_tensor.unsqueeze(0))  # Add batch dimension
                embedding_np = embedding.detach().numpy().flatten()
            
            return embedding_np, "Face embedding extracted successfully"
            
        except Exception as e:
            return None, f"Error processing image: {str(e)}"
    
    def register_student(self, name, student_id, class_name, image):
        """
        Register a new student with face embedding
        """
        if not name or not student_id:
            return "❌ Please provide name and student ID"
        
        if image is None:
            return "❌ Please upload a face image"
        
        # Extract face embedding
        embedding, message = self.extract_face_embedding(image)
        if embedding is None:
            return f"❌ {message}"
        
        # Store in database
        try:
            self.cursor.execute('''
                INSERT INTO students (name, student_id, class, face_embedding)
                VALUES (?, ?, ?, ?)
            ''', (name, student_id, class_name, embedding.tobytes()))
            
            self.conn.commit()
            return f"✅ Student {name} registered successfully with face recognition!"
            
        except sqlite3.IntegrityError:
            return f"❌ Student ID {student_id} already exists"
        except Exception as e:
            return f"❌ Database error: {str(e)}"
    
    def recognize_face(self, image):
        """
        Recognize face from image and mark attendance
        """
        global FACE_NET_AVAILABLE
        
        if not FACE_NET_AVAILABLE:
            return image, "Face recognition not available"
        
        if image is None:
            return None, "Please upload an image"
        
        # Extract face embedding from input image
        embedding, message = self.extract_face_embedding(image)
        if embedding is None:
            return image, f"❌ {message}"
        
        # Get all registered students
        self.cursor.execute('SELECT student_id, name, class, face_embedding FROM students')
        students = self.cursor.fetchall()
        
        if not students:
            return image, "❌ No students registered in the system"
        
        # Find the best match
        best_match = None
        min_distance = float('inf')
        threshold = 1.0  # Similarity threshold
        
        for student_id, name, class_name, embedding_blob in students:
            known_embedding = np.frombuffer(embedding_blob, dtype=np.float32)
            
            # Calculate Euclidean distance
            distance = np.linalg.norm(embedding - known_embedding)
            
            if distance < min_distance and distance < threshold:
                min_distance = distance
                best_match = (student_id, name, class_name, distance)
        
        if best_match:
            student_id, name, class_name, distance = best_match
            
            # Check if already marked today
            self.cursor.execute('''
                SELECT * FROM attendance 
                WHERE student_id = ? AND DATE(timestamp) = DATE('now')
            ''', (student_id,))
            
            existing = self.cursor.fetchone()
            
            if not existing:
                # Mark attendance
                self.cursor.execute('''
                    INSERT INTO attendance (student_id, name, class, status)
                    VALUES (?, ?, ?, 'Present')
                ''', (student_id, name, class_name))
                self.conn.commit()
                
                result_message = f"✅ Attendance marked for {name} (ID: {student_id})"
                result_message += f"\nSimilarity: {1-distance:.3f}"
            else:
                result_message = f"⚠️ {name} already marked today (ID: {student_id})"
            
            # Convert image for display
            display_image = self.prepare_display_image(image)
            
            # Add text to image
            if display_image is not None:
                cv2.putText(display_image, f"Recognized: {name}", (10, 30), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                cv2.putText(display_image, f"ID: {student_id}", (10, 60), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
                cv2.putText(display_image, f"Class: {class_name}", (10, 90), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
                cv2.putText(display_image, f"Similarity: {1-distance:.3f}", (10, 120), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
                
                return display_image, result_message
            
        return image, "❌ No matching student found"
    
    def prepare_display_image(self, image):
        """Convert image to display format"""
        try:
            if isinstance(image, str):
                display_image = cv2.imread(image)
            elif isinstance(image, np.ndarray):
                display_image = image.copy()
                if len(display_image.shape) == 3 and display_image.shape[2] == 3:
                    display_image = cv2.cvtColor(display_image, cv2.COLOR_RGB2BGR)
            else:
                # Handle Gradio file object
                display_image = np.array(Image.open(image))
                if len(display_image.shape) == 3 and display_image.shape[2] == 3:
                    display_image = cv2.cvtColor(display_image, cv2.COLOR_RGB2BGR)
            
            return display_image
        except:
            return None
    
    def get_all_students(self):
        """Get all registered students"""
        self.cursor.execute('SELECT name, student_id, class FROM students ORDER BY name')
        students = self.cursor.fetchall()
        return students
    
    def get_attendance_records(self, time_range="Today"):
        """Get attendance records"""
        if time_range == "Today":
            self.cursor.execute('''
                SELECT name, student_id, class, timestamp 
                FROM attendance 
                WHERE DATE(timestamp) = DATE('now')
                ORDER BY timestamp DESC
            ''')
        else:
            self.cursor.execute('''
                SELECT name, student_id, class, timestamp 
                FROM attendance 
                ORDER BY timestamp DESC
                LIMIT 50
            ''')
        return self.cursor.fetchall()

# Import torch after FaceNet imports to avoid conflicts
try:
    import torch
except ImportError:
    torch = None

# Initialize the system
face_system = FaceRecognitionSystem()

def register_student_interface(name, student_id, class_name, image):
    return face_system.register_student(name, student_id, class_name, image)

def recognize_face_interface(image):
    return face_system.recognize_face(image)

def view_students_interface():
    students = face_system.get_all_students()
    if not students:
        return "No students registered yet"
    
    result = "## 📋 Registered Students\n\n"
    for name, student_id, class_name in students:
        result += f"**{name}**\n"
        result += f"- ID: {student_id}\n"
        result += f"- Class: {class_name}\n\n"
    return result

def view_attendance_interface(time_range):
    records = face_system.get_attendance_records(time_range)
    if not records:
        return f"No attendance records for {time_range.lower()}"
    
    result = f"## 📊 {time_range}'s Attendance\n\n"
    for name, student_id, class_name, timestamp in records:
        time_str = timestamp[:16] if timestamp else "Unknown"
        result += f"**{name}** ({student_id})\n"
        result += f"- Class: {class_name}\n"
        result += f"- Time: {time_str}\n\n"
    return result

# Create Gradio interface
with gr.Blocks(title="Face Recognition Attendance System") as demo:
    gr.Markdown("# 🎯 Face Recognition Attendance System")
    gr.Markdown("AI-powered facial recognition for automated attendance marking")
    
    with gr.Tab("👤 Register Student"):
        gr.Markdown("### Register New Student with Face Recognition")
        with gr.Row():
            with gr.Column():
                name_input = gr.Textbox(
                    label="Full Name", 
                    placeholder="Enter student's full name"
                )
                student_id_input = gr.Textbox(
                    label="Student ID", 
                    placeholder="Enter unique student ID"
                )
                class_input = gr.Textbox(
                    label="Class", 
                    placeholder="Enter class name"
                )
                image_input = gr.Image(
                    label="Upload Face Image",
                    type="filepath"
                )
                register_btn = gr.Button("Register Student with Face", variant="primary")
            
            with gr.Column():
                register_output = gr.Textbox(
                    label="Registration Status", 
                    lines=5
                )
        
        register_btn.click(
            register_student_interface,
            inputs=[name_input, student_id_input, class_input, image_input],
            outputs=register_output
        )
    
    with gr.Tab("📷 Mark Attendance"):
        gr.Markdown("### Mark Attendance using Face Recognition")
        with gr.Row():
            with gr.Column():
                attendance_image = gr.Image(
                    label="Upload Photo for Attendance",
                    type="filepath"
                )
                recognize_btn = gr.Button("Recognize Faces & Mark Attendance", variant="primary")
            
            with gr.Column():
                processed_image = gr.Image(
                    label="Processed Image"
                )
                recognition_output = gr.Textbox(
                    label="Recognition Results",
                    lines=4
                )
        
        recognize_btn.click(
            recognize_face_interface,
            inputs=[attendance_image],
            outputs=[processed_image, recognition_output]
        )
    
    with gr.Tab("📊 View Records"):
        gr.Markdown("### View Students and Attendance Records")
        with gr.Row():
            with gr.Column():
                gr.Markdown("#### Registered Students")
                students_btn = gr.Button("View All Students", variant="secondary")
                students_output = gr.Markdown()
            
            with gr.Column():
                gr.Markdown("#### Attendance Records")
                time_range = gr.Radio(
                    choices=["Today", "All Time"],
                    value="Today",
                    label="Select Time Range"
                )
                attendance_btn = gr.Button("View Attendance Records", variant="secondary")
                attendance_output = gr.Markdown()
        
        students_btn.click(view_students_interface, outputs=students_output)
        attendance_btn.click(
            view_attendance_interface,
            inputs=[time_range],
            outputs=attendance_output
        )
    
    with gr.Tab("ℹ️ System Info"):
        status_msg = "✅ **System Status: FULLY OPERATIONAL**" if FACE_NET_AVAILABLE else "⚠️ **System Status: FACE RECOGNITION UNAVAILABLE**"
        gr.Markdown(f"## System Status\n{status_msg}")
        
        gr.Markdown("""
        ## About Face Recognition Attendance System
        
        This system uses state-of-the-art facial recognition technology to automate attendance marking.
        
        ### Features:
        - **Face Registration**: Register students with their facial data
        - **Automatic Recognition**: Recognize students from photos and mark attendance
        - **Attendance Management**: View and manage attendance records
        
        ### How to Use:
        1. **Register Students**: Go to 'Register Student' tab, enter details and upload face photo
        2. **Mark Attendance**: Go to 'Mark Attendance' tab, upload photo with faces
        3. **View Records**: Check registered students and attendance records
        
        ### Technology:
        - FaceNet for face recognition
        - Gradio for web interface
        - SQLite for data storage
        """)

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