Update app.py

app.py

@@ -15,9 +15,20 @@ EXPECTED_SIZE = (64, 64)  # Update this based on your model's input shape
 
 def load_resources():
     """Load model and label encoder."""
-    model = tf.keras.models.load_model(MODEL_PATH)
-    with open(LABEL_ENCODER_PATH, "rb") as f:
-        label_encoder = pickle.load(f)
+    try:
+        # Try loading with the newer Keras API first
+        model = tf.keras.models.load_model(MODEL_PATH, compile=False)
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None, None
+    
+    try:
+        with open(LABEL_ENCODER_PATH, "rb") as f:
+            label_encoder = pickle.load(f)
+    except Exception as e:
+        st.error(f"Error loading label encoder: {str(e)}")
+        return None, None
+    
     return model, label_encoder
 
 # Load resources
@@ -27,15 +38,29 @@ def preprocess_image(image):
     """Resize image to match model input shape."""
     image = image.resize(EXPECTED_SIZE)  # Resize to match model input
     image_array = np.array(image)  # Convert to numpy array
+    
+    # Ensure image has 3 channels (for RGB)
+    if len(image_array.shape) == 2:  # Grayscale image
+        image_array = np.stack((image_array,)*3, axis=-1)
+    elif image_array.shape[2] == 4:  # RGBA image
+        image_array = image_array[:, :, :3]
+    
+    image_array = image_array.astype('float32') / 255.0  # Normalize to [0,1]
     image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
     return image_array
 
 def predict(image):
     """Predict the class of the uploaded image."""
+    if model is None or label_encoder is None:
+        return "Model or label encoder not loaded properly"
+    
     image_array = preprocess_image(image)
-    preds = model.predict(image_array)
-    class_index = np.argmax(preds)
-    return label_encoder.inverse_transform([class_index])[0]
+    try:
+        preds = model.predict(image_array)
+        class_index = np.argmax(preds)
+        return label_encoder.inverse_transform([class_index])[0]
+    except Exception as e:
+        return f"Error during prediction: {str(e)}"
 
 # Streamlit UI
 st.set_page_config(page_title="Image Classifier", layout="wide")
@@ -53,5 +78,8 @@ if uploaded_file:
     st.image(image, caption="📸 Uploaded Image", use_column_width=True)
     
     if st.button("🔍 Classify Image", use_container_width=True):
-        prediction = predict(image)
-        st.success(f"🎯 Predicted Class: {prediction}")
+        if model is None or label_encoder is None:
+            st.error("Model or label encoder failed to load. Please check the files.")
+        else:
+            prediction = predict(image)
+            st.success(f"🎯 Predicted Class: {prediction}")