Spaces:
Build error
Build error
| import streamlit as st | |
| import cv2 | |
| import numpy as np | |
| from datetime import datetime | |
| import torch | |
| from facenet_pytorch import MTCNN, InceptionResnetV1 | |
| from keras.models import load_model | |
| from PIL import Image | |
| import sqlite3 | |
| import os | |
| import tempfile | |
| # SQLite Database Connection | |
| DB_NAME = "emotion_detection.db" | |
| # 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() | |
| initialize_database() | |
| # Page title | |
| st.markdown("<h1 style='text-align: center;'>Emotion Detection with Face Recognition</h1>", unsafe_allow_html=True) | |
| st.markdown("<h3 style='text-align: center;'>angry, fear, happy, neutral, sad, surprise</h3>", unsafe_allow_html=True) | |
| # Load the emotion detection model | |
| def load_emotion_model(): | |
| model = load_model('CNN_Model_acc_75.h5') # Ensure this file is present | |
| 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() | |
| # Known faces and embeddings | |
| known_faces = [] | |
| known_names = [] | |
| def load_known_faces(): | |
| folder_path = "known_faces" # Folder containing known face images | |
| for image_name in os.listdir(folder_path): | |
| if image_name.endswith(('.jpg', '.jpeg', '.png')): | |
| image_path = os.path.join(folder_path, 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 for similarity | |
| min_distance = distance | |
| name = known_names[idx] | |
| return name | |
| 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: | |
| result_text = "No face detected!" | |
| return frame, result_text | |
| 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) | |
| # Sidebar for input source | |
| 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"]) | |
| 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"]) | |
| if uploaded_video: | |
| with tempfile.NamedTemporaryFile(delete=False) as tfile: | |
| tfile.write(uploaded_video.read()) | |
| video_source = cv2.VideoCapture(tfile.name) | |
| video_feed(video_source) | |
| #Display recent SQLite records | |
| st.markdown("### 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: | |
| col1, col2, col3 = st.columns(3) | |
| col1.write(f"**Name**: {record[0]}") | |
| col2.write(f"**Emotion**: {record[1]}") | |
| col3.write(f"**Timestamp**: {record[2]}") | |