app.py

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

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'
STUDENT_DB = "students.db"

# 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))

# Check if roll exists
def roll_exists(rno):
    with sqlite3.connect(STUDENT_DB) as conn:
        cur = conn.execute("SELECT * FROM students WHERE rno=?", (rno,))
        return cur.fetchone()

# Register new student
def register_student(rno, sname, sclass, image):
    image_path = os.path.join(data_path, f"{rno}.jpg")
    image.save(image_path)
    with sqlite3.connect(STUDENT_DB) as conn:
        conn.execute("INSERT INTO students (rno, sname, sclass, image_path) VALUES (?, ?, ?, ?)",
                     (rno, sname, sclass, image_path))


# Load all registered images
def load_registered_images():
    with sqlite3.connect(STUDENT_DB) as conn:
        return [row[3] for row in conn.execute("SELECT * FROM students")]
        
# 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)
    if not webcam_faces:
        return [], cv2_webcam
    
    results = []
    for webcam_face in webcam_faces:
        webcam_emb = torch.tensor(webcam_face.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)
        best_match_idx = torch.argmax(similarity_scores)
        best_score = similarity_scores[best_match_idx].item()
        
        # Get coordinates
        x1, y1, x2, y2 = [int(i) for i in webcam_face.bbox]
        if best_score >= 0.6:
            matched_name = os.path.basename(prod_path[best_match_idx]).split('.')[0]
        else:
            matched_name = "Unknown"
        results.append({'bbox': (x1, y1, x2, y2), 'name': matched_name})
        
    
    return results, cv2_webcam

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

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:
        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)
                matches, image_bgr = prod_function(app, image_paths, image_pil)
    
                if not matches:
                    st.warning("No face detected in the captured image.")
                else:
                    # st.write("Similarity Scores:", prediction_scores)
                    for match in matches:
                        x1, y1, x2, y2 = match['bbox']
                        matched_name = match['name']
                        color = (0, 255, 0) if matched_name != "Unknown" else (0, 0, 255)
                        cv2.rectangle(image_bgr, (x1, y1), (x2, y2), color, 2)
                        cv2.putText(image_bgr, matched_name, (x1, y2 + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
        
                        if matched_name != "Unknown":
                        # recognized = False
                        # for score, idx in matches:
                        #     if matched_score >= 0.6:
                        #         matched_name = os.path.basename(image_paths[match_idx]).split('.')[0]
                        #         st.success(f"✅ Welcome: {matched_name}")
                        #         recognizes=True
                            
                            # Send attendance via POST
                            url = "https://nielit-attendance.glitch.me/adds"
                            data = {'rno': 15, 'sname': matched_name, 'sclass': 7}
                            try:
                                response = requests.post(url, data=data)
                                if response.status_code == 200:
                                    st.success(f"Attendance marked successfully for {matched_name}.")
                                else:
                                    st.warning(f"Failed to update attendance for {matched_name}.")
                            except Exception as e:
                                st.error(f"Request failed: {e}")
                        else:
                            st.warning("❌ Face match not found or too low confidence.")
                    
                # Convert back to RGB for display
                image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
                st.image(image_rgb, caption="Detected Faces", use_container_width=True)

with register_tab:
    upload_mode = st.radio("Choose Method", ["📤 Upload Image", "📷 Take Photo"])
    image_pil = None
    if upload_mode == "📤 Upload Image":
        upload = st.file_uploader("Upload Student Image with plain background(.jpg)", type="jpg")
        if upload:
            image_pil = Image.open(upload)
    else:
        camera_capture = st.camera_input("Capture Image")
        if camera_capture:
            image_pil = Image.open(camera_capture)

    if 'image_pil' in locals():
        st.image(image_pil, caption="Student Image", use_container_width=True)
        rno = st.text_input("Roll Number")
        sname = st.text_input("Student Name")
        sclass = st.text_input("Class")

        if st.button("Register"):
            if not (rno and sname and sclass):
                st.warning("Fill all details.")
            # Save locally in employees folder
                local_path = os.path.join("employees", f"{sname}.jpg")
                image_pil.save(local_path)

                # Convert image to base64 for sending to Glitch
                buffered = BytesIO()
                image_pil.save(buffered, format="JPEG")
                img_str = base64.b64encode(buffered.getvalue()).decode()

                # Prepare data for POST request
                glitch_url = "https://your-glitch-app.glitch.me/register"  # Replace this
                data = {
                    "rno": rno,
                    "sname": sname,
                    "sclass": sclass,
                    "image_base64": img_str,
                }

                response = requests.post(glitch_url, json=data)
                result = response.json()

                if result["status"] == "success":
                    st.success(result["message"])
                else:
                    st.error(result["message"])