import streamlit as st
import cv2
import numpy as np
from PIL import Image
import tensorflow as tf
face_detector = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
# Load the Keras model
model = tf.keras.models.load_model("./affectnet_CNN_VGG_FIVEEMO_FINE_FINAL.h5")

# Mapping of emotion labels to their indices
emotion_label_dict = {
    0: 'neutral',
    1: 'happiness',
    2: 'sadness',
    3: 'surprise',
    4: 'fear',
}

# Function to detect faces in an image
def detect_face(image):
    img  =image
    face = face_detector.detectMultiScale(img, 1.1, 5, minSize=(40, 40))

    if len(face) > 0:
        x, y, w, h = face[0]
        crop_img = img[y:y+h, x:x+w]
        cropped = cv2.resize(crop_img, (224, 224))
        img_rgb = cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB)
        return img_rgb
    else:
        print("No face detected.")
        return None

# Function to classify emotion using the loaded model
def classify_emotion(image):
    # Preprocess the image
    image = detect_face(image)
    image = np.array(image)
    image = np.expand_dims(image, axis=0)
    # image = image / 255.0

    # Make prediction using the model
    predictions = model.predict(image)

    predictions = tf.nn.softmax(predictions)
    print(predictions)
    emotion_index = np.argmax(predictions)

    emotion_name = emotion_label_dict[emotion_index]

    return emotion_name

# Streamlit app
def main():
    st.title("Emotion Prediction App")

    uploaded_file = st.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        image_array = np.array(image)
        detected_face = detect_face(image_array)

        if detected_face is not None:
            predicted_emotion = classify_emotion(detected_face)
            st.write('Predicted Emotion:', predicted_emotion)

if __name__ == '__main__':
    main()