| import numpy as np | |
| import gradio as gr | |
| import tensorflow as tf | |
| # App title | |
| title = "Welcome to your first sketch recognition app!" | |
| # App description | |
| head = ( | |
| "<center>" | |
| "<img src='./mnist-classes.png' width=400>" | |
| "<p>The model is trained to classify numbers (from 0 to 9). " | |
| "To test it, draw your number in the space provided.</p>" | |
| "</center>" | |
| ) | |
| # GitHub repository link | |
| ref = "Find the complete code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)." | |
| # Image size: 28x28 | |
| img_size = 28 | |
| # Class names (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(data): | |
| print(data.shape) | |
| # Reshape image to 28x28 | |
| img = np.reshape(data['composite'], (1, img_size, img_size, 1)) | |
| # Make prediction | |
| pred = model.predict(img) | |
| # Get top class | |
| top_3_classes = np.argsort(pred[0])[-3:][::-1] | |
| # Get top 3 probabilities | |
| top_3_probs = pred[0][top_3_classes] | |
| # Get class names | |
| class_names = [labels[i] for i in top_3_classes] | |
| # Return class names and probabilities | |
| return {class_names[i]: top_3_probs[i] for i in range(3)} | |
| # Top 3 classes | |
| label = gr.Label(num_top_classes=3) | |
| # Open Gradio interface for sketch recognition | |
| interface = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Sketchpad(type='numpy'), | |
| outputs=label, | |
| title=title, | |
| description=head, | |
| article=ref | |
| ) | |
| interface.launch() |