app.py

import os
import cv2
import torch
import numpy as np
import streamlit as st
import requests

from PIL import Image
from glob import glob
from insightface.app import FaceAnalysis
import torch.nn.functional as F

# Set the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Global Variables
IMAGE_SHAPE = 640
data_path = 'employees'
webcam_path = 'captured_image.jpg'

# Set Streamlit title
st.title("AIML-Student Attendance System")

# Load employee image paths
image_paths = glob(os.path.join(data_path, '*.jpg'))

# Initialize Face Analysis
app = FaceAnalysis(name="buffalo_l")  # ArcFace model
app.prepare(ctx_id=0 if torch.cuda.is_available() else -1, det_size=(IMAGE_SHAPE, IMAGE_SHAPE))

# Define function to match face embeddings
def prod_function(app, prod_path, webcam_img_pil):
    np_webcam = np.array(webcam_img_pil)
    cv2_webcam = cv2.cvtColor(np_webcam, cv2.COLOR_RGB2BGR)
    
    webcam_faces = app.get(cv2_webcam, max_num=1)
    if not webcam_faces:
        return None, None
    
    webcam_emb = torch.tensor(webcam_faces[0].embedding, dtype=torch.float32)

    similarity_scores = []
    for path in prod_path:
        img = cv2.imread(path)
        faces = app.get(img, max_num=1)
        if not faces:
            similarity_scores.append(torch.tensor(-1.0))
            continue

        face_emb = torch.tensor(faces[0].embedding, dtype=torch.float32)
        score = F.cosine_similarity(face_emb, webcam_emb, dim=0)
        similarity_scores.append(score)

    similarity_scores = torch.stack(similarity_scores)
    return similarity_scores, torch.argmax(similarity_scores)

# Streamlit tabs
about_tab, app_tab = st.tabs(["About the app", "Face Recognition"])

with about_tab:
    st.markdown("""
    # 👁️‍🗨️ AI-Powered Face Recognition Attendance System
    Secure and Accurate Attendance using Vision Transformer + ArcFace Embeddings.
    
    - **Automated, contactless attendance logging**
    - **Uses InsightFace ArcFace embeddings for recognition**
    - **Real-time logging with confidence scoring**
    - **Future Scope: Mask-aware recognition, Group detection, and more**
    """)

with app_tab:
    trained_names = [os.path.basename(p).split('.')[0] for p in image_paths]
    st.subheader("📋 Trained Faces in System")
    st.write(", ".join(trained_names) if trained_names else "No faces found.")
    enable = st.checkbox("Enable camera")
    picture = st.camera_input("Take a picture", disabled=not enable)

    if picture is not None:
        with st.spinner("Analyzing face..."):
            image_pil = Image.open(picture)
            np_webcam = np.array(image_pil)
            cv2_webcam = cv2.cvtColor(np_webcam, cv2.COLOR_RGB2BGR)
    
            webcam_faces = app.get(cv2_webcam, max_num=1)

            if not webcam_faces:
                st.warning("No face detected in the captured image.")
            else:
                webcam_emb = torch.tensor(webcam_faces[0].embedding, dtype=torch.float32)

                similarity_scores = []
                for path in image_paths:
                    img = cv2.imread(path)
                    faces = app.get(img, max_num=1)
                    if not faces:
                        similarity_scores.append(torch.tensor(-1.0))
                        continue
    
                    face_emb = torch.tensor(faces[0].embedding, dtype=torch.float32)
                    score = F.cosine_similarity(face_emb, webcam_emb, dim=0)
                    similarity_scores.append(score)
    
                similarity_scores = torch.stack(similarity_scores)
                match_idx = torch.argmax(similarity_scores)
                matched_score = similarity_scores[match_idx].item()
    
                # Draw bounding box and name
                (x1, y1, x2, y2) = map(int, webcam_faces[0].bbox)
                cv2.rectangle(cv2_webcam, (x1, y1), (x2, y2), (0, 255, 0), 2)

                if matched_score >= 0.5:
                    matched_name = os.path.basename(image_paths[match_idx]).split('.')[0]
                    cv2.putText(cv2_webcam, matched_name, (x1, y1 - 10), 
                                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    
                    st.success(f"✅ Welcome: {matched_name}")

                # Send attendance via POST
                    url = "https://aiml2025.glitch.me/attend"
                    data = {'rollno': 15, 'Name': matched_name, 'Class': 7}
                    try:
                        response = requests.post(url, data=data)
                        if response.status_code == 200:
                            st.success("Attendance marked successfully.")
                        else:
                            st.warning("Failed to update attendance.")
                    except Exception as e:
                        st.error(f"Request failed: {e}")
                else:
                    st.error("❌ Match not found. Try again.")
    
                # Convert back to RGB for displaying in Streamlit
                final_img = cv2.cvtColor(cv2_webcam, cv2.COLOR_BGR2RGB)
                st.image(final_img, caption="Detected Face")