import gradio as gr
from transformers import pipeline

MODEL_NAME = "nexusbert/tomato-disease-vit"
classifier = pipeline("image-classification", model=MODEL_NAME)

def classify_tomato(image):
    if image is None:
        return "Please upload an image"
    
    predictions = classifier(image)
    predictions = sorted(predictions, key=lambda x: x['score'], reverse=True)
    
    result = "## 🍅 Classification Results\n\n"
    result += f"**Top Prediction:** {predictions[0]['label']}\n\n"
    result += f"**Confidence:** {predictions[0]['score']*100:.2f}%\n\n"
    result += "### All Predictions:\n"
    for pred in predictions:
        result += f"- **{pred['label']}**: {pred['score']*100:.2f}%\n"
    
    return result

demo = gr.Interface(
    fn=classify_tomato,
    inputs=gr.Image(type="pil", label="Upload Tomato Leaf Image"),
    outputs=gr.Markdown(label="Classification Results"),
    title="🍅 Tomato Disease Classifier",
    description="Classify tomato leaf diseases. Upload an image to detect Early Blight, Late Blight, or Healthy status.",
    theme="soft",
)

if __name__ == "__main__":
    demo.launch()