Update app.py

app.py

@@ -1,251 +1,136 @@
 import streamlit as st
 import cv2
 import numpy as np
-import os
-import sqlite3
-from PIL import Image
+import time
 from keras.models import load_model
-from huggingface_hub import HfApi
+from PIL import Image
+from huggingface_hub import HfApi, Repository
+import os
 import tempfile
 
-# Constants
-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 Hugging Face token is available
-if not hf_token:
-    st.error("Hugging Face token not found. Please set the environment variable.")
-    st.stop()
-
-# Initialize Hugging Face API
-api = HfApi()
-
-# Create 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}")
-
-# Load the emotion model once, using caching
+# 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():
-    try:
-        model = load_model(EMOTION_MODEL_FILE)
-        return model
-    except Exception as e:
-        st.error(f"Error loading emotion model: {e}")
-        st.stop()
+    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')
 
-emotion_model = load_emotion_model()
-
-# Initialize the face recognizer
-face_recognizer = cv2.face.LBPHFaceRecognizer_create()
-
-# Database functions
-def initialize_database():
-    """
-    Initializes the SQLite database by creating a table to store student data.
-    """
-    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()
-
-# Initialize the database
-initialize_database()
-
-def save_to_database(name, roll_no, image_path):
-    """
-    Saves student data (name, roll number, image path) to the SQLite database.
-    Ensures roll number is unique.
-    """
-    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.
-    """
-    if not os.path.exists(KNOWN_FACES_DIR):
-        os.makedirs(KNOWN_FACES_DIR)
-    
-    filename = f"{name}_{roll_no}.jpg"
-    local_path = os.path.join(KNOWN_FACES_DIR, filename)
-    
-    # Saving the image to the correct directory
-    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
-
-# Load known faces
-def load_known_faces():
-    """
-    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(KNOWN_FACES_DIR, image_name)
-            image = cv2.imread(image_path)
-            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-            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]
-                known_faces.append(roi_gray)
-                known_names.append(image_name.split('.')[0])  # Assuming file name is the person's name
-    
-    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.")
-    
-    return known_names
-
-# Load known faces at the start
-known_names = load_known_faces()
-
-# Process frame for both emotion detection and face recognition
 def process_frame(frame):
+    """Detect faces and predict emotions."""
     gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-    faces = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml').detectMultiScale(
-        gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30)
-    )
-
-    result_text = ""
+    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 = 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]
-
-        result_text = f"{name} is feeling {emotion}"
+        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, result_text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-
-    return frame, result_text
-
-# Video feed handler
-def video_feed(video_source):
-    frame_placeholder = st.empty()
-    text_placeholder = st.empty()
-
-    while True:
-        ret, frame = video_source.read()
-        if not ret:
-            break
-
-        frame, result_text = process_frame(frame)
-
-        frame_placeholder.image(frame, channels="BGR", use_column_width=True)
-        text_placeholder.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")
-
-    # Step 1: Student Registration
-    registration_mode = st.sidebar.radio("Choose an option", ["Register Student", "Face and Emotion Recognition"])
-
-    if registration_mode == "Register Student":
-        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}")
-
-    elif registration_mode == "Face and Emotion Recognition":
-        upload_choice = st.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)
-
-if __name__ == "__main__":
-    main()
+        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.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
+    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."""
+        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!")
+        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")