|
import streamlit as st
|
|
import cv2
|
|
import torch
|
|
from PIL import Image
|
|
from transformers import AutoImageProcessor, AutoModelForImageClassification
|
|
from io import BytesIO
|
|
import numpy as np
|
|
|
|
|
|
processor = AutoImageProcessor.from_pretrained("RickyIG/emotion_face_image_classification")
|
|
model = AutoModelForImageClassification.from_pretrained("RickyIG/emotion_face_image_classification")
|
|
|
|
|
|
st.title("Emotion Detection App")
|
|
|
|
|
|
option = st.radio("Select an option", ("Upload Image", "Use Live Camera"))
|
|
|
|
if option == "Upload Image":
|
|
|
|
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
|
|
|
|
if uploaded_file is not None:
|
|
|
|
image = Image.open(uploaded_file)
|
|
st.image(image, caption="Uploaded Image", use_column_width=True)
|
|
|
|
|
|
inputs = processor(images=image, return_tensors="pt")
|
|
|
|
|
|
with torch.no_grad():
|
|
outputs = model(**inputs)
|
|
logits = outputs.logits
|
|
predicted_class_idx = logits.argmax(-1).item()
|
|
|
|
|
|
label = model.config.id2label[predicted_class_idx]
|
|
|
|
|
|
st.write(f"Predicted Emotion: {label}")
|
|
|
|
elif option == "Use Live Camera":
|
|
|
|
cap = cv2.VideoCapture(0)
|
|
|
|
if not cap.isOpened():
|
|
st.error("Error: Could not open webcam.")
|
|
else:
|
|
stframe = st.empty()
|
|
|
|
while True:
|
|
|
|
ret, frame = cap.read()
|
|
|
|
if not ret:
|
|
st.error("Error: Failed to capture frame.")
|
|
break
|
|
|
|
|
|
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
|
|
|
|
|
inputs = processor(images=image, return_tensors="pt")
|
|
|
|
|
|
with torch.no_grad():
|
|
outputs = model(**inputs)
|
|
logits = outputs.logits
|
|
predicted_class_idx = logits.argmax(-1).item()
|
|
|
|
|
|
label = model.config.id2label[predicted_class_idx]
|
|
|
|
|
|
cv2.putText(frame, f"Emotion: {label}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)
|
|
|
|
|
|
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
stframe.image(frame_rgb, channels="RGB", use_column_width=True)
|
|
|
|
|
|
cap.release()
|
|
|