File size: 1,704 Bytes
231be40
1f130ae
 
7ffdd20
412adf0
 
1f130ae
 
46274ff
1f130ae
 
 
d71fc08
 
 
 
 
 
 
 
 
 
 
 
 
7ffdd20
1effd41
1f130ae
412adf0
 
 
 
 
 
1f130ae
 
412adf0
1f130ae
412adf0
c096457
7ffdd20
 
36dac79
1f130ae
 
7ffdd20
1f130ae
 
36dac79
1f130ae
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import onnxruntime as ort
import numpy as np
import cv2
from PIL import Image
import streamlit as st

# Load ONNX model
onnx_model = ort.InferenceSession("onnx_model.onnx")

# Preprocess image function
def preprocess_image(image):
    """Preprocess image to match model input requirements"""
    # Convert the image to grayscale
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    
    # Resize image to 48x48 (model's expected input size)
    image_resized = cv2.resize(image, (48, 48))
    
    # Add batch dimension and channels (for grayscale: 1 channel)
    image_input = np.expand_dims(image_resized, axis=0)  # Add batch dimension (1, 48, 48)
    image_input = np.expand_dims(image_input, axis=1)  # Add channel dimension (1, 1, 48, 48)
    
    # Normalize the image
    image_input = image_input.astype(np.float32) / 255.0
    
    return image_input

# Predict emotion using the ONNX model
def predict_emotion_onnx(onnx_model, image_input):
    input_name = onnx_model.get_inputs()[0].name
    output_name = onnx_model.get_outputs()[0].name
    prediction = onnx_model.run([output_name], {input_name: image_input})
    return prediction

# Streamlit interface
st.title("Emotion Recognition with ONNX")

uploaded_file = st.file_uploader("Upload 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)

    # Preprocess the image
    image_input = preprocess_image(image)

    # Predict the emotion
    emotion_prediction = predict_emotion_onnx(onnx_model, image_input)

    # Display the prediction
    st.write(f"Predicted Emotion: {emotion_prediction[0]}")