app.py

import os
import numpy as np
import cv2
import gradio as gr
import tensorflow as tf
from PIL import Image

# app title
title = "Welcome on your first sketch recognition app!"

# app description
head = (
  "<center>"
  "<img src='./mnist-classes.png' width=400>"
  "<p>The robot was trained to classify numbers (0 to 9). "
  "To test it, write your number in the space provided!</p>"
  "</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
img_size = 28

# Classes
labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]

# Load model
model_path = "./sketch_recognition_numbers_model.h5"
try:
    model = tf.keras.models.load_model(model_path)
except Exception as e:
    raise FileNotFoundError(f"Model file '{model_path}' not found or failed to load. {str(e)}")

def predict(img):
    # If no image is provided, return an error message
    if img is None:
        return {"error": "No image provided."}

    # Ensure the image is a PIL Image
    if not isinstance(img, Image.Image):
        img = Image.fromarray(np.uint8(img))

    # Convert to grayscale
    img = img.convert("L")
    
    # Convert PIL Image to a NumPy array of type uint8
    img = np.array(img, dtype=np.uint8)

    # Resize to (28x28)
    img = cv2.resize(img, (img_size, img_size))

    # Reshape to match model input shape (1, 28, 28, 1)
    img = img.reshape(1, img_size, img_size, 1)

    # Model predictions
    preds = model.predict(img)[0]

    # Return probabilities for each class
    return {label: float(pred) for label, pred in zip(labels, preds)}

# Use gr.Sketchpad to ensure a PIL image is returned
interface = gr.Interface(
    fn=predict,
    inputs=gr.Sketchpad(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title=title,
    description=head,
    article=ref
)

interface.launch()