Update app.py

app.py

@@ -1,15 +1,18 @@
 import streamlit as st
-import pandas as pd
 import numpy as np
 import cv2
 from PIL import Image
 import tensorflow as tf
 
 # Function to load the model
-@st.cache_resource
+@st.cache
 def load_model():
-    model = tf.keras.models.load_model('path_to_your_saved_model.h5')  # Provide the path to your model
-    return model
+    try:
+        model = tf.keras.models.load_model('path_to_your_saved_model.h5')  # Provide the path to your model
+        return model
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return None
 
 # Function to preprocess the image
 def preprocess_image(image):
@@ -32,6 +35,10 @@ def main():
 
     model = load_model()
 
+    if model is None:
+        st.write("Model could not be loaded.")
+        return
+
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
     if uploaded_file is not None:
         image = Image.open(uploaded_file)
@@ -41,7 +48,9 @@ def main():
         st.write("Classifying...")
         prediction = predict(image, model)
         
-        st.write(f"Predicted class: {np.argmax(prediction)}")  # Update with your model's prediction logic
+        # Adjust based on your model's output
+        predicted_class = np.argmax(prediction, axis=1)[0]  # Assuming your model outputs probabilities for each class
+        st.write(f"Predicted class: {predicted_class}")
 
 if __name__ == "__main__":
     main()