Ghiffari4869's picture
Create app.py
b6198b2 verified
# 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 ----