Update app.py

app.py

@@ -1,219 +1,150 @@
-import os
-import sqlite3
 import streamlit as st
-from PIL import Image
-import numpy as np
 import cv2
-import dlib
+import os
+import numpy as np
 from keras.models import load_model
-from datetime import datetime
-from huggingface_hub import HfApi
+from PIL import Image
+import sqlite3
+import requests
+from io import BytesIO
 
 # Constants
-KNOWN_FACES_DIR = "known_faces"  # Directory to save user images
-DATABASE = "students.db"  # SQLite database file to store student information
+DB_FILE = "students.db"
+KNOWN_FACES_DIR = "known_faces"
+EMOTION_MODEL_FILE = "CNN_Model_acc_75.h5"
+HUGGING_FACE_TOKEN = os.getenv("HUGGING_FACE_TOKEN")
+HUGGING_FACE_REPO = "username/repo_name"
 
-# Ensure the directory exists
+# Create directories
 os.makedirs(KNOWN_FACES_DIR, exist_ok=True)
 
-# Initialize Hugging Face API
-hf_token = os.getenv("upload")  # Ensure this is set correctly as a secret in Hugging Face
-if not hf_token:
-    raise ValueError("Hugging Face token not found. Ensure it's set as a secret in Hugging Face")
-api = HfApi()
-
-# Repository Details on Hugging Face
-REPO_NAME = "face_and_emotion_detection"  # Replace with your Hugging Face repository name
-REPO_ID = "LovnishVerma/" + REPO_NAME  # Replace "LovnishVerma" with your Hugging Face username
-REPO_TYPE = "space"  # 'space' type for Streamlit-based projects
-
-# Load emotion detection model
-model = load_model('CNN_Model_acc_75.h5')
-emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
-face_detector = dlib.get_frontal_face_detector()  # Use Dlib's face detector
-face_predictor = dlib.shape_predictor('shape_predictor_68_face_landmarks.dat')  # Landmarks model for face recognition
-face_rec_model = dlib.face_recognition_model_v1('dlib_face_recognition_resnet_model_v1.dat')  # Face recognition model
-
-# Initialize the SQLite database
-def initialize_database():
-    conn = sqlite3.connect(DATABASE)
-    cursor = conn.cursor()
-    cursor.execute("""
-        CREATE TABLE IF NOT EXISTS students (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            name TEXT NOT NULL,
-            roll_no TEXT NOT NULL UNIQUE,
-            image_path TEXT NOT NULL,
-            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
-        )
-    """)
-    conn.commit()
-    conn.close()
-
-# Save student information in the SQLite database
-def save_to_database(name, roll_no, image_path):
-    conn = sqlite3.connect(DATABASE)
-    cursor = conn.cursor()
-    try:
+# Load models
+try:
+    emotion_model = load_model(EMOTION_MODEL_FILE)
+except Exception as e:
+    st.error(f"Error loading emotion model: {e}")
+
+# Database Functions
+def create_table():
+    with sqlite3.connect(DB_FILE) as conn:
+        cursor = conn.cursor()
         cursor.execute("""
-            INSERT INTO students (name, roll_no, image_path)
-            VALUES (?, ?, ?)
-        """, (name, roll_no, image_path))
+            CREATE TABLE IF NOT EXISTS students (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                name TEXT NOT NULL,
+                roll_number TEXT NOT NULL UNIQUE,
+                image_path TEXT NOT NULL
+            )
+        """)
         conn.commit()
-        st.success("Data saved successfully!")
+
+def insert_student(name, roll_number, image_path):
+    try:
+        with sqlite3.connect(DB_FILE) as conn:
+            cursor = conn.cursor()
+            cursor.execute("INSERT INTO students (name, roll_number, image_path) VALUES (?, ?, ?)",
+                           (name, roll_number, image_path))
+            conn.commit()
     except sqlite3.IntegrityError:
-        st.error("Roll number already exists!")
-    finally:
-        conn.close()
-
-# Save the captured image locally in known_faces directory and upload to Hugging Face
-def save_image_to_hugging_face(image, name, roll_no):
-    # Create a filename based on the student name and roll number
-    filename = f"{name}_{roll_no}.jpg"
-    local_path = os.path.join(KNOWN_FACES_DIR, filename)
-    
-    # Save the image locally
-    image.save(local_path)
-    st.success(f"Image saved locally to {local_path}")
-    
+        st.warning("Roll number already exists!")
+
+# Hugging Face Functions
+def upload_to_hugging_face(file_path, file_name):
+    if not HUGGING_FACE_TOKEN:
+        st.error("Hugging Face token not found.")
+        return
+
+    url = f"https://huggingface.co/api/repos/{HUGGING_FACE_REPO}/uploads/{file_name}"
+    headers = {"Authorization": f"Bearer {HUGGING_FACE_TOKEN}"}
+
     try:
-        # Upload the image to Hugging Face repository
-        api.upload_file(
-            path_or_fileobj=local_path,
-            path_in_repo=filename,
-            repo_id=REPO_ID,
-            repo_type=REPO_TYPE,
-            token=hf_token
-        )
-        st.success(f"Image uploaded to Hugging Face: {filename}")
+        with open(file_path, "rb") as file:
+            response = requests.post(url, headers=headers, files={"file": file})
+        if response.status_code == 200:
+            st.success(f"Uploaded {file_name} to Hugging Face!")
+        else:
+            st.error(f"Failed to upload: {response.content}")
     except Exception as e:
-        st.error(f"Error uploading image to Hugging Face: {e}")
-    
-    return local_path
-
-# Process each frame for emotion detection and face recognition
-def process_frame(frame, known_face_encodings, known_face_names):
-    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-    faces = face_detector(gray_frame)
-
-    face_locations = []
-    face_encodings = []
-    face_names = []
-    for face in faces:
-        # Get the landmarks of the face
-        shape = face_predictor(gray_frame, face)
-        face_encoding = np.array(face_rec_model.compute_face_descriptor(frame, shape))
-
-        # Check if the detected face matches any known face
-        matches = [np.allclose(face_encoding, known_encoding) for known_encoding in known_face_encodings]
-        name = "Unknown"
-
-        if True in matches:
-            first_match_index = matches.index(True)
-            name = known_face_names[first_match_index]
-
-        # Emotion detection
-        roi_color = frame[face.top():face.bottom(), face.left():face.right()]
-        face_roi = cv2.resize(roi_color, (48, 48))
-        face_roi = np.expand_dims(face_roi, axis=0)
-        face_roi = face_roi / float(48)
-        predictions = model.predict(face_roi)
-        emotion = emotion_labels[np.argmax(predictions[0])]
-
-        # Display name and emotion text on face
-        cv2.rectangle(frame, (face.left(), face.top()), (face.right(), face.bottom()), (0, 255, 0), 2)
-        cv2.putText(frame, f"{name} - {emotion}", (face.left(), face.bottom() + 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-
-    return frame
-
-# User Interface for registration
-st.title("Student Registration and Attendance")
-
-# Initialize the database and create the table if it does not exist
-initialize_database()
-
-# Load known faces from the database
-known_face_encodings = []
-known_face_names = []
-
-conn = sqlite3.connect(DATABASE)
-cursor = conn.cursor()
-cursor.execute("SELECT name, image_path FROM students")
-rows = cursor.fetchall()
-
-for row in rows:
-    name, image_path = row
-    image = Image.open(image_path)
-    encoding = dlib.face_recognition_model_v1.compute_face_descriptor(image)[0]
-    known_face_encodings.append(encoding)
-    known_face_names.append(name)
-
-conn.close()
-
-# Choose input method for the image (webcam or file upload)
-capture_mode = st.radio("Choose an option to upload your image", ["Use Webcam", "Upload File"])
-
-if capture_mode == "Use Webcam":
-    picture = st.camera_input("Take a picture")  # Capture image using webcam
-elif capture_mode == "Upload File":
-    picture = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-
-# Input fields for student details
-name = st.text_input("Enter your name")
-roll_no = st.text_input("Enter your roll number")
-
-# Handle image upload or webcam capture
-if st.button("Register"):
-    if not name or not roll_no:
-        st.error("Please fill in both name and roll number.")
-    elif not picture:
-        st.error("Please upload or capture an image.")
-    else:
-        try:
-            # Open the image based on capture mode
-            if capture_mode == "Use Webcam" and picture:
-                image = Image.open(picture)
-            elif capture_mode == "Upload File" and picture:
-                image = Image.open(picture)
-
-            # Save the image locally and upload it to Hugging Face
-            image_path = save_image_to_hugging_face(image, name, roll_no)
-
-            # Save user data to the database
-            save_to_database(name, roll_no, image_path)
-
-            # Update the known faces list
-            known_face_encodings.append(dlib.face_recognition_model_v1.compute_face_descriptor(image)[0])
-            known_face_names.append(name)
-
-            st.success(f"Student {name} registered successfully!")
-
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
-
-# Detect faces and emotions from webcam
-cap = cv2.VideoCapture(0)
-while True:
-    ret, frame = cap.read()
-    if not ret:
-        break
-
-    frame = process_frame(frame, known_face_encodings, known_face_names)
-    st.image(frame, channels="BGR", use_column_width=True)
-    break  # Stop after capturing one frame
-
-cap.release()
-
-# Display registered students and attendance history
-if st.checkbox("Show registered students"):
-    conn = sqlite3.connect(DATABASE)
-    cursor = conn.cursor()
-    cursor.execute("SELECT name, roll_no, image_path, timestamp FROM students")
-    rows = cursor.fetchall()
-    conn.close()
-
-    st.write("### Registered Students")
-    for row in rows:
-        name, roll_no, image_path, timestamp = row
-        st.write(f"**Name:** {name}, **Roll No:** {roll_no}, **Timestamp:** {timestamp}")
-        st.image(image_path, caption=f"{name} ({roll_no})", use_column_width=True)
+        st.error(f"Error uploading file: {e}")
+
+# Image Processing Functions
+def detect_faces_and_emotions(image):
+    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
+    faces = face_cascade.detectMultiScale(gray_image, scaleFactor=1.3, minNeighbors=5)
+
+    for (x, y, w, h) in faces:
+        face = gray_image[y:y+h, x:x+w]
+        resized_face = cv2.resize(face, (48, 48))
+        normalized_face = resized_face / 255.0
+        reshaped_face = np.reshape(normalized_face, (1, 48, 48, 1))
+        emotion_prediction = emotion_model.predict(reshaped_face)
+        emotion_label = np.argmax(emotion_prediction)
+        return emotion_label
+
+    return None
+
+# UI Design
+st.title("Student Registration and Emotion Detection")
+create_table()
+
+menu = ["Register Student", "Face Recognition and Emotion Detection"]
+choice = st.sidebar.selectbox("Menu", menu)
+
+if choice == "Register Student":
+    st.subheader("Register a New Student")
+
+    with st.form("register_form"):
+        name = st.text_input("Name")
+        roll_number = st.text_input("Roll Number")
+        image_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
+        submitted = st.form_submit_button("Register")
+
+        if submitted:
+            if name and roll_number and image_file:
+                try:
+                    img = Image.open(image_file)
+                    img_path = os.path.join(KNOWN_FACES_DIR, f"{roll_number}.png")
+                    img.save(img_path)
+                    insert_student(name, roll_number, img_path)
+                    upload_to_hugging_face(img_path, f"{roll_number}.png")
+                    st.success("Student Registered Successfully!")
+                except Exception as e:
+                    st.error(f"Error: {e}")
+            else:
+                st.warning("Please fill in all fields and upload an image.")
+
+elif choice == "Face Recognition and Emotion Detection":
+    st.subheader("Recognize Faces and Detect Emotions")
+
+    action = st.radio("Choose Action", ["Upload Image", "Use Webcam"])
+
+    if action == "Upload Image":
+        uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
+        if uploaded_file:
+            try:
+                img = Image.open(uploaded_file)
+                img_array = np.array(img)
+                emotion_label = detect_faces_and_emotions(img_array)
+                if emotion_label is not None:
+                    st.success(f"Emotion Detected: {emotion_label}")
+                else:
+                    st.warning("No face detected.")
+            except Exception as e:
+                st.error(f"Error: {e}")
+
+    elif action == "Use Webcam":
+        cap = cv2.VideoCapture(0)
+        if st.button("Start Webcam"):
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                emotion_label = detect_faces_and_emotions(frame)
+                if emotion_label is not None:
+                    st.success(f"Emotion Detected: {emotion_label}")
+                cv2.imshow("Webcam", frame)
+                if cv2.waitKey(1) & 0xFF == ord('q'):
+                    break
+            cap.release()
+            cv2.destroyAllWindows()