Views
No views yet
deeplabv3p-resnet50-human.onnx file and use it with ONNXRuntime package.model.run is a (1, 1, 512, 512, 20) tensor:1import onnxruntime
2import numpy as np
3from PIL import Image
4
5model = onnxruntime.InferenceSession("deeplabv3p-resnet50-human.onnx")
6
7img = Image.open(sys.argv[1] if len(sys.argv) > 1 else "image.jpg")
8img = img.resize((512, 512))
9img = np.array(img).astype(np.float32) / 127.5 - 1
10
11# infer
12input_name = model.get_inputs()[0].name
13output_name = model.get_outputs()[0].name
14result = model.run([output_name], {input_name: img})
15
16# squeeze, argmax...
17result = np.array(result[0])
18# argmax the classes, remove the batch size
19result = result.argmax(axis=3).squeeze(0)
20
21# get the masks
22for i in range(20):
23 detected = result == i # get the detected pixels for the class i
24 # detected is a 512, 512 boolean array
25 mask = np.zeros_like(img)
26 mask[detected] = 255
27 Image.fromarray(mask).show() # or save, or return the mask...Anyway, thanks to the authors of the model for sharing it and to leave it open to use.