import gradio as gr
from transformers import pipeline

MODEL_AGE = pipeline('image-classification', model='nateraw/vit-age-classifier', device=-1)

def predict(image):
    age_result = MODEL_AGE(image)
    top_age = age_result[0]['label']
    age_map = {
        "3-9": 6,
        "10-19": 15,
        "20-29": 25,
        "30-39": 35,
        "40-49": 45,
        "50-59": 55,
        "60-69": 65,
        "more than 70": 75
    }
    return {"data": [age_map.get(top_age, 30)]}

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Json()
)

iface.launch()