Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import numpy as np | |
| import cv2 | |
| 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 | |
| def load_resources(): | |
| model = load_model("captains_cv2_model.keras") | |
| with open("label_encoder.pkl", "rb") as f: | |
| le = pickle.load(f) | |
| return model, le | |
| 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()) | |
| # Read the image using cv2.imread | |
| img = cv2.imread(temp_path) | |
| # Resize to the model's expected input size (64, 64) | |
| img = cv2.resize(img, (64, 64)) # Note: cv2.resize takes (width, height), not (height, width, channels) | |
| # Add new axis for batch dimension | |
| img = img[np.newaxis, :, :, :] | |
| # Clean up the temporary file | |
| os.remove(temp_path) | |
| return img | |
| # 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: | |
| # Display the 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..."): | |
| # Predict and decode as per your specified steps | |
| 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) | |
| else: | |
| st.info("Please upload an image to get started.") | |
| # Footer | |
| st.markdown("---") | |
| st.markdown("Created with ❤️ using Streamlit | Hosted on [Hugging Face Spaces](https://huggingface.co/spaces)") |