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 | |
| def load_resources(): | |
| # Custom loading to handle compatibility | |
| try: | |
| model = load_model('captains_cv2_model.keras', compile=False) # Load without compiling first | |
| except Exception as e: | |
| st.error(f"Model loading failed: {str(e)}") | |
| raise | |
| with open('label_encoder.pkl', 'rb') as file: | |
| le = pickle.load(file) | |
| return model, le | |
| # Preprocess the image | |
| def preprocess_image(image_path): | |
| img1 = cv2.imread(image_path) | |
| img1 = cv2.resize(img1, (64, 64)) # Resize to 64x64 | |
| img1 = np.asarray(img1) # Shape: (64, 64, 3) | |
| img1 = img1[np.newaxis, :, :, :] # Shape: (1, 64, 64, 3) | |
| return img1 | |
| # Main app | |
| def main(): | |
| model, le = load_resources() | |
| st.title("Image Classification App") | |
| st.write("Upload an image to get a prediction") | |
| uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
| if uploaded_file is not None: | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption='Uploaded Image', use_column_width=True) | |
| 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: | |
| processed_img = preprocess_image(temp_filename) | |
| st.write(f"Processed image shape: {processed_img.shape}") | |
| prediction = model.predict(processed_img) | |
| predicted_class = le.inverse_transform([np.argmax(prediction)]) | |
| st.write("Prediction:", predicted_class[0]) | |
| 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)}") | |
| if os.path.exists(temp_filename): | |
| os.remove(temp_filename) | |
| if __name__ == '__main__': | |
| main() |