Update app.py

app.py

@@ -1,102 +1,105 @@
-import cv2
 import gradio as gr
-import tensorflow as tf
 import numpy as np
+import tensorflow as tf
 from PIL import Image
 
-title = "Welcome on your first sketch recognition app!"
-
-head = (
-    "<center>"
-    "The robot was trained to classify numbers (from 0 to 9). To test it, write your number in the space provided."
-    "</center>"
-)
-
-# Model yükleniyor
-model = tf.keras.models.load_model("number_recognition_model_colab.keras")
+# Load the pre-trained model
+model = tf.keras.models.load_model('number_recognition_model_colab.h5')
 
+# Class names for MNIST digits
+classes_names = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
 
-img_size = 28
+def preprocess_image(image):
+    """
+    Preprocess the input image to match the model's expected input
+    - Resize to 28x28
+    - Convert to grayscale
+    - Normalize pixel values
+    """
+    # Convert to grayscale if image is RGB
+    if len(image.shape) == 3:
+        image = np.mean(image, axis=2)
+    
+    # Resize to 28x28
+    image = Image.fromarray(image.astype('uint8')).resize((28, 28), Image.LANCZOS)
+    image = np.array(image)
+    
+    # Normalize
+    image = image.astype("float32") / 255.0
+    
+    # Reshape to match model input shape
+    image = image.reshape(1, 28, 28, 1)
+    
+    return image
 
-labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
+def predict_digit(image):
+    """
+    Predict the digit in the input image
+    """
+    # Preprocess the image
+    processed_image = preprocess_image(image)
+    
+    # Make prediction
+    predictions = model.predict(processed_image)
+    
+    # Get the predicted class
+    predicted_class = np.argmax(predictions[0])
+    predicted_label = classes_names[predicted_class]
+    confidence = predictions[0][predicted_class] * 100
+    
+    # Create a more detailed output
+    result_text = f"Predicted Digit: {predicted_label}\nConfidence: {confidence:.2f}%"
+    
+    # Create bar chart of probabilities
+    probabilities = predictions[0] * 100
+    
+    return result_text, probabilities
 
-def predict(img):
-    try:
-        # Enhanced image validation and conversion
-        if img is None:
-            raise ValueError("No image provided")
-        
-        # Convert to numpy array if it's a PIL Image
-        if isinstance(img, Image.Image):
-            img = np.array(img)
-        
-        # Handle base64 image strings
-        elif isinstance(img, str):
-            # Check if it's a base64 data URL
-            if img.startswith('data:image'):
-                # Split and decode base64 part
-                img = img.split(',')[1]
-            
-            # Decode base64 to image
-            try:
-                img = Image.open(io.BytesIO(base64.b64decode(img)))
-                img = np.array(img)
-            except Exception as e:
-                print(f"Base64 decoding error: {e}")
-                raise ValueError("Invalid base64 image")
-        
-        # Validate numpy array
-        if not isinstance(img, np.ndarray):
-            raise ValueError("Input could not be converted to a valid image")
-        
-        # Print initial image details for debugging
-        print(f"Initial image type: {type(img)}, shape: {img.shape}")
-        
-        # Handle color channels
-        if img.ndim == 3:
-            if img.shape[-1] == 3:  # Color image
-                img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-            elif img.shape[-1] == 4:  # RGBA image
-                img = cv2.cvtColor(img, cv2.COLOR_RGBA2GRAY)
-        
-        # Ensure single channel
-        if img.ndim == 2:
-            img = np.expand_dims(img, axis=-1)
-        
-        # Resize and normalize
-        img = cv2.resize(img, (img_size, img_size))
-        img = img.astype('float32') / 255.0
-        img = img.reshape(1, img_size, img_size, 1)
-        
-        # Print processed image details
-        print(f"Processed image shape: {img.shape}")
-        
-        # Get predictions from the model
-        preds = model.predict(img)[0]
-        
-        # Print predictions for debugging
-        print("Predictions:", preds)
-        
-        # Return predictions as a dictionary
-        return {label: float(pred) for label, pred in zip(labels, preds)}
+def create_probability_plot(probabilities):
+    """
+    Create a bar plot of digit probabilities
+    """
+    import matplotlib.pyplot as plt
+    
+    plt.figure(figsize=(10, 5))
+    plt.bar(classes_names, probabilities)
+    plt.title('Digit Probability Distribution')
+    plt.xlabel('Digits')
+    plt.ylabel('Probability (%)')
+    plt.ylim(0, 100)
+    
+    # Rotate x-axis labels
+    plt.xticks(rotation=45)
     
-    except Exception as e:
-        # Comprehensive error logging
-        print(f"Full error during prediction: {e}")
-        return {"Error": str(e)}
-            
-   
+    return plt
 
-# Set up the Gradio interface with the input as a sketchpad and output as labels
-label = gr.Label(num_top_classes=3)
+# Create Gradio interface
+def gradio_predict(image):
+    """
+    Wrapper function for Gradio interface
+    """
+    result_text, probabilities = predict_digit(image)
+    prob_plot = create_probability_plot(probabilities)
+    return result_text, prob_plot
 
-# Gradio arayüzü
-interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Sketchpad(type="pil"),
-    outputs=label,
-    title="Sketch Recognition App",
-    description="Draw a number (0-9) and see the model's top predictions."
+# Set up the Gradio interface
+iface = gr.Interface(
+    fn=gradio_predict,
+    inputs=gr.Image(type="numpy", image_mode="L"),
+    outputs=[
+        gr.Textbox(label="Prediction"),
+        gr.Plot(label="Probability Distribution")
+    ],
+    title="MNIST Digit Recognizer",
+    description="Draw a single-digit number (0-9) and the model will predict which digit it is!",
+    allow_flagging="never",
+    examples=[
+        ["example_zero.png"],
+        ["example_one.png"],
+        ["example_two.png"]
+    ]
 )
 
-interface.launch(debug=True, share=True)
+# Launch the interface
+if __name__ == "__main__":
+    iface.launch()