app.py

import gluoncv
import matplotlib.image as mpimg
import mxnet as mx
from gluoncv.data.transforms.presets.segmentation import test_transform
from gluoncv.utils.viz import get_color_pallete
from matplotlib import pyplot as plt
from mxnet import image
import gradio as gr
import numpy as np

# using cpu
ctx = mx.cpu(0)


FILE_NAME = "result.png"

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

def segmentation(img):
    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()