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


 
classifier = pipeline("image-classification", model="Docty/mangoes")

def classify_image(img):
    if not isinstance(img, Image.Image):
        img = Image.fromarray(img)
    results = classifier(img)
    return {res["label"]: float(res["score"]) for res in results}


theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="lime",
    neutral_hue="slate"
)


with gr.Blocks(theme=theme) as demo:
    gr.Markdown("## Mango Image Classifier")
    gr.Markdown("Upload an image of a mango to classify it using a fine-tuned model.")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Mango Image")
        label_output = gr.Label(num_top_classes=3, label="Predictions")

    classify_btn = gr.Button("Classify Image", variant="primary")
    gr.Examples(
        examples=[
            "0.jpg",
            "1.jpg",
            "2.jpg",
            "3.jpg",
            "4.jpg",
            "5.jpg",
            "6.jpg",
            "7.jpg"
        ],
        inputs=image_input,
        outputs=label_output,
        fn=classify_image,
        cache_examples=False  # set True if you want cached predictions
    )
    classify_btn.click(fn=classify_image, inputs=image_input, outputs=label_output)

demo.launch(share=True)