Vitpose Plus Base

🚀 VitPose模型卡

ViTPose是用于人体姿态估计的简单视觉Transformer基线模型，ViTPose+则是用于通用人体姿态估计的视觉Transformer基础模型。该模型在MS COCO关键点测试开发集上达到了81.1的平均精度（AP）。

🚀 快速开始

使用以下代码开始使用该模型：

import torch
import requests
import numpy as np

from PIL import Image

from transformers import (
    AutoProcessor,
    RTDetrForObjectDetection,
    VitPoseForPoseEstimation,
)

device = "cuda" if torch.cuda.is_available() else "cpu"

url = "http://images.cocodataset.org/val2017/000000000139.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# ------------------------------------------------------------------------
# 阶段1. 在图像上检测人体
# ------------------------------------------------------------------------

# 你可以自行选择检测器
person_image_processor = AutoProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
person_model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365", device_map=device)

inputs = person_image_processor(images=image, return_tensors="pt").to(device)

with torch.no_grad():
    outputs = person_model(**inputs)

results = person_image_processor.post_process_object_detection(
    outputs, target_sizes=torch.tensor([(image.height, image.width)]), threshold=0.3
)
result = results[0]  # 获取第一张图像的结果

# 在COCO数据集中，人体标签对应索引0
person_boxes = result["boxes"][result["labels"] == 0]
person_boxes = person_boxes.cpu().numpy()

# 将边界框从VOC (x1, y1, x2, y2) 格式转换为COCO (x1, y1, w, h) 格式
person_boxes[:, 2] = person_boxes[:, 2] - person_boxes[:, 0]
person_boxes[:, 3] = person_boxes[:, 3] - person_boxes[:, 1]

# ------------------------------------------------------------------------
# 阶段2. 为检测到的每个人检测关键点
# ------------------------------------------------------------------------

image_processor = AutoProcessor.from_pretrained("usyd-community/vitpose-plus-base")
model = VitPoseForPoseEstimation.from_pretrained("usyd-community/vitpose-plus-base", device_map=device)

inputs = image_processor(image, boxes=[person_boxes], return_tensors="pt").to(device)

# 这是MOE架构，我们需要为每张图像指定数据集索引，范围为0..5
inputs["dataset_index"] = torch.tensor([0], device=device)

with torch.no_grad():
    outputs = model(**inputs)

pose_results = image_processor.post_process_pose_estimation(outputs, boxes=[person_boxes], threshold=0.3)
image_pose_result = pose_results[0]  # 第一张图像的结果

for i, person_pose in enumerate(image_pose_result):
    print(f"人物 #{i}")
    for keypoint, label, score in zip(
        person_pose["keypoints"], person_pose["labels"], person_pose["scores"]
    ):
        keypoint_name = model.config.id2label[label.item()]
        x, y = keypoint
        print(f" - {keypoint_name}: x={x.item():.2f}, y={y.item():.2f}, 分数={score.item():.2f}")

输出结果：

人物 #0
 - 鼻子: x=428.81, y=171.53, 分数=0.92
 - 左眼: x=429.32, y=168.30, 分数=0.92
 - 右眼: x=428.84, y=168.47, 分数=0.82
 - 左耳: x=434.60, y=166.54, 分数=0.90
 - 右耳: x=440.14, y=165.80, 分数=0.80
 - 左肩膀: x=440.74, y=176.95, 分数=0.96
 - 右肩膀: x=444.06, y=177.52, 分数=0.68
 - 左肘: x=436.30, y=197.08, 分数=0.91
 - 右肘: x=432.29, y=201.22, 分数=0.79
 - 左手腕: x=429.91, y=217.90, 分数=0.84
 - 右手腕: x=421.08, y=212.72, 分数=0.90
 - 左髋: x=446.15, y=223.88, 分数=0.74
 - 右髋: x=449.32, y=223.45, 分数=0.65
 - 左膝: x=443.73, y=255.72, 分数=0.76
 - 右膝: x=450.72, y=255.21, 分数=0.73
 - 左踝: x=452.14, y=287.30, 分数=0.66
 - 右踝: x=456.02, y=285.99, 分数=0.72
人物 #1
 - 鼻子: x=398.22, y=181.60, 分数=0.88
 - 左眼: x=398.67, y=179.84, 分数=0.87
 - 右眼: x=396.07, y=179.44, 分数=0.87
 - 右耳: x=388.94, y=180.38, 分数=0.87
 - 左肩膀: x=397.11, y=194.19, 分数=0.71
 - 右肩膀: x=384.75, y=190.74, 分数=0.55

✨ 主要特性

尽管在设计中未考虑特定领域知识，但普通视觉Transformer在视觉识别任务中表现出了出色的性能。然而，在揭示这种简单结构在姿态估计任务中的潜力方面，所做的工作还很少。在本文中，通过一个名为ViTPose的简单基线模型，从模型结构的简单性、模型大小的可扩展性、训练范式的灵活性以及模型间知识的可迁移性等方面，展示了普通视觉Transformer在姿态估计方面令人惊讶的良好能力。具体来说，ViTPose采用普通的非分层视觉Transformer作为骨干网络，为给定的人物实例提取特征，并使用轻量级解码器进行姿态估计。通过利用Transformer可扩展的模型容量和高并行性，其参数可以从1亿扩展到10亿，在吞吐量和性能之间建立了新的帕累托最优前沿。此外，ViTPose在注意力类型、输入分辨率、预训练和微调策略以及处理多个姿态任务方面都非常灵活。还通过实验证明，大型ViTPose模型的知识可以通过简单的知识令牌轻松转移到小型模型。实验结果表明，基本的ViTPose模型在具有挑战性的MS COCO关键点检测基准上优于代表性方法，而最大的模型则创造了新的技术水平，即在MS COCO测试开发集上达到了80.9的AP。代码和模型可在此处获取。

📚 详细文档

模型详情

这是一个已发布在Hugging Face Hub上的🤗 transformers模型的模型卡，该模型卡是自动生成的。

• 开发者：Yufei Xu, Jing Zhang, Qiming Zhang, Dacheng Tao
• 资助方：ARC FL - 170100117和IH - 180100002
• 许可证：Apache - 2.0
• 移植到🤗 Transformers的人员：Sangbum Choi和Niels Rogge

模型来源

• 原始仓库：https://github.com/ViTAE-Transformer/ViTPose
• 论文：https://arxiv.org/pdf/2204.12484
• 演示：https://huggingface.co/spaces?sort=trending&search=vitpose

模型用途

由ViTAE - Transformer团队开发的ViTPose模型主要用于姿态估计任务。以下是该模型的一些直接用途：

• 人体姿态估计：该模型可用于估计图像或视频中人体的姿态，包括识别头部、肩膀、肘部、手腕、臀部、膝盖和脚踝等关键身体关节的位置。
• 动作识别：通过分析一段时间内的姿态，该模型有助于识别各种人类动作和活动。
• 监控：在安全和监控应用中，ViTPose可用于监控和分析公共场所或私人场所中的人类行为。
• 健康与健身：该模型可用于健身应用程序，以跟踪和分析运动姿态，提供关于姿势和技术的反馈。
• 游戏和动画：ViTPose可集成到游戏和动画系统中，以创建更逼真的角色动作和交互。

偏差、风险和局限性

在本文中，提出了一个简单而有效的用于姿态估计的视觉Transformer基线模型，即ViTPose。尽管在结构上没有精心设计，但ViTPose在MS COCO数据集上取得了最先进的性能。然而，ViTPose的潜力尚未通过更先进的技术（如复杂的解码器或FPN结构）得到充分挖掘，这些技术可能会进一步提高其性能。此外，尽管ViTPose展示了诸如简单性、可扩展性、灵活性和可迁移性等令人兴奋的特性，但还可以进行更多的研究工作，例如探索基于提示的调优以进一步展示ViTPose的灵活性。此外，相信ViTPose也可以应用于其他姿态估计数据集，例如动物姿态估计 [47, 9, 45] 和面部关键点检测 [21, 6]，这些将作为未来的工作。

训练详情

训练数据

使用MS COCO [28]、AI Challenger [41]、MPII [3] 和CrowdPose [22] 数据集进行训练和评估。OCHuman [54] 数据集仅用于评估阶段，以衡量模型处理遮挡人物的性能。MS COCO数据集包含118K张图像和150K个人体实例，每个实例最多有17个关键点标注，用于训练，该数据集遵循CC - BY - 4.0许可证。MPII数据集遵循BSD许可证，包含15K张图像和22K个人体实例用于训练，该数据集中每个实例最多标注16个人体关键点。AI Challenger数据集更大，包含超过200K张训练图像和350个人体实例，每个实例最多标注14个关键点。OCHuman数据集包含大量遮挡的人体实例，仅用于验证和测试集，包括4K张图像和8K个实例。

训练超参数

• 训练机制：

速度、大小、时间

评估

OCHuman验证集和测试集

为了评估人体姿态估计模型在严重遮挡人体实例上的性能，在OCHuman验证集和测试集上使用真实边界框对ViTPose变体和代表性模型进行了测试。由于OCHuman数据集中并非所有人体实例都有标注，使用额外的人体检测器会导致大量“误报”边界框，无法反映姿态估计模型的真实能力，因此未采用额外的人体检测器。具体来说，使用了对应于MS COCO数据集的ViTPose解码器头，因为MS COCO和OCHuman数据集中的关键点定义相同。

MPII验证集

在MPII验证集上使用真实边界框评估了ViTPose和代表性模型的性能。遵循MPII的默认设置，使用PCKh作为性能评估指标。

评估结果

模型架构和目标

硬件

模型基于mmpose代码库在8个A100 GPU上进行训练。

🔧 技术细节

信息表格

📄 许可证

该模型遵循Apache - 2.0许可证。

🔖 引用

@article{xu2022vitposesimplevisiontransformer,
  title={ViTPose: Simple Vision Transformer Baselines for Human Pose Estimation},
  author={Yufei Xu and Jing Zhang and Qiming Zhang and Dacheng Tao},
  year={2022},
  eprint={2204.12484},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2204.12484}
}