Update app.py

app.py

@@ -1,57 +1,93 @@
-# Import necessary libraries
-import tensorflow as tf
-from tensorflow.keras.models import load_model
-from tensorflow.keras.preprocessing import image
-import numpy as np
-from google.colab import files
+import streamlit as st
 import cv2
-from google.colab.patches import cv2_imshow
+import numpy as np
+from tensorflow.keras.models import load_model
 import pickle
+from PIL import Image
+import os
+
+# Load the model and label encoder (cached for performance)
+@st.cache_resource
+def load_resources():
+    try:
+        model = load_model('captains_cv2_model.keras', compile=False)
+        with open('label_encoder.pkl', 'rb') as file:
+            label_encoder = pickle.load(file)
+        return model, label_encoder
+    except Exception as e:
+        st.error(f"Error loading resources: {str(e)}")
+        return None, None
+
+# Preprocess the image (adjusted to match your Colab code)
+def preprocess_image(image_path, target_size=(64, 64)):
+    try:
+        img = cv2.imread(image_path)
+        if img is None:
+            raise ValueError("Failed to load image")
+        img = cv2.resize(img, target_size)  # Resize to 64x64
+        img_array = np.asarray(img)         # Convert to numpy array, shape: (64, 64, 3)
+        img_array = img_array[np.newaxis, :, :, :]  # Add batch dimension, shape: (1, 64, 64, 3)
+        return img_array
+    except Exception as e:
+        st.error(f"Error preprocessing image: {str(e)}")
+        return None
 
-# Upload your model and label encoder files to Colab
-print("Please upload your captains_cv2_model.keras file")
-uploaded_model = files.upload()
-model_filename = list(uploaded_model.keys())[0]
+# Main Streamlit app
+def main():
+    # Load model and label encoder
+    model, label_encoder = load_resources()
+    if model is None or label_encoder is None:
+        st.error("Failed to load model or label encoder. Please check the files and try again.")
+        return
 
-print("Please upload your label_encoder.pkl file")
-uploaded_encoder = files.upload()
-encoder_filename = list(uploaded_encoder.keys())[0]
+    # UI setup
+    st.title("Image Classification App")
+    st.write("Upload an image to get a prediction using the pre-trained CNN model.")
 
-# Load the pretrained model
-model = load_model(model_filename)
+    # File uploader
+    uploaded_file = st.file_uploader("Choose an image to classify...", type=["jpg", "jpeg", "png"])
 
-# Load the label encoder
-with open(encoder_filename, 'rb') as file:
-    label_encoder = pickle.load(file)
+    if uploaded_file is not None:
+        # Display the uploaded image
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
 
-# Function to preprocess the image (adjust based on your model's requirements)
-def preprocess_image(img_path, target_size=(64, 64)):  # Adjust target_size as per your model
-    
-    img = cv2.imread(img_path)
-    img_array = np.asarray(img)
-    img_array = img_array[np.newaxis, :, :, :]
-    return img_array
+        # Save the uploaded file temporarily with its original extension
+        file_extension = os.path.splitext(uploaded_file.name)[1].lower()
+        temp_filename = f"temp_image{file_extension}"
+        with open(temp_filename, "wb") as f:
+            f.write(uploaded_file.getvalue())
 
-# Upload an image to predict
-print("Please upload an image to classify")
-uploaded_image = files.upload()
+        try:
+            # Preprocess the image
+            processed_image = preprocess_image(temp_filename)
+            if processed_image is None:
+                raise ValueError("Image preprocessing failed")
 
-# Get the uploaded image filename
-img_filename = list(uploaded_image.keys())[0]
+            # Display processed image shape for debugging
+            st.write(f"Processed image shape: {processed_image.shape}")
 
-# Preprocess the image
-processed_image = preprocess_image(img_filename)
+            # Make prediction
+            prediction = model.predict(processed_image)
+            predicted_class_index = np.argmax(prediction, axis=1)[0]
+            predicted_class = label_encoder.inverse_transform([predicted_class_index])[0]
 
-# Make prediction
-prediction = model.predict(processed_image)
+            # Display results
+            st.subheader("Prediction Results")
+            st.write(f"**Predicted Class:** {predicted_class}")
+            st.write(f"**Predicted Class Index:** {predicted_class_index}")
+            
+            # Display prediction probabilities
+            st.write("**Prediction Probabilities:**")
+            for class_name, prob in zip(label_encoder.classes_, prediction[0]):
+                st.write(f"{class_name}: {prob:.4f}")
 
-# Get the predicted class index
-predicted_class_index = np.argmax(prediction, axis=1)[0]
+        except Exception as e:
+            st.error(f"Error during prediction: {str(e)}")
 
-# Decode the prediction using the label encoder
-predicted_class = label_encoder.inverse_transform([predicted_class_index])[0]
+        # Clean up temporary file
+        if os.path.exists(temp_filename):
+            os.remove(temp_filename)
 
-# Output the prediction
-print("Prediction probabilities:", prediction)
-print("Predicted class index:", predicted_class_index)
-print("Predicted class:", predicted_class)
+if __name__ == "__main__":
+    main()