| import numpy as np | |
| import cv2 | |
| import gradio as gr | |
| import tensorflow as tf | |
| # app title | |
| title = "Welcome on your first sketch recognition app!" | |
| # app description | |
| head = ( | |
| "<center>" | |
| "<img src='./mnist-classes.png' width=400>" | |
| "The robot was trained to classify numbers (from 0 to 9). To test it, write your number in the space provided." | |
| "</center>" | |
| ) | |
| # GitHub repository link | |
| ref = "Find the whole code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)." | |
| # image size: 28x28 | |
| img_size = 28 | |
| # classes name (from 0 to 9) | |
| labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"] | |
| # load model (trained on MNIST dataset) | |
| model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5") | |
| # prediction function for sketch recognition | |
| def predict(img): | |
| # Convert from PIL to NumPy | |
| img = np.array(img) | |
| # If the image is in RGB format, convert it to grayscale | |
| if len(img.shape) == 3: | |
| img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) | |
| # Resize the image to 28x28 | |
| img = cv2.resize(img, (img_size, img_size)) | |
| # Reshape to the model's input shape (1,28,28,1) | |
| img = img.reshape(1, img_size, img_size, 1) | |
| # model predictions | |
| preds = model.predict(img)[0] | |
| # return the probability for each class | |
| return {label: float(pred) for label, pred in zip(labels, preds)} | |
| # top 3 of classes | |
| label = gr.Label(num_top_classes=3) | |
| # open Gradio interface for sketch recognition | |
| interface = gr.Interface(fn=predict, inputs="sketchpad", outputs=label, title=title, description=head, article=ref) | |
| interface.launch() |