import gradio as gr
from transformers import pipeline
from PIL import Image

# 1. Initialize the pipeline
pipe = pipeline("image-classification", model="dima806/fairface_age_image_detection")

# 2. Define the prediction function
def predict_age(input_img):
    # The pipeline can take a PIL image directly
    results = pipe(input_img)
    
    # Format the results for Gradio's Label component
    # Returns a dictionary like {"20-29": 0.85, "30-39": 0.10, ...}
    return {result['label']: result['score'] for result in results}

# 3. Build the Gradio Interface
demo = gr.Interface(
    fn=predict_age,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=gr.Label(num_top_classes=3, label="Predicted Age Range"),
    title="FairFace Age Detection",
    description="Upload a photo to estimate the age range of the person in the image."
)

# 4. Launch the app
if __name__ == "__main__":
    demo.launch(share=True)