Update app.py

app.py

@@ -1,25 +1,23 @@
+import os
 import streamlit as st
 import requests
 import numpy as np
 from PIL import Image
-from tensorflow import keras
+import pickle  # Using pickle since the model is saved as a .pkl file
 
-# Cache the model download and loading to avoid reloading on every interaction
+# Define the absolute path for the model
+MODEL_PATH = "/app/trained_model.pkl"
+
+# Ensure the model has the correct read permissions
+if os.path.exists(MODEL_PATH):
+    os.chmod(MODEL_PATH, 0o644)
+
+# Cache model loading to avoid repeated downloads
 @st.cache_resource
 def load_model():
-    # URL of the model file on Hugging Face
-    model_url = "https://huggingface.co/spaces/langatnaomi10/CNN/resolve/main/trained_model.pkl"
-    
-    # Download the file
-    model_path = "trained_model.pkl"
-    response = requests.get(model_url, stream=True)
-    with open(model_path, "wb") as file:
-        for chunk in response.iter_content(chunk_size=1024):
-            if chunk:
-                file.write(chunk)
-    
-    # Load the saved model
-    model = keras.models.load_model(model_path)
+    # Load the trained model from the saved .pkl file
+    with open(MODEL_PATH, "rb") as file:
+        model = pickle.load(file)
     return model
 
 # Load the model
@@ -29,7 +27,7 @@ except Exception as e:
     st.error(f"Failed to load the model: {e}")
     st.stop()
 
-# Define a prediction function
+# Define the prediction function
 def predict_axle_configuration(image):
     # Resize and preprocess the image
     image = image.resize((128, 128))  # Resize to match model input size
@@ -49,12 +47,11 @@ if uploaded_file:
         img = Image.open(uploaded_file)
         st.image(img, caption='Uploaded Image', use_column_width=True)
         st.write("Classifying...")
-        
+
         # Get prediction
         result = predict_axle_configuration(img)
-        
+
         # Display prediction (assuming result is a probability or class index)
         st.write(f"Predicted Axle Configuration: {result}")
     except Exception as e:
         st.error(f"An error occurred during prediction: {e}")
-