Update app.py

app.py

@@ -1,55 +1,33 @@
 import streamlit as st
-import numpy as np
-import tensorflow as tf
-from tensorflow.keras.models import load_model
+import torch
+from transformers import pipeline
 from PIL import Image
 import os
-
-def load_model_from_huggingface():
-    """Loads the model from Hugging Face Hub."""
-    from huggingface_hub import from_pretrained_keras
-    model = from_pretrained_keras("shaheer-data/Yellow-Rust-Prediction")
-    return model
-
-def preprocess_image(image):
-    """Preprocesses the uploaded image for model prediction."""
-    image = image.resize((224, 224))  # Assuming input size of 224x224 for the model
-    image_array = np.array(image)
-    image_array = image_array / 255.0  # Normalize pixel values to [0, 1]
-    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
-    return image_array
-
-def predict_severity(model, image):
-    """Predicts the severity using the model."""
-    predictions = model.predict(image)
-    class_names = ['0', 'MR', 'MRMS', 'MS', 'R', 'S']
-    predicted_class = class_names[np.argmax(predictions)]
-    return predicted_class, predictions
-
-# Streamlit App
-st.title("Disease Severity Prediction App")
-st.write("Upload an image to predict the severity of the disease.")
-
-# Image Upload
-uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-
-if uploaded_image:
-    st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
-
-    # Load model
-    with st.spinner("Loading model..."):
-        model = load_model_from_huggingface()
-
-    # Preprocess image
-    image = Image.open(uploaded_image)
-    preprocessed_image = preprocess_image(image)
-
-    # Predict severity
-    with st.spinner("Predicting severity..."):
-        predicted_class, prediction_scores = predict_severity(model, preprocessed_image)
-
-    # Display results
-    st.success(f"Predicted Class: {predicted_class}")
-    st.write("Prediction Scores:")
-    for class_name, score in zip(['0', 'MR', 'MRMS', 'MS', 'R', 'S'], prediction_scores[0]):
-        st.write(f"{class_name}: {score:.4f}")
+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}")