Spaces:
Runtime error
Runtime error
File size: 1,027 Bytes
4c9c630 2e4b1e8 871fb7e 2e4b1e8 4c9c630 2e4b1e8 d498fcd 2e4b1e8 4c9c630 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import gluoncv
import mxnet as mx
from gluoncv.utils.viz import get_color_pallete
import gradio as gr
import numpy as np
from PIL import 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() |