Update app.py

app.py

@@ -13,9 +13,19 @@ def load_model(model_path='onnx_model.onnx'):
 # Preprocess the image
 def preprocess_image(image):
     image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)  # Convert to BGR (OpenCV format)
-    image_resized = cv2.resize(image, (224, 224))  # Resize to 224x224
-    image_input = np.expand_dims(image_resized, axis=0)  # Add batch dimension
-    image_input = image_input.transpose(0, 3, 1, 2)  # Change dimensions to (1, 3, 224, 224)
+    
+    # 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
 
@@ -34,32 +44,4 @@ def predict_emotion_onnx(model, image_input):
     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}")
+    prediction = model.run([out