Update app.py

app.py

@@ -1,118 +1,87 @@
 import streamlit as st
-import numpy as np
 import cv2
+import numpy as np
 from tensorflow.keras.models import load_model
 import pickle
 from PIL import Image
 import os
 
-# Set page configuration
-st.set_page_config(
-    page_title="Image Detection App",
-    page_icon="📸",
-    layout="centered",
-    initial_sidebar_state="expanded"
-)
-
-# Load the trained model and label encoder with error handling
+# Load the model and label encoder
 @st.cache_resource
 def load_resources():
-    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)
+    model = load_model('captains_cv2_model.keras')
+    with open('label_encoder.pkl', 'rb') as file:
+        le = pickle.load(file)
     return model, le
 
-# Load resources
-model, label_encoder = load_resources()
-
-# Function to preprocess the uploaded image
-def preprocess_image(uploaded_file):
-    # Save the uploaded file temporarily to disk
-    temp_path = "temp_image.jpg"
-    with open(temp_path, "wb") as f:
-        f.write(uploaded_file.read())
+# Preprocess the image
+def preprocess_image(image_path):
+    # Read and convert image
+    img1 = cv2.imread(image_path)
+    img1 = cv2.resize(img1, (64, 64))  # Resize to 64x64
+    img1 = np.asarray(img1)  # Convert to numpy array, shape will be (64, 64, 3)
     
-    # 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.")
+    # Add batch dimension to get (1, 64, 64, 3)
+    img1 = img1[np.newaxis, :, :, :]
     
-    # Resize to the model's expected input size (64, 64)
-    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, :, :, :]
+    # Verify shape
+    if len(img1.shape) != 4:
+        raise ValueError(f"Expected shape (batch_size, 64, 64, 3), got {img1.shape}")
+    if img1.shape[1:] != (64, 64, 3):
+        raise ValueError(f"Image dimensions should be (64, 64, 3), got {img1.shape[1:]}")
     
-    # Clean up the temporary file
-    os.remove(temp_path)
-    return img
+    return img1
 
-# Sidebar
-st.sidebar.title("About")
-st.sidebar.info(
-    "This app uses a Convolutional Neural Network (CNN) to classify images into one of 10 categories. "
-    "Upload an image, and the model will predict its class!"
-)
-st.sidebar.markdown("### Classes")
-st.sidebar.write(
-    "The model can predict: lifeboat, ladybug, pizza, bell pepper, school bus, koala, espresso, red panda, orange, sports car."
-)
-
-# Main content
-st.title("📸 Image Classification App")
-st.markdown("Upload an image below, and let the model predict its class!")
-
-# File uploader
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-    try:
-        # Display the uploaded image
+# Main app
+def main():
+    # Load resources
+    model, le = load_resources()
+    
+    # Streamlit UI
+    st.title("Image Classification App")
+    st.write("Upload an image to get a prediction")
+    
+    # File uploader
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    
+    if uploaded_file is not None:
+        # Display 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)
-
-        # 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]
-
-        # 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)
-
-    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.")
+        st.image(image, caption='Uploaded Image', use_column_width=True)
+        
+        # Get original file extension
+        file_extension = os.path.splitext(uploaded_file.name)[1].lower()
+        temp_filename = f"temp_image{file_extension}"
+        
+        # Save temporary file with original extension
+        with open(temp_filename, "wb") as f:
+            f.write(uploaded_file.getvalue())
+        
+        try:
+            # Preprocess image
+            processed_img = preprocess_image(temp_filename)
+            
+            # Display shape for debugging
+            st.write(f"Processed image shape: {processed_img.shape}")
+            
+            # Make prediction
+            prediction = model.predict(processed_img)
+            predicted_class = le.inverse_transform([np.argmax(prediction)])
+            
+            # Display prediction
+            st.write("Prediction:", predicted_class[0])
+            
+            # Display prediction probabilities
+            st.write("Prediction Probabilities:")
+            for class_name, prob in zip(le.classes_, prediction[0]):
+                st.write(f"{class_name}: {prob:.4f}")
+                
+        except Exception as e:
+            st.error(f"An error occurred: {str(e)}")
+        
+        # Clean up temporary file
+        if os.path.exists(temp_filename):
+            os.remove(temp_filename)
 
-# Footer
-st.markdown("---")
-st.markdown("Created with ❤️ using Streamlit | Hosted on [Hugging Face Spaces](https://huggingface.co/spaces)")
+if __name__ == '__main__':
+    main()