Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| import onnxruntime as ort | |
| import streamlit as st | |
| from PIL import Image | |
| # Load the ONNX model | |
| def load_model(model_path='onnx_model.onnx'): | |
| # Load the ONNX model | |
| model = ort.InferenceSession(model_path) | |
| return model | |
| # Preprocess the image | |
| def preprocess_image(image): | |
| image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # Convert to BGR (OpenCV format) | |
| # Resize the image to 48x48 as per the error message (model's expected input size) | |
| image_resized = cv2.resize(image, (48, 48)) # Resize to 48x48 | |
| # Convert to grayscale if the model expects a single channel | |
| image_gray = cv2.cvtColor(image_resized, cv2.COLOR_BGR2GRAY) # Convert to grayscale | |
| # If the model expects 3 channels, keep the image in RGB (3 channels) | |
| # image_resized = cv2.cvtColor(image_resized, cv2.COLOR_BGR2RGB) # For RGB input | |
| # Add batch dimension | |
| image_input = np.expand_dims(image_gray, axis=0) # Add batch dimension | |
| image_input = np.expand_dims(image_input, axis=0) # Add channel dimension (for grayscale) | |
| image_input = image_input.astype(np.float32) / 255.0 # Normalize the image | |
| return image_input | |
| # Map the raw output to emotions | |
| def get_emotion_from_output(output): | |
| emotion_labels = ['Anger', 'Disgust', 'Fear', 'Happiness', 'Sadness', 'Surprise', 'Neutral'] | |
| # Get the index of the highest value in the output (i.e., predicted emotion) | |
| emotion_index = np.argmax(output) | |
| confidence = output[0][emotion_index] # Confidence of the prediction | |
| emotion = emotion_labels[emotion_index] # Corresponding emotion label | |
| return emotion, confidence | |
| # Predict emotion using the ONNX model | |
| def predict_emotion_onnx(model, image_input): | |
| # Get the input name and output name for the ONNX model | |
| input_name = model.get_inputs()[0].name | |
| output_name = model.get_outputs()[0].name | |
| # Run the model | |
| prediction = model.run([output_name], {input_name: image_input}) | |
| return prediction[0] | |
| # Streamlit UI | |
| st.title("Emotion Detection") | |
| # Upload an image | |
| uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"]) | |
| if uploaded_file is not None: | |
| # Open and display the uploaded image | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_column_width=True) | |
| # Load model | |
| onnx_model = load_model() | |
| # Preprocess the image | |
| image_input = preprocess_image(image) | |
| # Get emotion prediction | |
| emotion_prediction = predict_emotion_onnx(onnx_model, image_input) | |
| # Get the emotion label and confidence | |
| emotion_label, confidence = get_emotion_from_output(emotion_prediction) | |
| # Display the predicted emotion and confidence | |
| st.write(f"Predicted Emotion: {emotion_label}") | |
| st.write(f"Confidence: {confidence:.2f}") | |