app.py

import gradio as gr
from PIL import Image
import torchvision
import torch

# load model

def predict(input_image):
    pil_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
    # transform image to torch and do preprocessing
    torch_image = torchvision.transforms.ToTensor()(pil_image)
    # model predict
    prediction = torch.rand(torch_image.shape)
    # transform torch to image
    predicted_pil_image = torchvision.transforms.ToPILImage()(prediction)
    # return correct image
    return predicted_pil_image

iface = gr.Interface(
    fn=predict, 
    inputs=gr.Image(shape=(512,512)), 
    outputs=gr.Image(shape=(512,512)),
    examples=[
        ["demo_imgs/fake.jpg"] # use real image
    ],
    title="DTM Estimation",
    description="This demo predict a DTM..."
)
iface.launch()