Harshad Bhandwaldar commited on
Commit
2e4b1e8
1 Parent(s): 3fa6693
Files changed (2) hide show
  1. app.py +9 -1
  2. requirements.txt +2 -1
app.py CHANGED
@@ -2,7 +2,9 @@ import gluoncv
2
  import mxnet as mx
3
  from gluoncv.utils.viz import get_color_pallete
4
  import gradio as gr
5
-
 
 
6
 
7
  # using cpu
8
  ctx = mx.cpu(0)
@@ -13,6 +15,12 @@ FILE_NAME = "result.png"
13
  model = gluoncv.model_zoo.get_model("psp_resnet101_ade", pretrained=True)
14
 
15
  def segmentation(img):
 
 
 
 
 
 
16
  output = model.predict(img)
17
  predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
18
  mask = get_color_pallete(predict, "ade20k")
 
2
  import mxnet as mx
3
  from gluoncv.utils.viz import get_color_pallete
4
  import gradio as gr
5
+ import numpy as np
6
+ import PIL as Image
7
+ from gluoncv.data.transforms.presets.segmentation import test_transform
8
 
9
  # using cpu
10
  ctx = mx.cpu(0)
 
15
  model = gluoncv.model_zoo.get_model("psp_resnet101_ade", pretrained=True)
16
 
17
  def segmentation(img):
18
+ img = Image.open(img)
19
+ img.load()
20
+ img = np.asarray(img, dtype="int32")
21
+ img = mx.ndarray.array(img)
22
+ ctx = mx.cpu(0)
23
+ img = test_transform(img, ctx)
24
  output = model.predict(img)
25
  predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
26
  mask = get_color_pallete(predict, "ade20k")
requirements.txt CHANGED
@@ -3,4 +3,5 @@ gluoncv==0.10.5
3
  matplotlib>=3.5.3
4
  opencv-python>=4.6.0
5
  gradio
6
- numpy==1.23.1
 
 
3
  matplotlib>=3.5.3
4
  opencv-python>=4.6.0
5
  gradio
6
+ numpy==1.23.1
7
+ pillow