这篇论文的核心内容是关于一种用于全范围旋转头部姿态估计的新方法。以下是关键点的总结：

标题

• Towards Robust and Unconstrained Full Range of Rotation Head Pose Estimation

摘要

• 提出了一种新的方法，用于在没有约束的情况下估计全范围旋转的头部姿态。
• 引入了旋转矩阵形式来解决模糊的旋转标签问题，并提出了一种连续的6D旋转矩阵表示方法，以实现高效且鲁棒的直接回归。
• 通过新累积的训练数据集和基于测地线的损失函数，设计了一个能够预测更广泛头部姿态的高级模型。
• 在公共数据集上的广泛评估表明，该方法在效率和鲁棒性方面显著优于其他最先进的方法。

关键词

• 头部姿态估计
• 全范围旋转
• 旋转矩阵
• 6D表示
• 测地线损失

主要贡献

1. 提出了一种简化且高效的6参数旋转矩阵表示方法，用于准确回归头部姿态，避免了歧义问题。
2. 提出了一种基于测地距离的方法，用于网络惩罚，以封装训练损失在特殊正交群SO(3)流形几何内。
3. 利用CMU Panoptic数据集扩展了传统的300W-LP头部姿态数据集，包含了全头部姿态旋转数据。
4. 创建了一个新的头部姿态预测模型，其预测范围超过了当前方法，并且在常见测试数据集上实现了更低的误差。
5. 在多个实验设置中展示了该方法在准确性和鲁棒性方面的优势。
6. 进行了消融研究，以评估模型的每个组成部分对结果的影响。

方法概述

• 使用了旋转矩阵作为旋转表示，以克服欧拉角和四元数表示中的歧义和不连续性问题。
• 提出了一种基于测地距离的损失函数，以稳定学习过程。
• 利用CMU Panoptic数据集和300W-LP数据集，创建了一个新的数据集，用于训练模型。

实验

• 在多个公共数据集上进行了广泛的评估，包括AFLW2000、BIWI和CMU Panoptic数据集。
• 与当前最先进的方法进行了比较，证明了该方法在准确性和鲁棒性方面的优势。

结论

• 该研究解决了全范围头部姿态估计的挑战，并提出了一种新的6D旋转矩阵表示方法，该方法在准确性和鲁棒性方面都取得了显著的改进。

• 源码：https://github.com/thohemp/6DRepNet360

模型转换和onnxruntime推理

将官网的模型Fine-tuned on 300W-LP + Panoptic，Fine-tuned on 300W-LP转为onnx后，使用onnxruntime运行，全部实现代码如下；

​​

• 安装依赖，只需要安装一下几个包即可运行

  ​pip install numpy onnxruntime-gpu tqdm opencv-python​

• 可以在添加人头检测或者人脸检测模型首先提取人头框，并外扩一定区域后送入到人头姿态估计模型中推理；

模型和代码下载

• 6DRepNet360全范围旋转头部姿态估计onnx模型+onnxruntime推理代码+测试视频
  资源中的内容如下：
  

​

可视化结果

‍

代码

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import copy
import os
import sys
import cv2
import numpy as np
import torch
import onnxruntime as ort
from math import cos, sin
#from torch.utils.serialization import load_lua
import scipy.io as sio
from tqdm import tqdmsys.path.append(os.getcwd())def plot_pose_cube(img, yaw, pitch, roll, tdx=None, tdy=None, size=150.):# Input is a cv2 image# pose_params: (pitch, yaw, roll, tdx, tdy)# Where (tdx, tdy) is the translation of the face.# For pose we have [pitch yaw roll tdx tdy tdz scale_factor]p = pitch * np.pi / 180y = -(yaw * np.pi / 180)r = roll * np.pi / 180if tdx != None and tdy != None:face_x = tdx - 0.50 * size face_y = tdy - 0.50 * sizeelse:height, width = img.shape[:2]face_x = width / 2 - 0.5 * sizeface_y = height / 2 - 0.5 * sizex1 = size * (cos(y) * cos(r)) + face_xy1 = size * (cos(p) * sin(r) + cos(r) * sin(p) * sin(y)) + face_y x2 = size * (-cos(y) * sin(r)) + face_xy2 = size * (cos(p) * cos(r) - sin(p) * sin(y) * sin(r)) + face_yx3 = size * (sin(y)) + face_xy3 = size * (-cos(y) * sin(p)) + face_y# Draw base in redcv2.line(img, (int(face_x), int(face_y)), (int(x1),int(y1)),(0,0,255),3, cv2.LINE_AA)cv2.line(img, (int(face_x), int(face_y)), (int(x2),int(y2)),(0,0,255),3, cv2.LINE_AA)cv2.line(img, (int(x2), int(y2)), (int(x2+x1-face_x),int(y2+y1-face_y)),(0,0,255),3, cv2.LINE_AA)cv2.line(img, (int(x1), int(y1)), (int(x1+x2-face_x),int(y1+y2-face_y)),(0,0,255),3, cv2.LINE_AA)# Draw pillars in bluecv2.line(img, (int(face_x), int(face_y)), (int(x3),int(y3)),(255,0,0),2, cv2.LINE_AA)cv2.line(img, (int(x1), int(y1)), (int(x1+x3-face_x),int(y1+y3-face_y)),(255,0,0),2, cv2.LINE_AA)cv2.line(img, (int(x2), int(y2)), (int(x2+x3-face_x),int(y2+y3-face_y)),(255,0,0),2, cv2.LINE_AA)cv2.line(img, (int(x2+x1-face_x),int(y2+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(255,0,0),2, cv2.LINE_AA)# Draw top in greencv2.line(img, (int(x3+x1-face_x),int(y3+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2, cv2.LINE_AA)cv2.line(img, (int(x2+x3-face_x),int(y2+y3-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2, cv2.LINE_AA)cv2.line(img, (int(x3), int(y3)), (int(x3+x1-face_x),int(y3+y1-face_y)),(0,255,0),2, cv2.LINE_AA)cv2.line(img, (int(x3), int(y3)), (int(x3+x2-face_x),int(y3+y2-face_y)),(0,255,0),2, cv2.LINE_AA)return imgdef draw_axis(img, yaw, pitch, roll, tdx=None, tdy=None, size = 100):pitch = pitch * np.pi / 180yaw = -(yaw * np.pi / 180)roll = roll * np.pi / 180if tdx != None and tdy != None:tdx = tdxtdy = tdyelse:height, width = img.shape[:2]tdx = width / 2tdy = height / 2# X-Axis pointing to right. drawn in redx1 = size * (cos(yaw) * cos(roll)) + tdxy1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy# Y-Axis | drawn in green#        vx2 = size * (-cos(yaw) * sin(roll)) + tdxy2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy# Z-Axis (out of the screen) drawn in bluex3 = size * (sin(yaw)) + tdxy3 = size * (-cos(yaw) * sin(pitch)) + tdycv2.line(img, (int(tdx), int(tdy)), (int(x1),int(y1)),(0,0,255),4, cv2.LINE_AA)cv2.line(img, (int(tdx), int(tdy)), (int(x2),int(y2)),(0,255,0),4, cv2.LINE_AA)cv2.line(img, (int(tdx), int(tdy)), (int(x3),int(y3)),(255,0,0),4, cv2.LINE_AA)return imgdef get_pose_params_from_mat(mat_path):# This functions gets the pose parameters from the .mat# Annotations that come with the Pose_300W_LP dataset.mat = sio.loadmat(mat_path)# [pitch yaw roll tdx tdy tdz scale_factor]pre_pose_params = mat['Pose_Para'][0]# Get [pitch, yaw, roll, tdx, tdy]pose_params = pre_pose_params[:5]return pose_paramsdef get_ypr_from_mat(mat_path):# Get yaw, pitch, roll from .mat annotation.# They are in radiansmat = sio.loadmat(mat_path)# [pitch yaw roll tdx tdy tdz scale_factor]pre_pose_params = mat['Pose_Para'][0]# Get [pitch, yaw, roll]pose_params = pre_pose_params[:3]return pose_paramsdef get_pt2d_from_mat(mat_path):# Get 2D landmarksmat = sio.loadmat(mat_path)pt2d = mat['pt2d']return pt2d# batch*n
def normalize_vector(v):batch = v.shape[0]v_mag = torch.sqrt(v.pow(2).sum(1))# batchgpu = v_mag.get_device()if gpu < 0:eps = torch.autograd.Variable(torch.FloatTensor([1e-8])).to(torch.device('cpu'))else:eps = torch.autograd.Variable(torch.FloatTensor([1e-8])).to(torch.device('cuda:%d' % gpu))v_mag = torch.max(v_mag, eps)v_mag = v_mag.view(batch,1).expand(batch,v.shape[1])v = v/v_magreturn v# u, v batch*n
def cross_product(u, v):batch = u.shape[0]#print (u.shape)#print (v.shape)i = u[:,1]*v[:,2] - u[:,2]*v[:,1]j = u[:,2]*v[:,0] - u[:,0]*v[:,2]k = u[:,0]*v[:,1] - u[:,1]*v[:,0]out = torch.cat((i.view(batch,1), j.view(batch,1), k.view(batch,1)),1) #batch*3return out#poses batch*6
#poses
def compute_rotation_matrix_from_ortho6d(poses):x_raw = poses[:,0:3] #batch*3y_raw = poses[:,3:6] #batch*3x = normalize_vector(x_raw) #batch*3z = cross_product(x,y_raw) #batch*3z = normalize_vector(z) #batch*3y = cross_product(z,x) #batch*3x = x.view(-1,3,1)y = y.view(-1,3,1)z = z.view(-1,3,1)matrix = torch.cat((x,y,z), 2) #batch*3*3return matrix#input batch*4*4 or batch*3*3
#output torch batch*3 x, y, z in radiant
#the rotation is in the sequence of x,y,z
def compute_euler_angles_from_rotation_matrices(rotation_matrices):batch = rotation_matrices.shape[0]R = rotation_matricessy = torch.sqrt(R[:,0,0]*R[:,0,0]+R[:,1,0]*R[:,1,0])singular = sy<1e-6singular = singular.float()x = torch.atan2(R[:,2,1], R[:,2,2])y = torch.atan2(-R[:,2,0], sy)z = torch.atan2(R[:,1,0],R[:,0,0])xs = torch.atan2(-R[:,1,2], R[:,1,1])ys = torch.atan2(-R[:,2,0], sy)zs = R[:,1,0]*0gpu = rotation_matrices.get_device()if gpu < 0:out_euler = torch.autograd.Variable(torch.zeros(batch,3)).to(torch.device('cpu'))else:out_euler = torch.autograd.Variable(torch.zeros(batch,3)).to(torch.device('cuda:%d' % gpu))out_euler[:,0] = x*(1-singular)+xs*singularout_euler[:,1] = y*(1-singular)+ys*singularout_euler[:,2] = z*(1-singular)+zs*singularreturn out_eulerdef get_R(x,y,z):''' Get rotation matrix from three rotation angles (radians). right-handed.Args:angles: [3,]. x, y, z anglesReturns:R: [3, 3]. rotation matrix.'''# xRx = np.array([[1, 0, 0],[0, np.cos(x), -np.sin(x)],[0, np.sin(x), np.cos(x)]])# yRy = np.array([[np.cos(y), 0, np.sin(y)],[0, 1, 0],[-np.sin(y), 0, np.cos(y)]])# zRz = np.array([[np.cos(z), -np.sin(z), 0],[np.sin(z), np.cos(z), 0],[0, 0, 1]])R = Rz.dot(Ry.dot(Rx))return Rif __name__ == '__main__':onnx_path = "./snapshots/6DRepNet360_Full-Rotation_300W_LP+Panoptic_1.onnx"vid_file = "./test_vid/11-FemaleGlasses.avi.avi"vid_result_path = './test_vid_result'vid_save_file = os.path.join(vid_result_path, os.path.basename(vid_file))if not os.path.exists(vid_result_path):os.makedirs(vid_result_path)device=torch.device(0) if torch.cuda.is_available() else torch.device('cpu')ort_session = ort.InferenceSession(onnx_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])cap = cv2.VideoCapture(vid_file)nums = cap.get(cv2.CAP_PROP_FRAME_COUNT)height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))fps = cap.get(cv2.CAP_PROP_FPS)ret, frame = cap.read()# import pdb;pdb.set_trace()out = cv2.VideoWriter(vid_save_file, cv2.VideoWriter_fourcc(*'XVID'), fps, (width, height))for i in tqdm(range(int(nums))):ret, img = cap.read()if not ret or img is None:break# 125, 45 410, 350image = copy.deepcopy(img)img = img[45:250, 125:410]  # 根据人头位置更改img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_LINEAR)img = np.array(img)img = torch.tensor(img.transpose(2, 0, 1).astype('float32')).unsqueeze(0)# mean = [0., 0., 0.]# std = [1.0, 1.0, 1.0]mean=[0.485, 0.456, 0.406]std=[0.229, 0.224, 0.225]img = img / 255mean = torch.tensor(mean).reshape(1, 3, 1, 1)std = torch.tensor(std).reshape(1, 3, 1, 1)img = (img - mean) / stdinputs = {ort_session.get_inputs()[0].name: img.numpy()}outs = ort_session.run(None, inputs)R_pred = compute_rotation_matrix_from_ortho6d(torch.Tensor(outs[0]).to(device))euler = compute_euler_angles_from_rotation_matrices(R_pred) * 180 / np.pip_pred_deg = euler[:, 0].cpu()y_pred_deg = euler[:, 1].cpu()r_pred_deg = euler[:, 2].cpu()draw_axis(image, y_pred_deg[0], p_pred_deg[0], r_pred_deg[0], tdx=60, tdy=100, size=50)plot_pose_cube(image, y_pred_deg[0], p_pred_deg[0], r_pred_deg[0] ,size=50)txt = "pitch: %.4f, yaw: %.4f, roll: %.4f" % (float(p_pred_deg[0]), float(y_pred_deg[0]), float(r_pred_deg[0]))image = cv2.putText(image, txt, (int(50), int(40)), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5,color=(0, 255, 0), thickness=2)out.write(image)out.release()