# Put your code here:
import gradio as gr
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification

# Load emotional classification model published in hugging face
processor_emotion = AutoImageProcessor.from_pretrained("Ghiffari4869/emotional_classification")
model_emotion = AutoModelForImageClassification.from_pretrained("Ghiffari4869/emotional_classification")

# Load age classification model
model_name_age = "nateraw/vit-age-classifier"
processor_age = AutoImageProcessor.from_pretrained(model_name_age)
model_age = AutoModelForImageClassification.from_pretrained(model_name_age)

# Define the age categories
age_categories = ['0-2', '3-9', '10-19', '20-29', '30-39', '40-49', '50-59', '60-69', 'more than 70']

# Define function to predict emotional label and age label
def predict_emotion_and_age(input_image):
    # Preprocess input image for emotional classification
    inputs_emotion = processor_emotion(images=input_image, return_tensors="pt")

    # Make predictions for emotional classification
    outputs_emotion = model_emotion(**inputs_emotion)
    logits_emotion = outputs_emotion.logits
    predicted_class_emotion = logits_emotion.argmax(dim=1).item()

    # Get emotional label and score
    emotional_labels = ["anger", "contempt", "disgust", "fear", "happy", "sad", "surprise"]
    emotional_score = logits_emotion.softmax(dim=1).max().item()
    predicted_emotion_label = emotional_labels[predicted_class_emotion]

    # Preprocess input image for age classification
    inputs_age = processor_age(images=input_image, return_tensors="pt")

    # Make predictions for age classification
    outputs_age = model_age(**inputs_age)
    logits_age = outputs_age.logits
    predicted_class_age = logits_age.argmax(dim=1).item()

    # Get age label and score
    age_score = logits_age.softmax(dim=1).max().item()
    predicted_age_label = age_categories[predicted_class_age]
    output10 = {"score": emotional_score, "label": predicted_emotion_label}
    output20 = {"score": age_score, "label": predicted_age_label}
    return output10, output20


# Define Gradio interface
iface = gr.Interface(
    fn=predict_emotion_and_age,  # Function to be called for prediction
    inputs=gr.Image(),  # Image input
    outputs=[gr.JSON(label="output 0"), gr.JSON(label="output 1")]  # JSON outputs for emotional label and age label
)

# Launch the Gradio interface
iface.launch()
# ---- End of your code ----