Spaces:
Running
Running
| import streamlit as st | |
| import tensorflow as tf | |
| import numpy as np | |
| from PIL import Image | |
| import io | |
| import logging | |
| # Set up logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # Load the pre-trained MNIST model | |
| def load_model(): | |
| try: | |
| model = tf.keras.models.load_model('mnist_cnn.h5') | |
| logger.info("MNIST model loaded successfully") | |
| return model | |
| except Exception as e: | |
| logger.error(f"Error loading model: {e}") | |
| st.error("Failed to load the model. Please check the model file.") | |
| return None | |
| # Preprocess the uploaded image | |
| def preprocess_image(image): | |
| try: | |
| # Convert to grayscale | |
| img = image.convert('L') | |
| # Resize to 28x28 (MNIST model input size) | |
| img = img.resize((28, 28), Image.Resampling.LANCZOS) | |
| # Convert to numpy array and normalize | |
| img_array = np.array(img) | |
| # Ensure the image is inverted if necessary (MNIST expects white digits on black background) | |
| img_array = 255 - img_array # Invert colors | |
| img_array = img_array / 255.0 # Normalize to [0, 1] | |
| # Reshape for model input (1, 28, 28, 1) | |
| img_array = img_array.reshape(1, 28, 28, 1) | |
| logger.info("Image preprocessed successfully") | |
| return img_array | |
| except Exception as e: | |
| logger.error(f"Error preprocessing image: {e}") | |
| st.error("Failed to preprocess the image. Please ensure it's a valid image.") | |
| return None | |
| # Streamlit app | |
| st.title("AutoWeightLogger - Number Detection") | |
| st.write("Upload an image containing a single handwritten digit to detect the number.") | |
| # File uploader | |
| uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"]) | |
| if uploaded_file is not None: | |
| try: | |
| # Display the uploaded image | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_column_width=True) | |
| # Preprocess the image | |
| processed_image = preprocess_image(image) | |
| if processed_image is None: | |
| st.stop() | |
| # Load the model | |
| model = load_model() | |
| if model is None: | |
| st.stop() | |
| # Make prediction | |
| with st.spinner("Detecting number..."): | |
| prediction = model.predict(processed_image) | |
| predicted_digit = np.argmax(prediction, axis=1)[0] | |
| confidence = np.max(prediction) * 100 | |
| # Display result | |
| st.success(f"Detected Number: {predicted_digit}") | |
| st.write(f"Confidence: {confidence:.2f}%") | |
| # Provide feedback if confidence is low | |
| if confidence < 70: | |
| st.warning("Low confidence in prediction. Please ensure the image contains a clear, single handwritten digit.") | |
| except Exception as e: | |
| logger.error(f"Error processing image: {e}") | |
| st.error("An error occurred while processing the image. Please try again with a different image.") | |
| else: | |
| st.info("Please upload an image to proceed.") | |
| # Instructions for users | |
| st.markdown(""" | |
| ### Instructions | |
| 1. Upload an image containing a single handwritten digit (0-9). | |
| 2. Ensure the digit is clear, centered, and on a plain background for best results. | |
| 3. The model expects white digits on a black background, similar to MNIST dataset images. | |
| """) |