EasyMocap/easymocap/dataset/mirror.py
2023-02-10 08:10:48 -06:00

139 lines
5.0 KiB
Python

'''
@ Date: 2021-01-21 19:34:48
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-03-06 18:57:47
@ FilePath: /EasyMocap/code/dataset/mirror.py
'''
import numpy as np
from os.path import join
import os
import cv2
FLIP_BODY25 = [0,1,5,6,7,2,3,4,8,12,13,14,9,10,11,16,15,18,17,22,23,24,19,20,21]
FLIP_BODYHAND = [
0,1,5,6,7,2,3,4,8,12,13,14,9,10,11,16,15,18,17,22,23,24,19,20,21, # body 25
22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, # right hand
]
FLIP_SMPL_VERTICES = np.loadtxt(join(os.path.dirname(__file__), 'smpl_vert_sym.txt'), dtype=int)
def flipPoint2D(point):
if point.shape[-2] == 25:
return point[..., FLIP_BODY25, :]
elif point.shape[-2] == 15:
return point[..., FLIP_BODY25[:15], :]
elif point.shape[-2] == 6890:
return point[..., FLIP_SMPL_VERTICES, :]
import ipdb; ipdb.set_trace()
elif point.shape[-1] == 67:
import ipdb; ipdb.set_trace()
# Permutation of SMPL pose parameters when flipping the shape
_PERMUTATION = {
'smpl': [0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, 20, 23, 22],
'smplh': [0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, 20, 24, 25, 23, 24],
'smplx': [0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, 20, 24, 25, 23, 24, 26, 28, 27],
'smplhfull': [
0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, 20, # body
37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51,
22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36
],
'smplxfull': [
0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, 20, # body
22, 24, 23, # jaw, left eye, right eye
40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, # right hand
25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, # left hand
]
}
PERMUTATION = {}
for key in _PERMUTATION.keys():
res = []
for i in _PERMUTATION[key]:
res.extend([3*i + j for j in range(3)])
PERMUTATION[max(res)+1] = res
def flipSMPLPoses(pose):
"""Flip pose.
const input: (N, 72) -> (N, 72)
The flipping is based on SMPL parameters.
"""
pose = pose[:, PERMUTATION[pose.shape[-1]]]
if pose.shape[1] in [72, 156, 165]:
pose[:, 1::3] = -pose[:, 1::3]
pose[:, 2::3] = -pose[:, 2::3]
elif pose.shape[1] in [78, 87]:
pose[:, 1:66:3] = -pose[:, 1:66:3]
pose[:, 2:66:3] = -pose[:, 2:66:3]
else:
import ipdb; ipdb.set_trace()
# we also negate the second and the third dimension of the axis-angle
return pose
def mirrorPoint3D(point, M):
point_homo = np.hstack([point, np.ones([point.shape[0], 1])])
point_m = (M @ point_homo.T).T[..., :3]
return flipPoint2D(point_m)
def calc_mirror_transform(m):
coeff_mat = np.eye(4)[None, :, :]
coeff_mat = coeff_mat.repeat(m.shape[0], 0)
norm = np.linalg.norm(m[:, :3], keepdims=True, axis=1)
m[:, :3] /= norm
coeff_mat[:, 0, 0] = 1 - 2*m[:, 0]**2
coeff_mat[:, 0, 1] = -2*m[:, 0]*m[:, 1]
coeff_mat[:, 0, 2] = -2*m[:, 0]*m[:, 2]
coeff_mat[:, 0, 3] = -2*m[:, 0]*m[:, 3]
coeff_mat[:, 1, 0] = -2*m[:, 1]*m[:, 0]
coeff_mat[:, 1, 1] = 1-2*m[:, 1]**2
coeff_mat[:, 1, 2] = -2*m[:, 1]*m[:, 2]
coeff_mat[:, 1, 3] = -2*m[:, 1]*m[:, 3]
coeff_mat[:, 2, 0] = -2*m[:, 2]*m[:, 0]
coeff_mat[:, 2, 1] = -2*m[:, 2]*m[:, 1]
coeff_mat[:, 2, 2] = 1-2*m[:, 2]**2
coeff_mat[:, 2, 3] = -2*m[:, 2]*m[:, 3]
return coeff_mat
def get_rotation_from_two_directions(direc0, direc1):
direc0 = direc0/np.linalg.norm(direc0)
direc1 = direc1/np.linalg.norm(direc1)
rotdir = np.cross(direc0, direc1)
if np.linalg.norm(rotdir) < 1e-2:
return np.eye(3)
rotdir = rotdir/np.linalg.norm(rotdir)
rotdir = rotdir * np.arccos(np.dot(direc0, direc1))
rotmat, _ = cv2.Rodrigues(rotdir)
return rotmat
def mirror_Rh(Rh, normals):
rvecs = np.zeros_like(Rh)
for nf in range(Rh.shape[0]):
normal = normals[nf]
rotmat = cv2.Rodrigues(Rh[nf])[0]
rotmat_m = np.zeros((3, 3))
for i in range(3):
rot = rotmat[:, i] - 2*(rotmat[:, i] * normal).sum()*normal
rotmat_m[:, i] = rot
rotmat_m[:, 0] *= -1
rvecs[nf] = cv2.Rodrigues(rotmat_m)[0].T
return rvecs
def flipSMPLParams(params, mirror):
"""Flip pose.
const input: (1, 72) -> (1, 72)
The flipping is based on SMPL parameters.
"""
mirror[:, :3] /= np.linalg.norm(mirror[:, :3], keepdims=True, axis=1)
if mirror.shape[0] == 1 and mirror.shape[0] != params['Rh'].shape[0]:
mirror = mirror.repeat(params['Rh'].shape[0], 0)
M = calc_mirror_transform(mirror)
T = params['Th']
rvecm = mirror_Rh(params['Rh'], mirror[:, :3])
Tnew = np.einsum('bmn,bn->bm', M[:, :3, :3], params['Th']) + M[:, :3, 3]
params = {
'poses': flipSMPLPoses(params['poses']),
'shapes': params['shapes'],
'Rh': rvecm,
'Th': Tnew
}
return params