用于无约束头部姿态估计的6D旋转表示法(Pytorch)
引用
如果您觉得我们的工作有用,请引用以下论文:
@ARTICLE{10477888,
author={Hempel, Thorsten and Abdelrahman, Ahmed A. and Al-Hamadi, Ayoub},
journal={IEEE Transactions on Image Processing},
title={Toward Robust and Unconstrained Full Range of Rotation Head Pose Estimation},
year={2024},
volume={33},
number={},
pages={2377-2387},
keywords={Head;Training;Predictive models;Pose estimation;Quaternions;Three-dimensional displays;Training data;Head pose estimation;full range of rotation;rotation matrix;6D representation;geodesic loss},
doi={10.1109/TIP.2024.3378180}}
@INPROCEEDINGS{9897219,
author={Hempel, Thorsten and Abdelrahman, Ahmed A. and Al-Hamadi, Ayoub},
booktitle={2022 IEEE International Conference on Image Processing (ICIP)},
title={6d Rotation Representation For Unconstrained Head Pose Estimation},
year={2022},
volume={},
number={},
pages={2496-2500},
doi={10.1109/ICIP46576.2022.9897219}}
更新
2023年9月18日
- 我们推出了6DRepNet360!查看我们处理全范围头部姿态方向预测的新版本:https://github.com/thohemp/6DRepNet360
2022年9月13日
- 6DRepNet现已作为pip包提供,使用更加便捷:pip3 install SixDRepNet
2022年6月20日
- 6DRepNet已被ICIP 2022接收。
2022年5月29日
- 简化了训练脚本
- 更新了默认训练配置,以获得更稳健的结果
论文
Thorsten Hempel、Ahmed A. Abdelrahman和Ayoub Al-Hamadi,"用于无约束头部姿态估计的6D旋转表示法",已被ICIP 2022接收。[ResearchGate][Arxiv]
摘要
在本文中,我们提出了一种无约束的端到端头部姿态估计方法。我们通过引入旋转矩阵形式来处理模糊旋转标签的问题,并提出了一种连续的6D旋转矩阵表示法,用于高效和稳健的直接回归。这样,我们的方法可以学习完整的旋转外观,这与之前为了获得令人满意的结果而将姿态预测限制在窄角度范围内的方法相反。此外,我们提出了一种基于测地距离的损失函数,以根据流形几何来惩罚我们的网络。在公开的AFLW2000和BIWI数据集上的实验表明,我们提出的方法显著优于其他最先进的方法,性能提升幅度高达20%。
## **在300W-LP上训练,在AFLW2000和BIWI上测试**
| | | || | | | | | | |
| --------------------- | -------------- |--- | ------- | ------ | ------ | ------ | ------ | ------ | ------ | ------ |
| | **全范围** | **偏航角** | **俯仰角** | **横滚角** | **平均绝对误差** | | **偏航角** | **俯仰角** | **横滚角** | **平均绝对误差** |
| HopeNet ( =2) | 否 | 6.47 | 6.56 | 5.44 | 6.16 || 5.17 | 6.98 | 3.39 | 5.18 |
| HopeNet ( =1)| 否 | 6.92 | 6.64 | 5.67 | 6.41 || 4.81 | 6.61 | 3.27 | 4.90 |
| FSA-Net | 否 | 4.50 | 6.08 | 4.64 | 5.07 || 4.27 | 4.96 | 2.76 | 4.00 |
| HPE | 否 | 4.80 | 6.18 | 4.87 | 5.28 || 3.12 | 5.18 | 4.57 | 4.29 |
| QuatNet | 否 | 3.97 | 5.62 | 3.92 | 4.50 || **2.94** | 5.49 | 4.01 | 4.15 |
| WHENet-V | 否 | 4.44 | 5.75 | 4.31 | 4.83 || 3.60 | **4.10** | 2.73 | 3.48 |
| WHENet | 是/否 | 5.11 | 6.24 | 4.92 | 5.42 || 3.99 | 4.39 | 3.06 | 3.81 |
| TriNet | 是 | 4.04 | 5.77 | 4.20 | 4.67 || 4.11 | 4.76 | 3.05 | 3.97 |
| FDN | 否 | 3.78 | 5.61 | 3.88 | 4.42 | | 4.52 | 4.70 | **2.56** | 3.93 |
| | | | | | | | | | |
| **6DRepNet** | 是 | **3.63** | **4.91** | **3.37** | **3.97** || 3.24 | 4.48 | 2.68 | **3.47** |
| | | | | | | | | | | |
BIWI 70/30
偏航 | 俯仰 | 翻滚 | 平均绝对误差 | |
HopeNet ( =1) | 3.29 | 3.39 | 3.00 | 3.23 |
FSA-Net | 2.89 | 4.29 | 3.60 | 3.60 |
TriNet | 2.93 | 3.04 | 2.44 | 2.80 |
FDN | 3.00 | 3.98 | 2.88 | 3.29 |
6DRepNet | 2.69 | 2.92 | 2.36 | 2.66 |
微调模型
可以从这里下载微调模型:https://drive.google.com/drive/folders/1V1pCV0BEW3mD-B9MogGrz_P91UhTtuE_?usp=sharing
快速开始:
使用pip安装:
pip3 install sixdrepnet
使用示例:
# 导入 SixDRepNet
from sixdrepnet import SixDRepNet
import cv2
# 创建模型
# 权重会自动下载
model = SixDRepNet()
img = cv2.imread('/path/to/image.jpg')
pitch, yaw, roll = model.predict(img)
model.draw_axis(img, yaw, pitch, roll)
cv2.imshow("test_window", img)
cv2.waitKey(0)
自行设置:
git clone https://github.com/thohemp/6DRepNet
cd 6DRepNet
设置虚拟环境:
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt # 安装所需包
要运行演示脚本,您需要安装人脸检测器
pip install git+https://github.com/elliottzheng/face-detection.git@master
相机演示:
python ./sixdrepnet/demo.py --snapshot 6DRepNet_300W_LP_AFLW2000.pth \
--cam 0
测试/训练 3DRepNet
准备数据集
下载数据集:
将它们存储在 datasets 目录中。
对于300W-LP和AFLW2000,我们需要创建一个文件名列表。
python create_filename_list.py --root_dir datasets/300W_LP
BIWI数据集需要通过人脸检测器预处理,以从图像中裁剪出人脸。您可以使用这里提供的脚本。对于BIWI数据集的7:3分割,您可以使用这里的等效脚本。我们将裁剪后的图像大小设置为256。
测试:
python test.py --batch_size 64 \
--dataset AFLW2000 \
--data_dir datasets/AFLW2000 \
--filename_list datasets/AFLW2000/files.txt \
--snapshot output/snapshots/1.pth \
--show_viz False
训练
从这里下载预训练的RepVGG模型'RepVGG-B1g2-train.pth',并将其保存在根目录中。
python sixdrepnet/train.py
部署模型
使用转换脚本将训练好的模型重新参数化为推理模型。
python convert.py input-model.tar output-model.pth
加载推理模型时需要设置标志deploy=True
。
model = SixDRepNet(backbone_name='RepVGG-B1g2',
backbone_file='',
deploy=True,
pretrained=False)