Update app.py

app.py

@@ -1,33 +1,60 @@
 import streamlit as st
-import torch
-from transformers import pipeline
+import tensorflow as tf
 from PIL import Image
 import os
-from dotenv import load_dotenv
+from huggingface_hub import hf_hub_url, set_access_token  # Import Hugging Face utilities
 
-# Load environment variables
-load_dotenv()
+# Title of the Streamlit app
+st.title("Yellow Rust Severity Prediction")
 
-# Load Hugging Face model
-model_url = os.getenv('HUGGINGFACE_MODEL_URL')
-model = torch.hub.load(model_url, 'model', source='hf')
+# Load Hugging Face API token from environment
+huggingface_api_token = os.getenv("HUGGINGFACE_TOKEN")
+if huggingface_api_token is None:
+    st.error("YellowRust API token not found in environment. Please set it.")
+    st.stop()
 
-# Setup Streamlit
-st.title('Yellow Rust Severity Prediction')
+# Set Hugging Face token for authentication
+set_access_token(huggingface_api_token)
 
-# File uploader
-uploaded_file = st.file_uploader("Upload an image of Yellow Rust", type=["jpg", "png"])
+# Model repository details
+model_repo_id = "your_huggingface_model_repo_id"
+model_file_path = "final_meta_model.keras"
+
+# Construct the model URL
+st.write("Loading model from Hugging Face repo:", model_repo_id)
+model_url = hf_hub_url(model_repo_id, model_file_path)
+loaded_model = tf.keras.models.load_model(model_url)  # Load model using tf.keras directly
+
+# Function to make predictions
+def predict_image(image):
+    image = image.resize((224, 224))  # Resize to match model input dimensions
+    image_array = tf.keras.preprocessing.image.img_to_array(image)
+    image_array = tf.expand_dims(image_array, axis=0)  # Expand dimensions for batch prediction
+    predictions = loaded_model.predict(image_array)
+    return predictions
+
+# Class labels for Yellow Rust severity levels
+CLASS_LABELS = [
+    "Healthy",
+    "Mild Severity",
+    "Moderate Severity",
+    "Severe Severity",
+    "Very Severe",
+    "Extreme Severity"
+]
+
+# Image upload widget
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "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}")
+
+    # Display progress bar
+    with st.spinner("Making predictions..."):
+        predictions = predict_image(image)
+        predicted_class = predictions.argmax(axis=-1)
+        st.write(f"Predicted Severity Level: {CLASS_LABELS[predicted_class[0]]} with confidence {predictions[0][predicted_class[0]]:.2f}")
+
+else:
+    st.write("Please upload an image file to make predictions.")