Fixes

- fixes the example
- the number of classes is 20, not 21

README.md

@@ -21,15 +21,16 @@ I translated the model to ONNX format.
 
 Get the `deeplabv3p-resnet50-human.onnx` file and use it with ONNXRuntime package.
 
-The result of `model.run` is a `(1, 1, 512, 512, 21)` tensor:
+The result of `model.run` is a `(1, 1, 512, 512, 20)` tensor:
 
 - 1: number of output (you can squeeze it)
 - 1: batch size (you can squeeze it)
 - 512, 512: the size of the image (fixed)
-- 21: number of classes, so you can take the `argmax`` of the tensor to get the class of each pixel
+- 20: number of classes, so you can take the `argmax`` of the tensor to get the class of each pixel
 
 ```python
 import onnxruntime
+import numpy as np
 from PIL import Image
 
 img = Image.open(sys.argv[1] if len(sys.argv) > 1 else "image.jpg")
@@ -44,7 +45,7 @@ result = np.array(result[0])
 result = result.argmax(axis=3).squeeze(0)
 
 # get the masks
-for i in range(21):
+for i in range(20):
     detected = result == i # get the detected pixels for the class i
     # detected  is a 512, 512 boolean array
     mask = np.zeros_like(img)