DeepLabV3+ ResNet50 人体部位分割模型

这是一个非常简单的 ONNX 模型，可以对人体部位进行分割。

为什么选择这个模型

这个模型是 keras-io/deeplabv3p-resnet50 的 ONNX 转换版本， 该模型可以对人体部位进行分割。我找到的其他所有模型都是在城市分割数据集上训练的。

原始模型是为旧版本的 Keras 构建的，无法与最新版本的 TensorFlow 一起使用。 我将模型转换为了 ONNX 格式。

使用方法

获取 deeplabv3p-resnet50-human.onnx 文件并配合 ONNXRuntime 包使用。

model.run 的结果是一个 (1, 1, 512, 512, 20) 张量：

• 1：输出数量（可以压缩）
• 1：批次大小（可以压缩）
• 512, 512：图像尺寸（固定）
• 20：类别数量，您可以对张量取 argmax 来获取每个像素的类别

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

# 推理
input_name = model.get_inputs()[0].name
output_name = model.get_outputs()[0].name
result = model.run([output_name], {input_name: img})

# 压缩，取 argmax...
result = np.array(result[0])
# 对类别取 argmax，移除批次大小
result = result.argmax(axis=3).squeeze(0)

# 获取掩码
for i in range(20):
    detected = result == i # 获取类别 i 的检测像素
    # detected 是一个 512, 512 的布尔数组
    mask = np.zeros_like(img)
    mask[detected] = 255
    Image.fromarray(mask).show() # 或保存，或返回掩码...

类别索引

这是模型可以检测的类别列表（有些类别没有明确标识，见下文）：

• 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"（右脚）,

已知限制

• 该模型在肖像图像上可能会失败，因为模型是在"全身"图像上训练的。
• 有些类别我不知道它们代表什么。我找不到类别列表（求助！）。
• 该模型并不完美，在某些图像上可能会失败。我不是模型的作者，所以我无法修复它。

许可证

原始模型卡 采用 "CC0-1.0" 许可证。我不知道这是否是该模型的正确许可证，但我保留了这个许可证。

无论如何，感谢模型的作者分享它并保持开放使用。

这意味着您可以不受任何限制地使用、分享、修改和分发该模型。