Update app.py

app.py

@@ -21,36 +21,29 @@ model_path = hf_hub_download(repo_id="shaheer-data/Yellow-Rust-Prediction", file
 
 loaded_model = load_model(model_path)  # 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"])
+# Function to preprocess the uploaded image
+def preprocess_image(image):
+    # Resize the image to match the model input size (e.g., 224x224 for many pre-trained models)
+    image = image.resize((224, 224))  # Adjust size based on your model input
+    image = img_to_array(image)       # Convert image to numpy array
+    image = image / 255.0            # Normalize pixel values to [0, 1]
+    image = np.expand_dims(image, axis=0)  # Add batch dimension
+    return image
+
+# Streamlit file uploader
+uploaded_file = st.file_uploader("Upload a wheat leaf 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)
 
-    # 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}")
+    # Preprocess the image
+    processed_image = preprocess_image(image)
 
-else:
-    st.write("Please upload an image file to make predictions.")
+    # Perform prediction
+    with st.spinner("Predicting..."):
+        prediction = loaded_model.predict(processed_image)
+        predicted_class = np.argmax(prediction, axis=1)[0]  # Get the class index
+        class_labels = ['0', 'MR', 'MRMS', 'MS', 'R', 'S']  # Update based on your classes
+        st.success(f"Predicted Severity Class: {class_labels[predicted_class]}")