Spaces:
Sleeping
Sleeping
File size: 3,501 Bytes
5a5b14c 1b7540f 5a5b14c a18495a 5a5b14c a18495a 5a5b14c a18495a 5a5b14c a18495a 5a5b14c 1b7540f 5a5b14c 1b7540f 5a5b14c 1b7540f 5a5b14c 1b7540f 5a5b14c 1b7540f 5a5b14c 1b7540f 5a5b14c 1b7540f 5a5b14c 1b7540f 5a5b14c a18495a 5a5b14c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
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)
@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
# 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() |