Update app.py

app.py

@@ -1,322 +1,173 @@
 import streamlit as st
 import cv2
 import numpy as np
-import time
 import os
-from keras.models import load_model
+import sqlite3
+import time
 from PIL import Image
-import tempfile
+from keras.models import load_model
 from huggingface_hub import HfApi
-
-# Larger title
-st.markdown("<h1 style='text-align: center;'>Emotion Detection with Face Recognition</h1>", unsafe_allow_html=True)
-
-# Smaller subtitle
-st.markdown("<h3 style='text-align: center;'>angry, fear, happy, neutral, sad, surprise</h3>", unsafe_allow_html=True)
-
-start = time.time()
+import tempfile
 
 # Constants
-KNOWN_FACES_DIR = "known_faces"  # Directory to save user images
-DATABASE = "students.db"  # SQLite database file to store student information
+KNOWN_FACES_DIR = "known_faces"
+DATABASE = "students.db"
+EMOTION_MODEL_FILE = "CNN_Model_acc_75.h5"
+EMOTION_LABELS = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
+REPO_NAME = "face_and_emotion_detection"
+REPO_ID = "LovnishVerma/" + REPO_NAME
+IMG_SHAPE = 48
+hf_token = os.getenv("upload")
 
-# Ensure the directory exists
-os.makedirs(KNOWN_FACES_DIR, exist_ok=True)
+# Initialize Hugging Face API
+api = HfApi()
 
-# Retrieve the Hugging Face token from environment variable
-hf_token = os.getenv("upload")  # The key must match the secret name set in Hugging Face
+# Ensure the Hugging Face token is available
 if not hf_token:
-    raise ValueError("Hugging Face token not found. Ensure it's set as a secret in the Hugging Face Space.")
+    st.error("Hugging Face token not found. Please set the environment variable.")
+    st.stop()
 
-# Initialize Hugging Face API
-api = HfApi()
+# Initialize the Hugging Face repository
+def create_hugging_face_repo():
+    try:
+        api.create_repo(repo_id=REPO_ID, repo_type="space", space_sdk="streamlit", token=hf_token, exist_ok=True)
+        st.success(f"Repository '{REPO_NAME}' is ready on Hugging Face!")
+    except Exception as e:
+        st.error(f"Error creating Hugging Face repository: {e}")
 
-# 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 the emotion model once, using caching
+@st.cache_resource
+def load_emotion_model():
+    try:
+        model = load_model(EMOTION_MODEL_FILE)
+        return model
+    except Exception as e:
+        st.error(f"Error loading emotion model: {e}")
+        st.stop()
 
-# Ensure the repository exists or create it
-try:
-    api.create_repo(repo_id=REPO_ID, repo_type=REPO_TYPE, space_sdk="streamlit", token=hf_token, exist_ok=True)
-    st.success(f"Repository '{REPO_NAME}' is ready on Hugging Face!")
-except Exception as e:
-    st.error(f"Error creating repository: {e}")
+emotion_model = load_emotion_model()
 
+# Initialize the face recognizer
+face_recognizer = cv2.face.LBPHFaceRecognizer_create()
 
-# Database setup
+# Database functions
 def initialize_database():
     """
-    Initializes the SQLite database by creating a table to store student data
-    such as name, roll number, image path, and registration timestamp.
+    Initializes the SQLite database by creating a table to store student data.
     """
-    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()
+    with sqlite3.connect(DATABASE) as conn:
+        conn.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()
 
 def save_to_database(name, roll_no, image_path):
     """
-    Saves the student's information (name, roll number, image path) to the SQLite database.
+    Saves student data (name, roll number, image path) to the SQLite database.
     Ensures roll number is unique.
-
-    Args:
-        name (str): The name of the student.
-        roll_no (str): The roll number of the student.
-        image_path (str): Path to the stored image of the student.
     """
-    conn = sqlite3.connect(DATABASE)
-    cursor = conn.cursor()
-    try:
-        cursor.execute("""
-            INSERT INTO students (name, roll_no, image_path)
-            VALUES (?, ?, ?)
-        """, (name, roll_no, image_path))
-        conn.commit()
-        st.success("Data saved successfully!")
-    except sqlite3.IntegrityError:
-        st.error("Roll number already exists!")
-    finally:
-        conn.close()
+    with sqlite3.connect(DATABASE) as conn:
+        try:
+            conn.execute("""
+                INSERT INTO students (name, roll_no, image_path)
+                VALUES (?, ?, ?)
+            """, (name, roll_no, image_path))
+            conn.commit()
+            st.success("Data saved successfully!")
+        except sqlite3.IntegrityError:
+            st.error("Roll number already exists!")
 
 def save_image_to_hugging_face(image, name, roll_no):
     """
     Saves the captured image locally in the 'known_faces' directory and uploads it to Hugging Face.
-    The image is renamed using the format 'UserName_RollNo.jpg'.
-
-    Args:
-        image (PIL Image): The image object captured by the user.
-        name (str): The name of the student.
-        roll_no (str): The roll number of the student.
-
-    Returns:
-        str: The local path where the image is saved.
     """
-    # Rename the image using the format 'UserName_RollNo.jpg'
     filename = f"{name}_{roll_no}.jpg"
     local_path = os.path.join(KNOWN_FACES_DIR, filename)
-    
-    # Save the image locally to the known_faces directory
     image.save(local_path)
 
-    # Try uploading the image to Hugging Face
     try:
         api.upload_file(
             path_or_fileobj=local_path,
             path_in_repo=filename,
             repo_id=REPO_ID,
-            repo_type=REPO_TYPE,
-            token=hf_token  # Pass the Hugging Face token directly
+            repo_type="space",
+            token=hf_token
         )
         st.success(f"Image uploaded to Hugging Face: {filename}")
     except Exception as e:
         st.error(f"Error uploading image to Hugging Face: {e}")
-
-    return local_path
-
-# Initialize the database when the app starts
-initialize_database()
-
-# Streamlit user interface (UI)
-st.title("Student Registration with Hugging Face Image Upload")
-
-# Input fields for student details
-name = st.text_input("Enter your name")
-roll_no = st.text_input("Enter your roll number")
-
-# 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"])
-
-# Handle webcam capture or file upload
-if capture_mode == "Use Webcam":
-    try:
-        picture = st.camera_input("Take a picture")  # Capture image using webcam
-    except Exception as e:
-        st.error(f"Error accessing webcam: {e}")
-        picture = None
-
-elif capture_mode == "Upload File":
-    picture = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])  # Upload image from file system
-
-# Save data and process image on button click
-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_to_database(name, roll_no, image_path)
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
-
-# Display registered student data
-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)
-
-
-# Constants
-DB_FILE = "students.db"
-KNOWN_FACES_DIR = "known_faces"
-EMOTION_MODEL_FILE = "CNN_Model_acc_75.h5"
-EMOTION_LABELS = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
-
-# Hugging Face Repository Details
-REPO_NAME = "face_and_emotion_detection"
-REPO_ID = f"LovnishVerma/{REPO_NAME}"  # Replace with your Hugging Face username and repository name
-
-# Ensure Directories
-os.makedirs(KNOWN_FACES_DIR, exist_ok=True)
-
-# Load Hugging Face Token
-hf_token = os.getenv("upload")
-if not hf_token:
-    st.error("Hugging Face token not found. Please set the environment variable.")
-
-# Initialize Hugging Face API
-api = HfApi()
-try:
-    api.create_repo(repo_id=REPO_ID, repo_type="space", space_sdk="streamlit", token=hf_token, exist_ok=True)
-    st.success(f"Repository '{REPO_NAME}' is ready on Hugging Face!")
-except Exception as e:
-    st.error(f"Error creating Hugging Face repository: {e}")
-
-# Load Emotion Model
-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:
-        conn.execute("""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()
-
-def insert_student(name, roll_number, image_path):
-    try:
-        with sqlite3.connect(DB_FILE) as conn:
-            conn.execute("INSERT INTO students (name, roll_number, image_path) VALUES (?, ?, ?)",
-                         (name, roll_number, image_path))
-            conn.commit()
-    except sqlite3.IntegrityError:
-        st.warning("Roll number already exists!")
-
-def get_all_students():
-    with sqlite3.connect(DB_FILE) as conn:
-        cursor = conn.execute("SELECT * FROM students")
-        return cursor.fetchall()
-
-# Load the emotion model
-@st.cache_resource
-def load_emotion_model():
-    model = load_model('CNN_Model_acc_75.h5')  # Ensure this file is in your Space
-    return model
-
-model = load_emotion_model()
-print("time taken to load model: ", time.time() - start)
-
-# Emotion labels
-emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
-
-# Load known faces (from images in a folder)
-known_faces = []
-known_names = []
-face_recognizer = cv2.face.LBPHFaceRecognizer_create()
+    return local_path
 
+# Load known faces
 def load_known_faces():
-    folder_path = "known_faces"  # Place your folder with known faces here
-    for image_name in os.listdir(folder_path):
+    """
+    Loads known faces from the 'known_faces' directory and trains the recognizer.
+    """
+    known_faces = []
+    known_names = []
+    
+    for image_name in os.listdir(KNOWN_FACES_DIR):
         if image_name.endswith(('.jpg', '.jpeg', '.png')):
-            image_path = os.path.join(folder_path, image_name)
+            image_path = os.path.join(KNOWN_FACES_DIR, image_name)
             image = cv2.imread(image_path)
             gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-            # Detect face in the image
-            faces = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml').detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
-
+            faces = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml').detectMultiScale(
+                gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30)
+            )
+            
             for (x, y, w, h) in faces:
                 roi_gray = gray[y:y+h, x:x+w]
-                # We only need the face, so we crop it and store it for training
                 known_faces.append(roi_gray)
                 known_names.append(image_name.split('.')[0])  # Assuming file name is the person's name
-
-    # Train the recognizer with the known faces
-    face_recognizer.train(known_faces, np.array([i for i in range(len(known_faces))]))
+    
+    if known_faces:
+        face_recognizer.train(known_faces, np.array([i for i in range(len(known_faces))]))
+    else:
+        st.warning("No known faces found for training.")
 
 load_known_faces()
 
-# Face detection using OpenCV
-face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-img_shape = 48
-
+# Process frame for both emotion detection and face recognition
 def process_frame(frame):
     gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-    faces = face_cascade.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
-
-    result_text = ""  # Initialize the result text for display
+    faces = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml').detectMultiScale(
+        gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30)
+    )
 
+    result_text = ""
     for (x, y, w, h) in faces:
         roi_gray = gray_frame[y:y+h, x:x+w]
         roi_color = frame[y:y+h, x:x+w]
-        face_roi = cv2.resize(roi_color, (img_shape, img_shape))  # Resize to 48x48
-        face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB)  # Convert to RGB (3 channels)
-        face_roi = np.expand_dims(face_roi, axis=0)  # Add batch dimension
-        face_roi = face_roi / 255.0  # Normalize the image
-
-        # Emotion detection
-        predictions = model.predict(face_roi)
-        emotion = emotion_labels[np.argmax(predictions[0])]
-
-        # Face recognition using LBPH
+        face_roi = cv2.resize(roi_color, (IMG_SHAPE, IMG_SHAPE))
+        face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB)
+        face_roi = np.expand_dims(face_roi, axis=0) / 255.0
+        
+        predictions = emotion_model.predict(face_roi)
+        emotion = EMOTION_LABELS[np.argmax(predictions[0])]
+        
         label, confidence = face_recognizer.predict(roi_gray)
         name = "Unknown"
         if confidence < 100:
             name = known_names[label]
 
-        # Format the result text as "Name is feeling Emotion"
         result_text = f"{name} is feeling {emotion}"
 
-        # Draw bounding box and label on the frame
         cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
         cv2.putText(frame, result_text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
 
     return frame, result_text
 
-# Video feed
+# Video feed handler
 def video_feed(video_source):
-    frame_placeholder = st.empty()  # This placeholder will be used to replace frames in-place
-    text_placeholder = st.empty()  # This placeholder will display the result text
+    frame_placeholder = st.empty()
+    text_placeholder = st.empty()
 
     while True:
         ret, frame = video_source.read()
@@ -325,42 +176,72 @@ def video_feed(video_source):
 
         frame, result_text = process_frame(frame)
 
-        # Display the frame in the placeholder
         frame_placeholder.image(frame, channels="BGR", use_column_width=True)
-
-        # Display the result text in the text placeholder
         text_placeholder.markdown(f"<h3 style='text-align: center;'>{result_text}</h3>", unsafe_allow_html=True)
 
-# Sidebar for video or image upload
-upload_choice = st.sidebar.radio("Choose input source", ["Upload Image", "Upload Video", "Camera"])
-
-if upload_choice == "Camera":
-    # Use Streamlit's built-in camera input widget for capturing images from the webcam
-    image = st.camera_input("Take a picture")
-
-    if image is not None:
-        # Convert the image to a numpy array
-        frame = np.array(Image.open(image))
-        frame, result_text = process_frame(frame)
-        st.image(frame, caption='Processed Image', use_column_width=True)
-        st.markdown(f"<h3 style='text-align: center;'>{result_text}</h3>", unsafe_allow_html=True)
+# Streamlit interface
+def main():
+    st.title("Student Registration with Face Recognition and Emotion Detection")
 
-elif upload_choice == "Upload Image":
-    uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg", "gif"])
-    if uploaded_image:
-        image = Image.open(uploaded_image)
-        frame = np.array(image)
-        frame, result_text = process_frame(frame)
-        st.image(frame, caption='Processed Image', use_column_width=True)
-        st.markdown(f"<h3 style='text-align: center;'>{result_text}</h3>", unsafe_allow_html=True)
-
-elif upload_choice == "Upload Video":
-    uploaded_video = st.file_uploader("Upload Video", type=["mp4", "mov", "avi", "mkv", "webm"])
-    if uploaded_video:
-        # Temporarily save the video to disk
-        with tempfile.NamedTemporaryFile(delete=False) as tfile:
-            tfile.write(uploaded_video.read())
-            video_source = cv2.VideoCapture(tfile.name)
+    name = st.text_input("Enter your name")
+    roll_no = st.text_input("Enter your roll number")
+    
+    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")
+    else:
+        picture = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+
+    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:
+                image = Image.open(picture)
+                image_path = save_image_to_hugging_face(image, name, roll_no)
+                save_to_database(name, roll_no, image_path)
+            except Exception as e:
+                st.error(f"An error occurred: {e}")
+
+    if st.checkbox("Show registered students"):
+        with sqlite3.connect(DATABASE) as conn:
+            cursor = conn.cursor()
+            cursor.execute("SELECT name, roll_no, image_path, timestamp FROM students")
+            rows = cursor.fetchall()
+
+        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)
+
+    upload_choice = st.sidebar.radio("Choose input source", ["Upload Image", "Upload Video", "Camera"])
+
+    if upload_choice == "Camera":
+        image = st.camera_input("Take a picture")
+        if image:
+            frame = np.array(Image.open(image))
+            frame, result_text = process_frame(frame)
+            st.image(frame, caption='Processed Image', use_column_width=True)
+            st.markdown(f"<h3 style='text-align: center;'>{result_text}</h3>", unsafe_allow_html=True)
+    elif upload_choice == "Upload Image":
+        uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg", "gif"])
+        if uploaded_image:
+            image = Image.open(uploaded_image)
+            frame = np.array(image)
+            frame, result_text = process_frame(frame)
+            st.image(frame, caption='Processed Image', use_column_width=True)
+            st.markdown(f"<h3 style='text-align: center;'>{result_text}</h3>", unsafe_allow_html=True)
+    elif upload_choice == "Upload Video":
+        video_file = st.file_uploader("Upload Video", type=["mp4", "mov", "avi"])
+        if video_file:
+            temp_video_file = tempfile.NamedTemporaryFile(delete=False)
+            temp_video_file.write(video_file.read())
+            temp_video_file.close()
+            video_source = cv2.VideoCapture(temp_video_file.name)
             video_feed(video_source)
 
-st.sidebar.write("Emotion Labels: Angry, Fear, Happy, Neutral, Sad, Surprise")
+if __name__ == "__main__":
+    main()