Spaces:
Sleeping
Sleeping
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) | |