import gradio as gr
from transformers import AutoModelForImageClassification, AutoProcessor
import torch

# Load the model and processor from Hugging Face
model_name = "dima806/facial_age_image_detection"
model = AutoModelForImageClassification.from_pretrained(model_name)
processor = AutoProcessor.from_pretrained(model_name)

# Define the prediction function
def predict(image):
    # Process the input image
    inputs = processor(images=image, return_tensors="pt")
    # Perform the prediction
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Get the model's original outputs (e.g., logits or probabilities)
    predictions = outputs.logits

    # Convert predictions to a list and round to 2 decimal places if necessary
    predictions_list = predictions.tolist()
    rounded_predictions = [[round(pred, 2) for pred in prediction] for prediction in predictions_list]
    
    return rounded_predictions

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs="image",
    outputs="label",  # Use the model's original output type
    title="Facial Age Prediction",
    description="This application predicts your age from a facial image."
)

# Launch the Gradio application
iface.launch(share=True)