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

# Load processor and model
processor = AutoImageProcessor.from_pretrained("RickyIG/emotion_face_image_classification")
model = AutoModelForImageClassification.from_pretrained("RickyIG/emotion_face_image_classification")

# Title of the Streamlit app
st.title("Emotion Detection App")

# Option to choose between uploading image or using live camera
option = st.radio("Select an option", ("Upload Image", "Use Live Camera"))

if option == "Upload Image":
    # Upload image
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    
    if uploaded_file is not None:
        # Display the uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)

        # Preprocess the image
        inputs = processor(images=image, return_tensors="pt")
        
        # Make predictions
        with torch.no_grad():
            outputs = model(**inputs)
        logits = outputs.logits  # raw model outputs (before softmax)
        predicted_class_idx = logits.argmax(-1).item()  # predicted class index
        
        # Get the label of the predicted class
        label = model.config.id2label[predicted_class_idx]

        # Display the result
        st.write(f"Predicted Emotion: {label}")

elif option == "Use Live Camera":
    # Use OpenCV to capture video from the front camera
    cap = cv2.VideoCapture(0)

    if not cap.isOpened():
        st.error("Error: Could not open webcam.")
    else:
        stframe = st.empty()  # Placeholder to display live camera feed
        
        while True:
            # Capture frame-by-frame
            ret, frame = cap.read()
            
            if not ret:
                st.error("Error: Failed to capture frame.")
                break
            
            # Convert frame (BGR) to RGB (PIL format)
            image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            
            # Preprocess the image
            inputs = processor(images=image, return_tensors="pt")
            
            # Make predictions
            with torch.no_grad():
                outputs = model(**inputs)
            logits = outputs.logits  # raw model outputs (before softmax)
            predicted_class_idx = logits.argmax(-1).item()  # predicted class index
            
            # Get the label of the predicted class
            label = model.config.id2label[predicted_class_idx]

            # Display the result
            cv2.putText(frame, f"Emotion: {label}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)
            
            # Convert the frame to RGB for Streamlit
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            stframe.image(frame_rgb, channels="RGB", use_column_width=True)

    # Release the capture when finished
    cap.release()