模型简介
模型特点
模型能力
使用案例
🚀 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)
# ------------------------------------------------------------------------
# Stage 1. Detect humans on the image
# ------------------------------------------------------------------------
# You can choose detector by your choice
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] # take first image results
# Human label refers 0 index in COCO dataset
person_boxes = result["boxes"][result["labels"] == 0]
person_boxes = person_boxes.cpu().numpy()
# Convert boxes from VOC (x1, y1, x2, y2) to COCO (x1, y1, w, h) format
person_boxes[:, 2] = person_boxes[:, 2] - person_boxes[:, 0]
person_boxes[:, 3] = person_boxes[:, 3] - person_boxes[:, 1]
# ------------------------------------------------------------------------
# Stage 2. Detect keypoints for each person found
# ------------------------------------------------------------------------
image_processor = AutoProcessor.from_pretrained("usyd-community/vitpose-base-coco-aic-mpii")
model = VitPoseForPoseEstimation.from_pretrained("usyd-community/vitpose-base-coco-aic-mpii", device_map=device)
inputs = image_processor(image, boxes=[person_boxes], return_tensors="pt").to(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] # results for first image
for i, person_pose in enumerate(image_pose_result):
print(f"Person #{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={score.item():.2f}")
输出示例:
Person #0
- Nose: x=428.70, y=170.20, score=0.90
- L_Eye: x=429.26, y=167.11, score=0.94
- R_Eye: x=429.36, y=167.39, score=0.78
- L_Ear: x=432.93, y=167.07, score=0.88
- R_Ear: x=441.39, y=166.26, score=0.86
- L_Shoulder: x=439.87, y=176.99, score=0.94
- R_Shoulder: x=444.96, y=177.49, score=0.70
- L_Elbow: x=436.33, y=196.93, score=0.98
- R_Elbow: x=431.81, y=200.50, score=0.84
- L_Wrist: x=430.75, y=217.52, score=0.92
- R_Wrist: x=421.83, y=212.19, score=0.86
- L_Hip: x=444.97, y=223.51, score=0.79
- R_Hip: x=452.21, y=222.88, score=0.70
- L_Knee: x=442.63, y=255.64, score=0.78
- R_Knee: x=452.44, y=255.15, score=0.83
- L_Ankle: x=444.95, y=288.12, score=0.63
- R_Ankle: x=456.43, y=286.81, score=0.86
Person #1
- Nose: x=398.27, y=181.73, score=0.84
- L_Eye: x=398.43, y=179.77, score=0.85
- R_Eye: x=396.03, y=179.55, score=0.85
- R_Ear: x=389.00, y=180.26, score=0.84
- L_Shoulder: x=397.21, y=194.18, score=0.73
- R_Shoulder: x=384.42, y=190.45, score=0.56
✨ 主要特性
- 结构简单:ViTPose采用简单且非分层的视觉Transformer作为骨干网络,用于提取给定人物实例的特征,并使用轻量级解码器进行姿态估计。
- 可扩展性:利用Transformer可扩展的模型容量和高并行性,模型参数可以从1亿扩展到10亿,在吞吐量和性能之间达到了新的帕累托最优。
- 灵活性:在注意力类型、输入分辨率、预训练和微调策略以及处理多个姿态任务方面具有很高的灵活性。
- 知识可迁移性:通过简单的知识令牌,可以轻松地将大型ViTPose模型的知识迁移到小型模型中。
📚 详细文档
模型详情
尽管在设计中没有考虑特定的领域知识,但普通的视觉Transformer在视觉识别任务中表现出了出色的性能。然而,很少有人努力揭示这种简单结构在姿态估计任务中的潜力。在本文中,作者通过一个名为ViTPose的简单基线模型,从模型结构的简单性、模型大小的可扩展性、训练范式的灵活性以及模型之间的知识可迁移性等方面,展示了普通视觉Transformer在姿态估计方面令人惊讶的良好能力。
模型描述
- 开发者: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
应用场景
- 人体姿态估计:该模型可用于估计图像或视频中人体的姿态,识别头部、肩部、肘部、手腕、臀部、膝盖和脚踝等关键身体关节的位置。
- 动作识别:通过分析一段时间内的姿态,该模型有助于识别各种人类动作和活动。
- 监控:在安全和监控应用中,ViTPose可用于监控和分析公共场所或私人场所中的人类行为。
- 健康与健身:该模型可用于健身应用程序,跟踪和分析运动姿态,提供关于姿势和技术的反馈。
- 游戏和动画:ViTPose可集成到游戏和动画系统中,创建更逼真的角色动作和交互。
偏差、风险和局限性
尽管ViTPose在MS COCO数据集上取得了最优性能,但该模型的潜力尚未通过更先进的技术(如复杂的解码器或FPN结构)得到充分挖掘,这些技术可能会进一步提高性能。此外,虽然ViTPose展示了简单性、可扩展性、灵活性和可迁移性等令人兴奋的特性,但还需要更多的研究工作,例如探索基于提示的调优,以进一步展示ViTPose的灵活性。作者认为ViTPose也可以应用于其他姿态估计数据集,如动物姿态估计和面部关键点检测。
训练详情
训练数据
使用MS COCO、AI Challenger、MPII和CrowdPose数据集进行训练和评估。OCHuman数据集仅用于评估阶段,以衡量模型处理遮挡人物的性能。
- 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许可证。
🔧 技术细节
在姿态估计任务中,ViTPose通过简单的结构和灵活的设计,充分发挥了Transformer的优势。其使用普通的视觉Transformer作为骨干网络,能够有效地提取特征,并且在不同的训练和应用场景中表现出了良好的适应性。通过实验验证,该模型在多个数据集上取得了优异的成绩,证明了其在人体姿态估计领域的有效性和潜力。
📚 引用
@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}
}









