Update app.py

app.py

@@ -1,8 +1,13 @@
 import os
+import streamlit as st
+
+# Must be the first Streamlit command
+st.set_page_config(page_title="Image Classifier", layout="wide")
+
+# Now import other libraries
 os.system("pip install tensorflow")
 os.system("pip install scikit-learn")
 
-import streamlit as st
 import tensorflow as tf
 import numpy as np
 import pickle
@@ -16,8 +21,15 @@ EXPECTED_SIZE = (64, 64)  # Update this based on your model's input shape
 def load_resources():
     """Load model and label encoder."""
     try:
-        # Try loading with the newer Keras API first
-        model = tf.keras.models.load_model(MODEL_PATH, compile=False)
+        # Try loading with custom InputLayer configuration
+        custom_objects = {
+            'InputLayer': lambda **kwargs: tf.keras.layers.InputLayer(**{k: v for k, v in kwargs.items() if k != 'batch_shape'})
+        }
+        model = tf.keras.models.load_model(
+            MODEL_PATH,
+            compile=False,
+            custom_objects=custom_objects
+        )
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")
         return None, None
@@ -63,7 +75,6 @@ def predict(image):
         return f"Error during prediction: {str(e)}"
 
 # Streamlit UI
-st.set_page_config(page_title="Image Classifier", layout="wide")
 st.markdown("""
     <h1 style='text-align: center; color: #4A90E2;'>🖼️ Image Classification App</h1>
     <p style='text-align: center; font-size: 18px;'>Upload an image and let our model classify it for you!</p>