Update app.py

app.py

@@ -8,8 +8,7 @@ import os
 # Function to load the model
 @st.cache_resource
 def load_model():
-    # Ensure this path is correct relative to your deployment setup
-    model_path = 'path_to_your_saved_model.h5'  
+    model_path = 'path_to_your_saved_model.h5'  # Ensure this path is correct relative to your deployment setup
     
     if not os.path.isfile(model_path):
         st.error(f"Model file not found: {model_path}")
@@ -17,6 +16,7 @@ def load_model():
     
     try:
         model = tf.keras.models.load_model(model_path)
+        st.success("Model loaded successfully!")
         return model
     except Exception as e:
         st.error(f"Error loading model: {e}")
@@ -54,11 +54,12 @@ def main():
         
         st.write("")
         st.write("Classifying...")
-        prediction = predict(image, model)
-        
-        # 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}")
+        try:
+            prediction = predict(image, model)
+            predicted_class = np.argmax(prediction, axis=1)[0]  # Adjust based on your model's output
+            st.write(f"Predicted class: {predicted_class}")
+        except Exception as e:
+            st.error(f"Error in prediction: {e}")
 
 if __name__ == "__main__":
     main()