import gradio as gr
from PIL import Image
import requests
from io import BytesIO
import os
from transformers import pipeline

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

def classify_image(input_img=None, img_url=None):
    if img_url: 
        try:
            response = requests.get(img_url)
            response.raise_for_status()
            img = Image.open(BytesIO(response.content)).convert("RGB")
        except Exception as e:
            return {"Error": f"Failed to load image from URL: {e}"}
    elif input_img:  
        if not isinstance(input_img, Image.Image):
            img = Image.fromarray(input_img)
        else:
            img = input_img
    else:
        return {"Error": "No image provided."}

    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("## Soil Type Classifier")
    gr.Markdown("Upload an image of a soil or enter an image URL to classify it.")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Soil Image")
        url_input = gr.Textbox(label="Soil Image URL (optional)", placeholder="Paste image URL here...")
        label_output = gr.Label(num_top_classes=3, label="Predictions")

    classify_btn = gr.Button("Submit", variant="primary")

    gr.Examples(
        examples=[f'./images_samples/{i}' for i in os.listdir('./images_samples')],
        inputs=image_input,
        outputs=label_output,
        fn=classify_image,
        cache_examples=False
    )

    classify_btn.click(
        fn=classify_image,
        inputs=[image_input, url_input],
        outputs=label_output
    )

demo.launch(share=True)