Update app.py

app.py

@@ -3,7 +3,6 @@ import cv2
 import numpy as np
 import os
 import sqlite3
-import time
 from PIL import Image
 from keras.models import load_model
 from huggingface_hub import HfApi
@@ -19,15 +18,15 @@ REPO_ID = "LovnishVerma/" + REPO_NAME
 IMG_SHAPE = 48
 hf_token = os.getenv("upload")
 
-# Initialize Hugging Face API
-api = HfApi()
-
 # 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 the Hugging Face repository
+# 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)
@@ -87,6 +86,9 @@ 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)
     image.save(local_path)
@@ -183,65 +185,59 @@ def video_feed(video_source):
 def main():
     st.title("Student Registration with Face Recognition and Emotion Detection")
 
-    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"])
+    # Step 1: Student Registration
+    registration_mode = st.sidebar.radio("Choose an option", ["Register Student", "Face and Emotion Recognition"])
 
-    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.")
+    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:
-            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)
+            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()