Update app.py

app.py

@@ -1,115 +1,207 @@
-import os
+import streamlit as st
 import cv2
-import sqlite3
 import numpy as np
-import streamlit as st
 from datetime import datetime
-from tensorflow.keras.models import load_model
+import torch
+from facenet_pytorch import MTCNN, InceptionResnetV1
+from keras.models import load_model
 from PIL import Image
+import sqlite3
+import os
+import tempfile
 
-# Initialize SQLite database
-def init_db():
-    conn = sqlite3.connect("face_records.db")
-    c = conn.cursor()
-    c.execute('''CREATE TABLE IF NOT EXISTS records (
-                     id INTEGER PRIMARY KEY AUTOINCREMENT,
-                     name TEXT,
-                     timestamp TEXT
-                 )''')
-    conn.commit()
-    conn.close()
+# SQLite Database
+DB_NAME = "emotion_detection.db"
 
-# Save record to SQLite database
-def save_record(name):
-    conn = sqlite3.connect("face_records.db")
-    c = conn.cursor()
-    c.execute("INSERT INTO records (name, timestamp) VALUES (?, ?)", (name, datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
+# Initialize SQLite Database
+def initialize_database():
+    conn = sqlite3.connect(DB_NAME)
+    cursor = conn.cursor()
+    cursor.execute("""
+    CREATE TABLE IF NOT EXISTS face_data (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        name TEXT NOT NULL,
+        emotion TEXT NOT NULL,
+        timestamp TEXT NOT NULL
+    )
+    """)
     conn.commit()
     conn.close()
 
-# Load face recognition model
+initialize_database()
+
+# Load the emotion detection model
+@st.cache_resource
 def load_emotion_model():
-    return load_model("CNN_Model_acc_75.h5")
-
-# Register new face
-def register_face(name):
-    st.info("Open your camera to capture your face.")
-    cap = cv2.VideoCapture(0)
-    if not cap.isOpened():
-        st.error("Unable to access the camera.")
-        return
-
-    st.write("Press 's' to save the image or 'q' to quit.")
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            st.error("Failed to capture frame from camera.")
-            break
-
-        cv2.imshow("Register Face", frame)
-
-        key = cv2.waitKey(1)
-        if key == ord("s"):
-            known_faces_dir = "known_faces"
-            os.makedirs(known_faces_dir, exist_ok=True)
-            image_path = os.path.join(known_faces_dir, f"{name}.jpg")
-            cv2.imwrite(image_path, frame)
-            st.success(f"Face saved successfully as {image_path}.")
-            save_record(name)
-            break
-        elif key == ord("q"):
-            break
-
-    cap.release()
-    cv2.destroyAllWindows()
-
-# Show registered faces
-def show_registered_faces():
-    st.header("Registered Faces")
-    known_faces_dir = "known_faces"
-    if not os.path.exists(known_faces_dir):
-        st.warning("No registered faces found.")
-        return
-
-    col1, col2, col3 = st.columns(3)
-    for i, file in enumerate(os.listdir(known_faces_dir)):
-        if file.endswith(".jpg"):
-            img_path = os.path.join(known_faces_dir, file)
-            img = Image.open(img_path)
-            with [col1, col2, col3][i % 3]:
-                st.image(img, caption=file, use_column_width=True)
-
-# Show recent records
-def show_recent_records():
-    st.header("Recent Records")
-    conn = sqlite3.connect("face_records.db")
-    c = conn.cursor()
-    c.execute("SELECT * FROM records ORDER BY timestamp DESC LIMIT 10")
-    records = c.fetchall()
-    conn.close()
+    model = load_model('CNN_Model_acc_75.h5')  # Ensure the file exists
+    return model
+
+emotion_model = load_emotion_model()
+emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
+
+# Initialize FaceNet model and MTCNN
+facenet = InceptionResnetV1(pretrained='vggface2').eval()
+mtcnn = MTCNN()
+
+# Directory for known faces
+KNOWN_FACES_DIR = "known_faces"
+if not os.path.exists(KNOWN_FACES_DIR):
+    os.makedirs(KNOWN_FACES_DIR)
+
+# Load known faces and embeddings
+known_faces = []
+known_names = []
+
+def load_known_faces():
+    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 = Image.open(image_path).convert("RGB")
+            face, _ = mtcnn.detect(image)
+
+            if face is not None:
+                face_box = face[0].astype(int)
+                cropped_face = image.crop((face_box[0], face_box[1], face_box[2], face_box[3]))
+                cropped_face = cropped_face.resize((160, 160))
+                face_tensor = np.array(cropped_face).transpose(2, 0, 1) / 255.0
+                face_tensor = torch.tensor(face_tensor, dtype=torch.float32).unsqueeze(0)
+                
+                with torch.no_grad():
+                    embedding = facenet(face_tensor).numpy()
+                
+                known_faces.append(embedding)
+                known_names.append(image_name.split('.')[0])
+
+load_known_faces()
+
+def recognize_face(embedding):
+    min_distance = float('inf')
+    name = "Unknown"
+    for idx, known_embedding in enumerate(known_faces):
+        distance = np.linalg.norm(known_embedding - embedding)
+        if distance < min_distance and distance < 0.6:  # Threshold
+            min_distance = distance
+            name = known_names[idx]
+    return name
 
-    if records:
-        for record in records:
-            st.write(f"Name: {record[1]}, Time: {record[2]}")
+def process_frame(frame):
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    faces, _ = mtcnn.detect(frame_rgb)
+    result_text = ""
+
+    if faces is not None:
+        for face_box in faces:
+            x1, y1, x2, y2 = map(int, face_box)
+            cropped_face = frame_rgb[y1:y2, x1:x2]
+            resized_face = cv2.resize(cropped_face, (48, 48))
+            face_normalized = resized_face / 255.0
+            face_array = np.expand_dims(face_normalized, axis=0)
+
+            # Emotion prediction
+            predictions = emotion_model.predict(face_array)
+            emotion = emotion_labels[np.argmax(predictions[0])]
+
+            # Face recognition
+            cropped_face_for_recognition = cv2.resize(cropped_face, (160, 160))
+            face_tensor = np.array(cropped_face_for_recognition).transpose(2, 0, 1) / 255.0
+            face_tensor = torch.tensor(face_tensor, dtype=torch.float32).unsqueeze(0)
+
+            with torch.no_grad():
+                face_embedding = facenet(face_tensor).numpy()
+
+            name = recognize_face(face_embedding)
+
+            # Save record in SQLite
+            if name != "Unknown":
+                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                conn = sqlite3.connect(DB_NAME)
+                cursor = conn.cursor()
+                cursor.execute("""
+                INSERT INTO face_data (name, emotion, timestamp)
+                VALUES (?, ?, ?)
+                """, (name, emotion, timestamp))
+                conn.commit()
+                conn.close()
+
+            # Display result
+            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            result_text = f"{name} is feeling {emotion}"
+            cv2.putText(frame, result_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
     else:
-        st.info("No recent records found.")
-
-# Main App
-init_db()
-st.title("Face Registration and Records")
-menu = ["Home", "Register Face", "View Registered Faces", "Recent Records"]
-choice = st.sidebar.selectbox("Menu", menu)
-
-if choice == "Home":
-    st.write("Welcome to the Face Registration System.")
-elif choice == "Register Face":
-    name = st.text_input("Enter your name:")
-    if st.button("Register"):
-        if name.strip():
-            register_face(name)
+        result_text = "No face detected!"
+
+    return frame, result_text
+
+# Sidebar menu
+menu = st.sidebar.selectbox("Menu", ["Home", "Register New Face", "View Records"])
+
+if menu == "Home":
+    st.title("Emotion Detection")
+    st.write("Choose input source to start detection.")
+    upload_choice = st.radio("Choose Input Source", ["Camera", "Upload Image", "Upload Video"])
+
+    if upload_choice == "Camera":
+        cap = cv2.VideoCapture(0)
+        if not cap.isOpened():
+            st.error("Unable to access the camera.")
         else:
-            st.error("Name cannot be empty.")
-elif choice == "View Registered Faces":
-    show_registered_faces()
-elif choice == "Recent Records":
-    show_recent_records()
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                frame, result_text = process_frame(frame)
+                st.image(frame, channels="BGR")
+                st.write(result_text)
+            cap.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, result_text = process_frame(frame)
+            st.image(frame)
+            st.write(result_text)
+
+    elif upload_choice == "Upload Video":
+        uploaded_video = st.file_uploader("Upload Video", type=["mp4", "mov", "avi"])
+        if uploaded_video:
+            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+                temp_file.write(uploaded_video.read())
+                video_source = cv2.VideoCapture(temp_file.name)
+                while video_source.isOpened():
+                    ret, frame = video_source.read()
+                    if not ret:
+                        break
+                    frame, result_text = process_frame(frame)
+                    st.image(frame, channels="BGR")
+                    st.write(result_text)
+                video_source.release()
+
+elif menu == "Register New Face":
+    st.title("Register New Face")
+    name = st.text_input("Enter Name")
+    if st.button("Capture Image"):
+        cap = cv2.VideoCapture(0)
+        if not cap.isOpened():
+            st.error("Unable to access the camera.")
+        else:
+            ret, frame = cap.read()
+            if ret:
+                image_path = os.path.join(KNOWN_FACES_DIR, f"{name}.jpg")
+                cv2.imwrite(image_path, frame)
+                st.success(f"Face registered successfully for {name}")
+                load_known_faces()  # Refresh known faces
+            cap.release()
+
+elif menu == "View Records":
+    st.title("View Records")
+    st.subheader("Recent Records")
+    conn = sqlite3.connect(DB_NAME)
+    cursor = conn.cursor()
+    cursor.execute("SELECT name, emotion, timestamp FROM face_data ORDER BY timestamp DESC LIMIT 5")
+    records = cursor.fetchall()
+    conn.close()
+
+    for record in records:
+        st.write(f"**Name**: {record[0]}, **Emotion**: {record[1]}, **Timestamp**: {record[2]}")