File size: 2,574 Bytes
051f92c a49ba8d d4b4b25 cf0b1f5 a49ba8d e6fdb4c a49ba8d e6fdb4c e1b26df e6fdb4c d4b4b25 e1b26df e6fdb4c e1b26df d4b4b25 e6fdb4c e1b26df d4b4b25 e6fdb4c c40b85e cf0b1f5 e6fdb4c 1944562 9cc90a0 1d98dcd 9cc90a0 cf0b1f5 4172ef1 cf0b1f5 363a74f d6bce58 363a74f d6bce58 cf0b1f5 d6bce58 051f92c e6fdb4c c40b85e 051f92c fec34a0 e6fdb4c 051f92c c40b85e cf0b1f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import numpy as np
import gradio as gr
import tensorflow as tf
import cv2
# 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)."
# 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['composite'].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)} """
def predict(data):
# Extract the 'composite' key from the input dictionary
img = data['composite']
# Convert to NumPy array
img = np.array(img)
# Handle RGBA or RGB images
if img.shape[-1] == 4: # RGBA
img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
if img.shape[-1] == 3: # RGB
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
# Resize image to 28x28
img = cv2.resize(img, (28, 28))
# Normalize pixel values to [0, 1]
img = img / 255.0
# Reshape to match model input
img = img.reshape(1, 28, 28, 1)
# Model predictions
preds = model.predict(img)
print(preds)
preds = preds[0]
print(preds)
top_3_classes = np.argsort(preds)[-3:][::-1]
top_3_probs = preds[top_3_classes]
class_names = [labels[i] for i in top_3_classes]
print(class_names, top_3_probs, top_3_classes)
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(share=True) |