Update app.py

app.py

@@ -116,36 +116,47 @@ with col2:
         st.subheader("Original Drawing")
         st.image(canvas_result.image_data, use_column_width=True)
 
+# === Image preprocessing and prediction ===
+# Initialize session state for draw count
+if "draw_count" not in st.session_state:
+    st.session_state.draw_count = 0
+    st.session_state.last_image = None
+
 # === Image preprocessing and prediction ===
 if canvas_result.image_data is not None:
-    st.markdown("---")
+    current_image = canvas_result.image_data
+
+    # Compare current image with previous
+    if st.session_state.last_image is None or not np.array_equal(current_image, st.session_state.last_image):
+        st.session_state.draw_count += 1
+        st.session_state.last_image = current_image
+
+    st.markdown(f"### ✏️ Draw Count: {st.session_state.draw_count}")
+
     st.subheader("Preprocessed Image & Prediction")
-    
+
     # Preprocess image
-    img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
+    img = cv2.cvtColor(current_image.astype("uint8"), cv2.COLOR_RGBA2GRAY)
     img = 255 - img  # Invert colors
-    
-    # Resize image to match the expected model input shape
-    img_resized = cv2.resize(img, (28, 28))  # Resize to 28x28 for single digit model
-    img_normalized = img_resized / 255.0
-    final_img = img_normalized.reshape(1, 28, 28, 1)  # Adjust to model input shape
-    
-    # Choose the model based on image characteristics (height/width ratio, etc.)
-    if img_resized.shape[1] < 50:  # If the image is narrow, assume it's a single digit
+
+    # === Model Selection Based on Draw Count ===
+    if st.session_state.draw_count == 1:
+        # Assume single-digit drawing
+        img_resized = cv2.resize(img, (28, 28))
+        img_normalized = img_resized / 255.0
+        final_img = img_normalized.reshape(1, 28, 28, 1)
         model_to_use = single_digit_model
+        preds = model_to_use.predict(final_img)
+        predicted_str = str(np.argmax(preds))
     else:
+        # Assume multi-digit drawing
+        img_resized = cv2.resize(img, (100, 28))
+        img_normalized = img_resized / 255.0
+        final_img = img_normalized.reshape(1, 28, 100, 1)
         model_to_use = multi_digit_model
-    
-    # Predict using the selected model
-    preds = model_to_use.predict(final_img)
-    
-    # For multi-digit model, decode and clean prediction
-    if model_to_use == multi_digit_model:
+        preds = model_to_use.predict(final_img)
         predicted_digits = [np.argmax(p[0]) for p in preds]
         predicted_str = ''.join([str(d) for d in predicted_digits])
-    else:
-        # For single digit model, directly decode
-        predicted_str = str(np.argmax(preds))
 
-    # Show prediction result
+    # Show result
     st.markdown(f"### 🧠 Predicted Number: **{predicted_str}**")