LovnishVerma's picture
Update app.py
bbe4d10 verified
raw
history blame
6.89 kB
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
@st.cache_resource
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]}")