import gluoncv
import mxnet as mx
from gluoncv.utils.viz import get_color_pallete
import gradio as gr


# 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()