Mnist-Digits / app.py
cisemh's picture
Update app.py
c25210a verified
raw
history blame
1.85 kB
import cv2
import gradio as gr
import tensorflow as tf
import numpy as np
# Title and description for the interface
title = "Welcome to your first sketch recognition app!"
head = "<center>The robot was trained to classify numbers (0 to 9). To test it, write your number in the space provided.</center>"
# Image size and label mapping
img_size = 28
labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
# Load the trained model
model = tf.keras.models.load_model("number_recognition_model_colab.keras")
def predict(img):
try:
# Convert the input image to a NumPy array if needed
if not isinstance(img, np.ndarray):
img = np.array(img)
# Convert the image to grayscale if it's not already
if img.ndim == 3 and img.shape[-1] == 3:
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
elif img.ndim == 2:
img = np.expand_dims(img, axis=-1)
# Resize the image to the expected input size
img = cv2.resize(img, (img_size, img_size))
# Normalize the image
img = img.astype('float32') / 255.0
img = img.reshape(1, img_size, img_size, 1)
# Get predictions from the model
preds = model.predict(img)[0]
# Return the predicted probabilities for each class
return {label: float(pred) for label, pred in zip(labels, preds)}
except Exception as e:
return {"Error": str(e)}
# Use a sketchpad as input for drawing
input_component = gr.Sketchpad()
# Output will show the top 3 predicted classes
output_component = gr.Label(num_top_classes=3)
# Create the Gradio interface
interface = gr.Interface(
fn=predict,
inputs=input_component,
outputs=output_component,
title=title,
description=head
)
# Launch the interface
interface.launch(debug=True)