app.py

import gluoncv
import mxnet as mx
from gluoncv.utils.viz import get_color_pallete
import gradio as gr
import numpy as np
import PIL as Image
from gluoncv.data.transforms.presets.segmentation import test_transform

# using cpu
ctx = mx.cpu(0)


FILE_NAME = "result.png"

model = gluoncv.model_zoo.get_model("psp_resnet101_ade", pretrained=True)

def segmentation(img):
    img = Image.open(img)
    img.load()
    img = np.asarray(img, dtype="int32")
    img = mx.ndarray.array(img)
    ctx = mx.cpu(0)
    img = test_transform(img, ctx)
    output = model.predict(img)
    predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
    mask = get_color_pallete(predict, "ade20k")
    # mask.save("result.png")
    # mmask = mpimg.imread("result.png")
    # plt.imshow(mmask)
    # plt.savefig("result.png")
    return mask

image_in = gr.Image()
image_out = gr.components.Image()

Iface = gr.Interface(
    fn=segmentation,
    inputs=image_in,
    outputs=image_out,
    title="Sementic Segmentation - MXNet",
).launch()