Update app.py

app.py

@@ -1,29 +1,94 @@
-import gradio as gr
-from datetime import datetime
-import pytz
-from ocr_engine import extract_weight_from_image
+import streamlit as st
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import io
+import logging
 
-def process_image(img):
-    if img is None:
-        return "No image uploaded", None, None
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-    ist_time = datetime.now(pytz.timezone("Asia/Kolkata")).strftime("%d-%m-%Y %I:%M:%S %p")
-    weight, confidence = extract_weight_from_image(img)
-    return f"{weight} kg (Confidence: {confidence}%)", ist_time, img
+# Load the pre-trained MNIST model
+@st.cache_resource
+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
 
-with gr.Blocks(title="⚖️ Auto Weight Logger") as demo:
-    gr.Markdown("## ⚖️ Auto Weight Logger")
-    gr.Markdown("📷 Upload or capture an image of a digital weight scale.")
+# 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
 
-    with gr.Row():
-        image_input = gr.Image(type="pil", label="Upload / Capture Image")
-        output_weight = gr.Textbox(label="⚖️ Detected Weight (in kg)")
+# Streamlit app
+st.title("AutoWeightLogger - Number Detection")
+st.write("Upload an image containing a single handwritten digit to detect the number.")
 
-    with gr.Row():
-        timestamp = gr.Textbox(label="🕒 Captured At (IST)")
-        snapshot = gr.Image(label="📸 Snapshot Image")
+# File uploader
+uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
 
-    submit = gr.Button("🔍 Detect Weight")
-    submit.click(process_image, inputs=image_input, outputs=[output_weight, timestamp, snapshot])
+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)
 
-demo.launch()
+        # 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.
+""")