import streamlit as st
import torch
from transformers import pipeline
from PIL import Image
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Load Hugging Face model
model_url = os.getenv('HUGGINGFACE_MODEL_URL')
model = torch.hub.load(model_url, 'model', source='hf')

# Setup Streamlit
st.title('Yellow Rust Severity Prediction')

# File uploader
uploaded_file = st.file_uploader("Upload an image of Yellow Rust", type=["jpg", "png"])

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    # Process the image and make prediction
    classifier = pipeline('image-classification', model=model_url)
    results = classifier(image)
    
    severity_level = results[0]['label']
    confidence = results[0]['score']
    
    st.write(f"Predicted Severity Level: {severity_level} with confidence: {confidence:.2f}")