Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import cv2 | |
| 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) | |
| 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 | |
| # 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 | |
| # UI setup | |
| st.title("Image Classification App") | |
| st.write("Upload an image to get a prediction using the pre-trained CNN model.") | |
| # File uploader | |
| uploaded_file = st.file_uploader("Choose an image to classify...", type=["jpg", "jpeg", "png"]) | |
| 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) | |
| # 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()) | |
| try: | |
| # Preprocess the image | |
| processed_image = preprocess_image(temp_filename) | |
| if processed_image is None: | |
| raise ValueError("Image preprocessing failed") | |
| # Display processed image shape for debugging | |
| st.write(f"Processed image shape: {processed_image.shape}") | |
| # 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] | |
| # 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}") | |
| except Exception as e: | |
| st.error(f"Error during prediction: {str(e)}") | |
| # Clean up temporary file | |
| if os.path.exists(temp_filename): | |
| os.remove(temp_filename) | |
| if __name__ == "__main__": | |
| main() |