import streamlit as st
from PIL import Image
from transformers import AutoModelForImageClassification, ViTImageProcessor
import torch

# Load the model and processor
@st.cache_resource
def load_model():
    model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224-in21k")
    processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
    return model, processor

model, processor = load_model()

# Streamlit app UI
st.title("🌱 Plant Identification App 🌱")
st.write("Upload a plant image and let the app identify its species!")

# File uploader for plant image
uploaded_file = st.file_uploader("Choose a plant image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Open and display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Plant Image.", use_column_width=True)

    # Preprocess the image using the processor
    inputs = processor(images=image, return_tensors="pt", padding=True)

    # Run the classification
    with st.spinner('Classifying plant species...'):
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()

    # Get the label of the predicted class
    label = model.config.id2label[predicted_class_idx]

    # Display prediction results
    st.write(f"Predicted Species: {label}")
    st.write(f"Confidence: {torch.softmax(logits, dim=-1)[0][predicted_class_idx]*100:.2f}%")