import gradio as gr
import cv2
import os
import numpy as np
from deepface import DeepFace

# Define the folder path where images will be saved
dataset_path = "/content/dataset"  # For Colab, this will save in your Colab environment

# Ensure the directory exists
if not os.path.exists(dataset_path):
    os.makedirs(dataset_path)

# Function to capture, save image, and predict emotion
def capture_and_predict(image, name):
    # Convert Gradio image (PIL format) to an OpenCV image
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    
    # Analyze the emotion using DeepFace
    result = DeepFace.analyze(img, actions=['emotion'])
    
    # Get the dominant emotion
    dominant_emotion = result[0]['dominant_emotion']
    
    # Save the image with a timestamp in the dataset folder
    person_folder = os.path.join(dataset_path, name)
    os.makedirs(person_folder, exist_ok=True)  # Create a folder for each person if not exists
    
    image_count = len(os.listdir(person_folder))
    image_path = os.path.join(person_folder, f"{image_count + 1}.jpg")
    cv2.imwrite(image_path, cv2.cvtColor(img, cv2.COLOR_BGR2RGB))  # Save the image in RGB format for consistency
    
    return f"Image saved for {name} with emotion: {dominant_emotion} at {image_path}"

# Define the Gradio interface
iface = gr.Interface(
    fn=capture_and_predict,
    inputs=[gr.Image(type="pil"), gr.Textbox(label="Enter your name")],
    outputs="text",
    title="Capture and Predict Facial Emotion",
    description="Capture an image from your webcam, enter your name, and the system will predict your emotion and save the image.",
    live=True
)

# Launch the Gradio app
iface.launch()