--- license: cc0-1.0 tags: - art --- # DeepLabV3+ ResNet50 for human body parts segmentation This is a very simple ONNX model that can segment human body parts. ## Why this model This model is a ONNX transposition of [keras-io/deeplabv3p-resnet50](https://huggingface.co/keras-io/deeplabv3p-resnet50) where the provided model can segment human body parts. All the others models that I found was trained on city segmentation. The original model is built for old version of Keras and cannot be used with recent version of TensorFlow. I translated the model to ONNX format. ## Usage Get the `deeplabv3p-resnet50-human.onnx` file and use it with ONNXRuntime package. 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) - 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 model = onnxruntime.InferenceSession("deeplabv3p-resnet50-human.onnx") img = Image.open(sys.argv[1] if len(sys.argv) > 1 else "image.jpg") img = img.resize((512, 512)) img = np.array(img).astype(np.float32) / 127.5 - 1 # infer input_name = model.get_inputs()[0].name output_name = model.get_outputs()[0].name result = model.run([output_name], {input_name: img}) # squeeze, argmax... result = np.array(result[0]) # argmax the classes, remove the batch size result = result.argmax(axis=3).squeeze(0) # get the masks 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) mask[detected] = 255 Image.fromarray(mask).show() # or save, or return the mask... ``` ## Classes index This is the list of classes that the model can detect (some classes are not specifically identified, see below): - 0: "background", - 1: "unknown", - 2: "hair", - 3: "unknown", - 4: "glasses", - 5: "top-clothes", - 6: "unknown", - 7: "unknown", - 8: "unknown", - 9: "bottom-clothes", - 10: "torso-skin", - 11: "unknown", - 12: "unknown", - 13: "face", - 14: "left-arm", - 15: "right-arm", - 16: "left-leg", - 17: "right-leg", - 18: "left-foot", - 19: "right-foot", ## Known limitation - The model could fail on portrait images, because the model was trained on "full body" images. - There are some classes that I don't know what they are. I can't find the list of classes (help !). - The model is not perfect, and can fail on some images. I'm not the author of the model, so I can't fix it. ## License The [original model card](https://huggingface.co/keras-io/deeplabv3p-resnet50/blob/main/README.md) proposes the "CC0-1.0" license. I don't know if it's the right license for the model, but I keep it. > Anyway, thanks to the authors of the model for sharing it and to leave it open to use. This means that you may use the model, share, modify, and distribute it without any restriction.