Update app.py

app.py

@@ -1,136 +1,188 @@
 import streamlit as st
 import cv2
+import os
 import numpy as np
-import time
 from keras.models import load_model
 from PIL import Image
-from huggingface_hub import HfApi, Repository
-import os
-import tempfile
-
-# Page configuration
-st.set_page_config(page_title="Emotion Detection", layout="centered")
-
-# Title and Subtitle
-st.markdown("<h1 style='text-align: center;'>Emotion Detection</h1>", unsafe_allow_html=True)
-st.markdown("<h3 style='text-align: center;'>angry, fear, happy, neutral, sad, surprise</h3>", unsafe_allow_html=True)
-
-# Load Model
-@st.cache_resource
-def load_emotion_model():
-    model = load_model('CNN_Model_acc_75.h5')
-    return model
-
-start_time = time.time()
-model = load_emotion_model()
-st.write(f"Model loaded in {time.time() - start_time:.2f} seconds.")
-
-# Emotion labels and constants
-emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
-img_shape = 48
-face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-
-def process_frame(frame):
-    """Detect faces and predict emotions."""
-    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-    faces = face_cascade.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
-    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))
-        face_roi = np.expand_dims(face_roi, axis=0)
-        face_roi = face_roi / float(img_shape)
-        predictions = model.predict(face_roi)
-        emotion = emotion_labels[np.argmax(predictions[0])]
-
-        # Draw rectangle and emotion label
-        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
-        cv2.putText(frame, emotion, (x, y + h + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
-    return frame
-
-# Sidebar for input selection
-st.sidebar.title("Choose Input Source")
-upload_choice = st.sidebar.radio("Select:", ["Camera", "Upload Video", "Upload Image", "Upload to Hugging Face"])
-
-if upload_choice == "Camera":
-    # Use Streamlit's camera input widget
-    st.sidebar.info("Click a picture to analyze emotion.")
-    picture = st.camera_input("Take a picture")
-    if picture:
-        image = Image.open(picture)
-        frame = np.array(image)
-        frame = process_frame(frame)
-        st.image(frame, caption="Processed Image", use_column_width=True)
-
-elif upload_choice == "Upload Video":
-    uploaded_video = st.file_uploader("Upload Video", type=["mp4", "mov", "avi", "mkv", "webm"])
-    if uploaded_video:
-        with tempfile.NamedTemporaryFile(delete=False) as tfile:
-            tfile.write(uploaded_video.read())
-            video_source = cv2.VideoCapture(tfile.name)
-            frame_placeholder = st.empty()
-            while video_source.isOpened():
-                ret, frame = video_source.read()
-                if not ret:
-                    break
-                frame = process_frame(frame)
-                frame_placeholder.image(frame, channels="BGR", use_column_width=True)
-            video_source.release()
-
-elif upload_choice == "Upload Image":
-    uploaded_image = st.camera_input("Take a picture")
-    if uploaded_image:
-        image = Image.open(uploaded_image)
-        frame = np.array(image)
-        frame = process_frame(frame)
-        st.image(frame, caption="Processed Image", use_column_width=True)
-
-elif upload_choice == "Upload to Hugging Face":
-    st.sidebar.info("Upload images to the 'known_faces' directory in the Hugging Face repository.")
-
-    # Configure Hugging Face Repository
-    REPO_NAME = "face_and_emotion_detection"
-    REPO_ID = "LovnishVerma/" + REPO_NAME
-    hf_token = os.getenv("upload")  # Set your Hugging Face token as an environment variable
-
-    if not hf_token:
-        st.error("Hugging Face token not found. Please set it as an environment variable named 'HF_TOKEN'.")
-        st.stop()
-
-    # Initialize Hugging Face API
-    api = HfApi()
-
-    def create_hugging_face_repo():
-        """Create or verify the Hugging Face repository."""
-        try:
-            api.create_repo(repo_id=REPO_ID, repo_type="dataset", 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}")
-
-    def upload_to_hugging_face(file):
-        """Upload a file to the Hugging Face repository."""
+import sqlite3
+from huggingface_hub import HfApi
+from datetime import datetime
+
+# Constants
+KNOWN_FACES_DIR = "known_faces"  # Directory to save user images
+DATABASE = "students.db"  # SQLite database file to store student information
+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 = f"LovnishVerma/{REPO_NAME}"
+
+# Ensure the directories exist
+os.makedirs(KNOWN_FACES_DIR, exist_ok=True)
+
+# Retrieve Hugging Face token from environment variable
+hf_token = os.getenv("upload")  # Replace with your actual Hugging Face token
+if not hf_token:
+    st.error("Hugging Face token not found. Please set the environment variable.")
+    st.stop()
+
+# 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 the emotion detection model
+try:
+    emotion_model = load_model(EMOTION_MODEL_FILE)
+except Exception as e:
+    st.error(f"Error loading emotion model: {e}")
+    st.stop()
+
+# Database Functions
+def initialize_database():
+    """ Initializes the SQLite database by creating the students table if it doesn't exist. """
+    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()
+
+def save_to_database(name, roll_no, image_path):
+    """ Saves the student's data to the database. """
+    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()
+
+def save_image_to_hugging_face(image, name, roll_no):
+    """ Saves the image locally and uploads it to Hugging Face. """
+    filename = f"{name}_{roll_no}.jpg"
+    local_path = os.path.join(KNOWN_FACES_DIR, filename)
+    image.save(local_path)
+    
+    try:
+        api.upload_file(path_or_fileobj=local_path, path_in_repo=filename, repo_id=REPO_ID, 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":
+    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"])  # 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:
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
-                temp_file.write(file.read())
-                temp_file_path = temp_file.name
-
-            api.upload_file(
-                path_or_fileobj=temp_file_path,
-                path_in_repo=f"known_faces/{os.path.basename(temp_file_path)}",
-                repo_id=REPO_ID,
-                token=hf_token,
-            )
-            st.success("File uploaded successfully to Hugging Face!")
+            # 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"Error uploading file to Hugging Face: {e}")
-
-    # Create the repository if it doesn't exist
-    create_hugging_face_repo()
-
-    # Upload image file
-    hf_uploaded_image = st.file_uploader("Upload Image to Hugging Face", type=["png", "jpg", "jpeg"])
-    if hf_uploaded_image:
-        upload_to_hugging_face(hf_uploaded_image)
-
-st.sidebar.write("Emotion Labels: Angry, Fear, Happy, Neutral, Sad, Surprise")
+            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)
+
+# Face and Emotion Detection Function
+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))  # Resize face to 48x48
+        rgb_face = cv2.cvtColor(resized_face, cv2.COLOR_GRAY2RGB)
+        normalized_face = rgb_face / 255.0
+        reshaped_face = np.reshape(normalized_face, (1, 48, 48, 3))
+        
+        # Predict the emotion
+        emotion_prediction = emotion_model.predict(reshaped_face)
+        emotion_label = np.argmax(emotion_prediction)
+        return EMOTION_LABELS[emotion_label]
+    return None
+
+# UI for Emotion Detection
+if st.sidebar.selectbox("Menu", ["Register Student", "Face Recognition and Emotion Detection", "View Attendance"]) == "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:
+            img = Image.open(uploaded_file)
+            img_array = np.array(img)
+            emotion_label = detect_faces_and_emotions(img_array)
+            if emotion_label:
+                st.success(f"Emotion Detected: {emotion_label}")
+            else:
+                st.warning("No face detected.")
+
+    elif action == "Use Webcam":
+        st.info("Use the camera input widget to capture an image.")
+        camera_image = st.camera_input("Take a picture")
+        if camera_image:
+            img = Image.open(camera_image)
+            img_array = np.array(img)
+            emotion_label = detect_faces_and_emotions(img_array)
+            if emotion_label:
+                st.success(f"Emotion Detected: {emotion_label}")
+            else:
+                st.warning("No face detected.")