Update app.py

app.py

@@ -14,14 +14,26 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-# Load the trained model and label encoder
+# Load the trained model and label encoder with error handling
 @st.cache_resource
 def load_resources():
-    model = load_model("captains_cv2_model.keras")
+    try:
+        # Load the model (assuming TensorFlow 2.6+ with batch_shape support)
+        model = load_model("captains_cv2_model.keras")
+    except TypeError as e:
+        # Fallback for compatibility issues
+        st.error(f"Model loading failed: {e}")
+        st.warning("Attempting to load model without compilation...")
+        model = load_model("captains_cv2_model.keras", compile=False)
+        # Recompile the model manually if needed
+        model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
+    
+    # Load the label encoder
     with open("label_encoder.pkl", "rb") as f:
         le = pickle.load(f)
     return model, le
 
+# Load resources
 model, label_encoder = load_resources()
 
 # Function to preprocess the uploaded image
@@ -33,9 +45,16 @@ def preprocess_image(uploaded_file):
     
     # Read the image using cv2.imread
     img = cv2.imread(temp_path)
+    if img is None:
+        raise ValueError("Failed to load image. Please ensure the file is a valid image.")
+    
     # Resize to the model's expected input size (64, 64)
-    img = cv2.resize(img, (64, 64))  # Note: cv2.resize takes (width, height), not (height, width, channels)
-    # Add new axis for batch dimension
+    img = cv2.resize(img, (64, 64))  # cv2 uses (width, height)
+    # Convert BGR (OpenCV default) to RGB if needed
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    # Normalize pixel values to [0, 1] (common for CNNs)
+    img = img / 255.0
+    # Add batch dimension
     img = img[np.newaxis, :, :, :]
     
     # Clean up the temporary file
@@ -61,31 +80,35 @@ st.markdown("Upload an image below, and let the model predict its class!")
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    # Display the uploaded image
-    image = Image.open(uploaded_file)
-    uploaded_file.seek(0)  # Reset file pointer after reading for display
-    st.image(image, caption="Uploaded Image", use_column_width=True)
+    try:
+        # Display the uploaded image
+        image = Image.open(uploaded_file)
+        uploaded_file.seek(0)  # Reset file pointer after reading for display
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+
+        # Preprocess the image
+        processed_image = preprocess_image(uploaded_file)
 
-    # Preprocess the image
-    processed_image = preprocess_image(uploaded_file)
+        # Make prediction
+        with st.spinner("Predicting..."):
+            prediction = model.predict(processed_image)
+            predicted_class_idx = np.argmax(prediction, axis=1)[0]
+            predicted_class = label_encoder.inverse_transform([predicted_class_idx])[0]
 
-    # Make prediction
-    with st.spinner("Predicting..."):
-        # Predict and decode as per your specified steps
-        prediction = model.predict(processed_image)
-        predicted_class_idx = np.argmax(prediction, axis=1)[0]
-        predicted_class = label_encoder.inverse_transform([predicted_class_idx])[0]
+        # Display the prediction
+        st.success("Prediction complete!")
+        st.markdown(f"### Predicted Class: **{predicted_class}**")
+        st.write(f"Prediction Confidence: {prediction[0][predicted_class_idx]:.4f}")
 
-    # Display the prediction
-    st.success("Prediction complete!")
-    st.markdown(f"### Predicted Class: **{predicted_class}**")
-    st.write(f"Prediction Confidence: {prediction[0][predicted_class_idx]:.4f}")
+        # Optional: Display confidence scores for all classes
+        if st.checkbox("Show confidence scores for all classes"):
+            class_names = label_encoder.classes_
+            confidence_scores = {class_names[i]: float(prediction[0][i]) for i in range(len(class_names))}
+            st.bar_chart(confidence_scores)
 
-    # Optional: Display confidence scores for all classes
-    if st.checkbox("Show confidence scores for all classes"):
-        class_names = label_encoder.classes_
-        confidence_scores = {class_names[i]: float(prediction[0][i]) for i in range(len(class_names))}
-        st.bar_chart(confidence_scores)
+    except Exception as e:
+        st.error(f"An error occurred: {e}")
+        st.info("Please try uploading a different image or check the model compatibility.")
 
 else:
     st.info("Please upload an image to get started.")