import gradio as gr
import numpy as np
import torch
import torchvision

def on_submit(img):
    img = img['composite'].astype(np.float32)
    img = torch.from_numpy(img)
    return f"Sum: {img.shape}"

iface = gr.Interface(
    title = "LeNet",
    fn = on_submit,
    inputs=gr.Sketchpad(image_mode='P'),
    outputs=gr.Label(),
)
iface.launch()