fulyvv/EasyMocap
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[vis] update realtime visualization

Browse Source
This commit is contained in:
shuaiqing 2021-06-25 21:17:22 +08:00
parent 0ab05bdafc
commit 0c2d5f2d64
11 changed files with 351 additions and 101 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
config/vis3d/o3d_scene_smpl.yml Normal file
View File

@ -0,0 +1,11 @@
parent: "config/vis3d/o3d_scene.yml"
body_model:
module: "easymocap.smplmodel.body_model.SMPLlayer"
args:
_no_merge_: True
model_path: "data/smplx/smpl"
model_type: "smpl"
gender: "neutral"
device: "cuda"
regressor_path: "data/smplx/J_regressor_body25.npy"

5
easymocap/config/vis_socket.py
View File

@ -2,8 +2,8 @@
@ Date: 2021-05-30 11:17:18
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-06-12 14:56:00
@ FilePath: /EasyMocapRelease/easymocap/config/vis_socket.py
@ LastEditTime: 2021-06-22 10:35:26
@ FilePath: /EasyMocap/easymocap/config/vis_socket.py
'''
from .baseconfig import CN
from .baseconfig import Config as BaseConfig
@ -35,6 +35,7 @@ class Config(BaseConfig):
# skel
cfg.skel = CN()
cfg.skel.joint_radius = 0.02
cfg.body_model_template = "none"
# camera
cfg.camera = CN()
cfg.camera.phi = 0

3
easymocap/config/yacs.py
View File

@ -382,6 +382,9 @@ def _merge_a_into_b(a, b, root, key_list):
isinstance(b, CfgNode),
"`b` (cur type {}) must be an instance of {}".format(type(b), CfgNode),
)
if '_no_merge_' in a.keys() and a['_no_merge_']:
b.clear()
a.pop('_no_merge_')
for k, v_ in a.items():
full_key = ".".join(key_list + [k])

116
easymocap/smplmodel/body_model.py
View File

@ -2,8 +2,8 @@
@ Date: 2020-11-18 14:04:10
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-05-27 20:35:10
@ FilePath: /EasyMocapRelease/easymocap/smplmodel/body_model.py
@ LastEditTime: 2021-06-22 13:44:10
@ FilePath: /EasyMocap/easymocap/smplmodel/body_model.py
'''
import torch
import torch.nn as nn
@ -39,18 +39,7 @@ def load_regressor(regressor_path):
import ipdb; ipdb.set_trace()
return X_regressor
NUM_POSES = {'smpl': 72, 'smplh': 78, 'smplx': 66 + 12 + 9, 'mano': 9}
NUM_SHAPES = 10
NUM_EXPR = 10
class SMPLlayer(nn.Module):
def __init__(self, model_path, model_type='smpl', gender='neutral', device=None,
regressor_path=None) -> None:
super(SMPLlayer, self).__init__()
dtype = torch.float32
self.dtype = dtype
self.device = device
self.model_type = model_type
# create the SMPL model
def load_bodydata(model_path, gender):
if osp.isdir(model_path):
model_fn = 'SMPL_{}.{ext}'.format(gender.upper(), ext='pkl')
smpl_path = osp.join(model_path, model_fn)
@ -61,23 +50,71 @@ class SMPLlayer(nn.Module):
with open(smpl_path, 'rb') as smpl_file:
data = pickle.load(smpl_file, encoding='latin1')
return data
NUM_POSES = {'smpl': 72, 'smplh': 78, 'smplx': 66 + 12 + 9, 'mano': 9}
NUM_SHAPES = 10
NUM_EXPR = 10
class SMPLlayer(nn.Module):
def __init__(self, model_path, model_type='smpl', gender='neutral', device=None,
regressor_path=None,
use_pose_blending=True, use_shape_blending=True, use_joints=True,
with_color=False,
**kwargs) -> None:
super(SMPLlayer, self).__init__()
dtype = torch.float32
self.dtype = dtype
self.use_pose_blending = use_pose_blending
self.use_shape_blending = use_shape_blending
self.use_joints = use_joints
if isinstance(device, str):
device = torch.device(device)
self.device = device
self.model_type = model_type
# create the SMPL model
data = load_bodydata(model_path, gender)
if with_color:
self.color = data['vertex_colors']
else:
self.color = None
self.faces = data['f']
self.register_buffer('faces_tensor',
to_tensor(to_np(self.faces, dtype=np.int64),
dtype=torch.long))
for key in ['J_regressor', 'v_template', 'weights']:
val = to_tensor(to_np(data[key]), dtype=dtype)
self.register_buffer(key, val)
# add poseblending
if use_pose_blending:
# Pose blend shape basis: 6890 x 3 x 207, reshaped to 6890*3 x 207
num_pose_basis = data['posedirs'].shape[-1]
# 207 x 20670
posedirs = data['posedirs']
data['posedirs'] = np.reshape(data['posedirs'], [-1, num_pose_basis]).T
val = to_tensor(to_np(data['posedirs']), dtype=dtype)
self.register_buffer('posedirs', val)
else:
self.posedirs = None
# add shape blending
if use_shape_blending:
val = to_tensor(to_np(data['shapedirs']), dtype=dtype)
self.register_buffer('shapedirs', val)
else:
self.shapedirs = None
if use_shape_blending:
self.J_shaped = None
else:
val = to_tensor(to_np(data['J']), dtype=dtype)
self.register_buffer('J_shaped', val)
for key in ['J_regressor', 'v_template', 'weights', 'posedirs', 'shapedirs']:
val = to_tensor(to_np(data[key]), dtype=dtype)
self.register_buffer(key, val)
self.nVertices = self.v_template.shape[0]
# indices of parents for each joints
parents = to_tensor(to_np(data['kintree_table'][0])).long()
parents[0] = -1
self.register_buffer('parents', parents)
if self.use_shape_blending:
if self.model_type == 'smplx':
# shape
self.num_expression_coeffs = 10
@ -86,7 +123,7 @@ class SMPLlayer(nn.Module):
elif self.model_type in ['smpl', 'smplh']:
self.shapedirs = self.shapedirs[:, :, :NUM_SHAPES]
# joints regressor
if regressor_path is not None:
if regressor_path is not None and use_joints:
X_regressor = load_regressor(regressor_path)
X_regressor = torch.cat((self.J_regressor, X_regressor), dim=0)
@ -95,14 +132,20 @@ class SMPLlayer(nn.Module):
j_J_regressor[i, i] = 1
j_v_template = X_regressor @ self.v_template
#
j_shapedirs = torch.einsum('vij,kv->kij', [self.shapedirs, X_regressor])
# (25, 24)
j_weights = X_regressor @ self.weights
if self.use_pose_blending:
j_posedirs = torch.einsum('ab, bde->ade', [X_regressor, torch.Tensor(posedirs)]).numpy()
j_posedirs = np.reshape(j_posedirs, [-1, num_pose_basis]).T
j_posedirs = to_tensor(j_posedirs)
self.register_buffer('j_posedirs', j_posedirs)
else:
self.j_posedirs = None
if self.use_shape_blending:
j_shapedirs = torch.einsum('vij,kv->kij', [self.shapedirs, X_regressor])
self.register_buffer('j_shapedirs', j_shapedirs)
else:
self.j_shapedirs = None
self.register_buffer('j_weights', j_weights)
self.register_buffer('j_v_template', j_v_template)
self.register_buffer('j_J_regressor', j_J_regressor)
@ -122,13 +165,22 @@ class SMPLlayer(nn.Module):
self.use_flat_mean = True
elif self.model_type == 'mano':
# TODO:write this into config file
# self.num_pca_comps = 12
# self.use_pca = True
# if self.use_pca:
# NUM_POSES['mano'] = self.num_pca_comps + 3
# else:
# NUM_POSES['mano'] = 45 + 3
# self.use_flat_mean = True
self.num_pca_comps = 12
self.use_pca = True
self.use_flat_mean = True
if self.use_pca:
NUM_POSES['mano'] = self.num_pca_comps + 3
else:
NUM_POSES['mano'] = 45 + 3
self.use_flat_mean = True
val = to_tensor(to_np(data['hands_mean'].reshape(1, -1)), dtype=dtype)
self.register_buffer('mHandsMean', val)
val = to_tensor(to_np(data['hands_components'][:self.num_pca_comps, :]), dtype=dtype)
@ -144,6 +196,7 @@ class SMPLlayer(nn.Module):
self.register_buffer('mHandsComponents'+key[0], val)
self.use_pca = True
self.use_flat_mean = True
self.to(self.device)
@staticmethod
def extend_hand(poses, use_pca, use_flat_mean, coeffs, mean):
@ -245,7 +298,8 @@ class SMPLlayer(nn.Module):
poses = self.extend_pose(poses)
return poses.detach().cpu().numpy()
def forward(self, poses, shapes, Rh=None, Th=None, expression=None, return_verts=True, return_tensor=True, only_shape=False, **kwargs):
def forward(self, poses, shapes, Rh=None, Th=None, expression=None,
return_verts=True, return_tensor=True, return_smpl_joints=False, only_shape=False, **kwargs):
""" Forward pass for SMPL model
Args:
@ -285,31 +339,41 @@ class SMPLlayer(nn.Module):
shapes = torch.cat([shapes, expression], dim=1)
# process poses
poses = self.extend_pose(poses)
if return_verts:
if return_verts or not self.use_joints:
vertices, joints = lbs(shapes, poses, self.v_template,
self.shapedirs, self.posedirs,
self.J_regressor, self.parents,
self.weights, pose2rot=True, dtype=self.dtype)
self.weights, pose2rot=True, dtype=self.dtype,
use_pose_blending=self.use_pose_blending, use_shape_blending=self.use_shape_blending, J_shaped=self.J_shaped)
if not self.use_joints and not return_verts:
vertices = joints
else:
vertices, joints = lbs(shapes, poses, self.j_v_template,
self.j_shapedirs, self.j_posedirs,
self.j_J_regressor, self.parents,
self.j_weights, pose2rot=True, dtype=self.dtype, only_shape=only_shape)
self.j_weights, pose2rot=True, dtype=self.dtype, only_shape=only_shape,
use_pose_blending=self.use_pose_blending, use_shape_blending=self.use_shape_blending, J_shaped=self.J_shaped)
if return_smpl_joints:
vertices = vertices[:, :self.J_regressor.shape[0], :]
else:
vertices = vertices[:, self.J_regressor.shape[0]:, :]
vertices = torch.matmul(vertices, rot.transpose(1, 2)) + transl
if not return_tensor:
vertices = vertices.detach().cpu().numpy()
return vertices
def init_params(self, nFrames):
def init_params(self, nFrames=1, nShapes=1, ret_tensor=False):
params = {
'poses': np.zeros((nFrames, NUM_POSES[self.model_type])),
'shapes': np.zeros((1, NUM_SHAPES)),
'shapes': np.zeros((nShapes, NUM_SHAPES)),
'Rh': np.zeros((nFrames, 3)),
'Th': np.zeros((nFrames, 3)),
}
if self.model_type == 'smplx':
params['expression'] = np.zeros((nFrames, NUM_EXPR))
if ret_tensor:
for key in params.keys():
params[key] = to_tensor(params[key], self.dtype, self.device)
return params
def check_params(self, body_params):

27
easymocap/smplmodel/lbs.py
View File

@ -140,7 +140,8 @@ def vertices2landmarks(vertices, faces, lmk_faces_idx, lmk_bary_coords):
def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents,
lbs_weights, pose2rot=True, dtype=torch.float32, only_shape=False):
lbs_weights, pose2rot=True, dtype=torch.float32, only_shape=False,
use_shape_blending=True, use_pose_blending=True, J_shaped=None):
''' Performs Linear Blend Skinning with the given shape and pose parameters
Parameters
@ -183,32 +184,34 @@ def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents,
device = betas.device
# Add shape contribution
if use_shape_blending:
v_shaped = v_template + blend_shapes(betas, shapedirs)
# Get the joints
# NxJx3 array
J = vertices2joints(J_regressor, v_shaped)
else:
v_shaped = v_template.unsqueeze(0).expand(batch_size, -1, -1)
assert J_shaped is not None
J = J_shaped[None].expand(batch_size, -1, -1)
if only_shape:
return v_shaped, J
# 3. Add pose blend shapes
# N x J x 3 x 3
ident = torch.eye(3, dtype=dtype, device=device)
if pose2rot:
rot_mats = batch_rodrigues(
pose.view(-1, 3), dtype=dtype).view([batch_size, -1, 3, 3])
pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
# (N x P) x (P, V * 3) -> N x V x 3
pose_offsets = torch.matmul(pose_feature, posedirs) \
.view(batch_size, -1, 3)
else:
pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
rot_mats = pose.view(batch_size, -1, 3, 3)
pose_offsets = torch.matmul(pose_feature.view(batch_size, -1),
posedirs).view(batch_size, -1, 3)
if use_pose_blending:
ident = torch.eye(3, dtype=dtype, device=device)
pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
pose_offsets = torch.matmul(pose_feature, posedirs) \
.view(batch_size, -1, 3)
v_posed = pose_offsets + v_shaped
else:
v_posed = v_shaped
# 4. Get the global joint location
J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)

11
easymocap/socket/base_client.py
View File

@ -2,11 +2,11 @@
@ Date: 2021-05-25 13:39:07
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-06-04 16:43:39
@ FilePath: /EasyMocapRelease/easymocap/socket/base_client.py
@ LastEditTime: 2021-06-16 14:42:35
@ FilePath: /EasyMocap/easymocap/socket/base_client.py
'''
import socket
from .utils import encode_detect
from .utils import encode_detect, encode_smpl
class BaseSocketClient:
def __init__(self, host, port) -> None:
@ -21,5 +21,10 @@ class BaseSocketClient:
self.s.send(bytes('{}\n'.format(len(val)), 'ascii'))
self.s.sendall(val)
def send_smpl(self, data):
val = encode_smpl(data)
self.s.send(bytes('{}\n'.format(len(val)), 'ascii'))
self.s.sendall(val)
def close(self):
self.s.close()

82
easymocap/socket/o3d.py
View File

@ -2,7 +2,7 @@
@ Date: 2021-05-25 11:15:53
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-06-12 14:54:43
@ LastEditTime: 2021-06-25 21:16:02
@ FilePath: /EasyMocapRelease/easymocap/socket/o3d.py
'''
import open3d as o3d
@ -16,17 +16,26 @@ import numpy as np
from os.path import join
import os
from ..assignment.criterion import CritRange
import copy
rotate = False
def o3d_callback_rotate(vis):
global rotate
rotate = not rotate
return False
class VisOpen3DSocket(BaseSocket):
def __init__(self, host, port, cfg) -> None:
# output
self.write = cfg.write
self.out = cfg.out
self.cfg = cfg
if self.write:
print('[Info] capture the screen to {}'.format(self.out))
os.makedirs(self.out, exist_ok=True)
# scene
vis = o3d.visualization.Visualizer()
vis = o3d.visualization.VisualizerWithKeyCallback()
vis.register_key_callback(ord('A'), o3d_callback_rotate)
vis.create_window(window_name='Visualizer', width=cfg.width, height=cfg.height)
self.vis = vis
# load the scene
@ -41,10 +50,16 @@ class VisOpen3DSocket(BaseSocket):
# create vis => update => super() init
super().__init__(host, port, debug=cfg.debug)
self.block = cfg.block
if os.path.exists(cfg.body_model_template):
body_template = o3d.io.read_triangle_mesh(cfg.body_model_template)
self.body_template = body_template
else:
self.body_template = None
self.body_model = load_object(cfg.body_model.module, cfg.body_model.args)
zero_params = self.body_model.init_params(1)
self.max_human = cfg.max_human
self.track = cfg.track
self.filter = cfg.filter
self.camera_pose = cfg.camera.camera_pose
self.init_camera(cfg.camera.camera_pose)
self.zero_vertices = Vector3dVector(np.zeros((self.body_model.nVertices, 3)))
@ -59,13 +74,35 @@ class VisOpen3DSocket(BaseSocket):
self.new_frames = cfg.new_frames
def add_human(self, zero_params):
vertices = self.body_model(zero_params)[0]
vertices = self.body_model(return_verts=True, return_tensor=False, **zero_params)[0]
self.vertices.append(vertices)
mesh = create_mesh(vertices=vertices, faces=self.body_model.faces)
if self.body_template is None: # create template
mesh = create_mesh(vertices=vertices, faces=self.body_model.faces, colors=self.body_model.color)
else:
mesh = copy.deepcopy(self.body_template)
self.meshes.append(mesh)
self.vis.add_geometry(mesh)
self.init_camera(self.camera_pose)
@staticmethod
def set_camera(cfg, camera_pose):
theta, phi = np.deg2rad(-(cfg.camera.theta + 90)), np.deg2rad(cfg.camera.phi)
theta = theta + np.pi
st, ct = np.sin(theta), np.cos(theta)
sp, cp = np.sin(phi), np.cos(phi)
rot_x = np.array([
[1., 0., 0.],
[0., ct, -st],
[0, st, ct]
])
rot_z = np.array([
[cp, -sp, 0],
[sp, cp, 0.],
[0., 0., 1.]
])
camera_pose[:3, :3] = rot_x @ rot_z
return camera_pose
def init_camera(self, camera_pose):
ctr = self.vis.get_view_control()
init_param = ctr.convert_to_pinhole_camera_parameters()
@ -73,6 +110,11 @@ class VisOpen3DSocket(BaseSocket):
init_param.extrinsic = np.array(camera_pose)
ctr.convert_from_pinhole_camera_parameters(init_param)
def get_camera(self):
ctr = self.vis.get_view_control()
init_param = ctr.convert_to_pinhole_camera_parameters()
return np.array(init_param.extrinsic)
def filter_human(self, datas):
datas_new = []
for data in datas:
@ -90,15 +132,33 @@ class VisOpen3DSocket(BaseSocket):
def main(self, datas):
if self.debug:log('[Info] Load data {}'.format(self.count))
if isinstance(datas, str):
datas = json.loads(datas)
for data in datas:
for key in data.keys():
if key == 'id':
continue
data[key] = np.array(data[key])
if 'keypoints3d' not in data.keys() and self.filter:
data['keypoints3d'] = self.body_model(return_verts=False, return_tensor=False, **data)[0]
if self.filter:
datas = self.filter_human(datas)
with Timer('forward'):
params = []
for i, data in enumerate(datas):
if i >= len(self.meshes):
print('[Error] the number of human exceeds!')
self.add_human(np.array(data['keypoints3d']))
vertices = self.body_model(data['keypoints3d'])
self.vertices[i] = Vector3dVector(vertices[0])
self.add_human(data)
if 'vertices' in data.keys():
vertices = data['vertices']
self.vertices[i] = Vector3dVector(vertices)
else:
params.append(data)
if len(params) > 0:
params = self.body_model.merge_params(params, share_shape=False)
vertices = self.body_model(return_verts=True, return_tensor=False, **params)
for i in range(vertices.shape[0]):
self.vertices[i] = Vector3dVector(vertices[i])
for i in range(len(datas), len(self.meshes)):
self.vertices[i] = self.zero_vertices
# Open3D will lock the thread here
@ -109,6 +169,12 @@ class VisOpen3DSocket(BaseSocket):
col = get_rgb_01(datas[i]['id'])
self.meshes[i].paint_uniform_color(col)
def o3dcallback(self):
if rotate:
self.cfg.camera.phi += np.pi/10
camera_pose = self.set_camera(self.cfg, self.get_camera())
self.init_camera(camera_pose)
def update(self):
if self.disconnect and not self.block:
self.previous.clear()
@ -123,6 +189,7 @@ class VisOpen3DSocket(BaseSocket):
for mesh in self.meshes:
mesh.compute_triangle_normals()
self.vis.update_geometry(mesh)
self.o3dcallback()
self.vis.poll_events()
self.vis.update_renderer()
if self.write:
@ -132,5 +199,6 @@ class VisOpen3DSocket(BaseSocket):
self.count += 1
else:
with Timer('update renderer', True):
self.o3dcallback()
self.vis.poll_events()
self.vis.update_renderer()

9
easymocap/socket/utils.py
View File

@ -2,8 +2,8 @@
@ Date: 2021-05-24 20:07:34
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-06-04 16:29:35
@ FilePath: /EasyMocapRelease/media/qing/Project/mirror/EasyMocap/easymocap/socket/utils.py
@ LastEditTime: 2021-06-16 14:42:23
@ FilePath: /EasyMocap/easymocap/socket/utils.py
'''
import cv2
import numpy as np
@ -14,6 +14,11 @@ def encode_detect(data):
res = res.replace('\r', '').replace('\n', '').replace(' ', '')
return res.encode('ascii')
def encode_smpl(data):
res = write_common_results(None, data, ['poses', 'shapes', 'Rh', 'Th'])
res = res.replace('\r', '').replace('\n', '').replace(' ', '')
return res.encode('ascii')
def encode_image(image):
fourcc = [int(cv2.IMWRITE_JPEG_QUALITY), 90]
#frame을 binary 형태로 변환 jpg로 decoding

54
easymocap/visualize/geometry.py
View File

@ -2,13 +2,65 @@
@ Date: 2021-01-17 22:44:34
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-05-25 14:01:24
@ LastEditTime: 2021-06-20 17:24:12
@ FilePath: /EasyMocap/easymocap/visualize/geometry.py
'''
import numpy as np
import cv2
import numpy as np
from tqdm import tqdm
from os.path import join
def load_sphere():
cur_dir = os.path.dirname(__file__)
faces = np.loadtxt(join(cur_dir, 'sphere_faces_20.txt'), dtype=np.int)
vertices = np.loadtxt(join(cur_dir, 'sphere_vertices_20.txt'))
return vertices, faces
def load_cylinder():
cur_dir = os.path.dirname(__file__)
faces = np.loadtxt(join(cur_dir, 'cylinder_faces_20.txt'), dtype=np.int)
vertices = np.loadtxt(join(cur_dir, 'cylinder_vertices_20.txt'))
return vertices, faces
def create_point(points, r=0.01):
""" create sphere
Args:
points (array): (N, 3)/(N, 4)
r (float, optional): radius. Defaults to 0.01.
"""
nPoints = points.shape[0]
vert, face = load_sphere()
nVerts = vert.shape[0]
vert = vert[None, :, :].repeat(points.shape[0], 0)
vert = vert + points[:, None, :]
verts = np.vstack(vert)
face = face[None, :, :].repeat(points.shape[0], 0)
face = face + nVerts * np.arange(nPoints).reshape(nPoints, 1, 1)
faces = np.vstack(face)
return {'vertices': verts, 'faces': faces, 'name': 'points'}
def calRot(axis, direc):
direc = direc/np.linalg.norm(direc)
axis = axis/np.linalg.norm(axis)
rotdir = np.cross(axis, direc)
rotdir = rotdir/np.linalg.norm(rotdir)
rotdir = rotdir * np.arccos(np.dot(direc, axis))
rotmat, _ = cv2.Rodrigues(rotdir)
return rotmat
def create_line(start, end, r=0.01, col=None):
length = np.linalg.norm(end[:3] - start[:3])
vertices, faces = load_cylinder()
vertices[:, :2] *= r
vertices[:, 2] *= length/2
rotmat = calRot(np.array([0, 0, 1]), end - start)
vertices = vertices @ rotmat.T + (start + end)/2
ret = {'vertices': vertices, 'faces': faces, 'name': 'line'}
if col is not None:
ret['colors'] = col.reshape(-1, 3).repeat(vertices.shape[0], 0)
return ret
def create_ground(
center=[0, 0, 0], xdir=[1, 0, 0], ydir=[0, 1, 0], # 位置

20
easymocap/visualize/o3dwrapper.py
View File

@ -2,17 +2,25 @@
@ Date: 2021-04-25 15:52:01
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-05-25 11:48:49
@ LastEditTime: 2021-06-20 15:22:59
@ FilePath: /EasyMocap/easymocap/visualize/o3dwrapper.py
'''
import open3d as o3d
import numpy as np
from .geometry import create_ground as create_ground_
from .geometry import create_point as create_point_
from .geometry import create_line as create_line_
Vector3dVector = o3d.utility.Vector3dVector
Vector3iVector = o3d.utility.Vector3iVector
Vector2iVector = o3d.utility.Vector2iVector
TriangleMesh = o3d.geometry.TriangleMesh
load_mesh = o3d.io.read_triangle_mesh
vis = o3d.visualization.draw_geometries
def _create_cylinder():
# create_cylinder(radius=1.0, height=2.0, resolution=20, split=4, create_uv_map=False)
pass
def create_mesh(vertices, faces, colors=None, **kwargs):
mesh = TriangleMesh()
@ -25,14 +33,20 @@ def create_mesh(vertices, faces, colors=None, **kwargs):
mesh.compute_vertex_normals()
return mesh
def create_point(**kwargs):
return create_mesh(**create_point_(**kwargs))
def create_line(**kwargs):
return create_mesh(**create_line_(**kwargs))
def create_ground(**kwargs):
from .geometry import create_ground as create_ground_
ground = create_ground_(**kwargs)
return create_mesh(**ground)
def create_coord(camera = [0,0,0], radius=1):
def create_coord(camera = [0,0,0], radius=1, scale=1):
camera_frame = TriangleMesh.create_coordinate_frame(
size=radius, origin=camera)
camera_frame.scale(scale)
return camera_frame
def create_bbox(min_bound=(-3., -3., 0), max_bound=(3., 3., 2), flip=False):

40
easymocap/visualize/skelmodel.py
View File

@ -2,7 +2,7 @@
@ Date: 2021-01-17 21:38:19
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-06-04 15:52:49
@ LastEditTime: 2021-06-18 18:48:37
@ FilePath: /EasyMocap/easymocap/visualize/skelmodel.py
'''
import numpy as np
@ -24,9 +24,12 @@ def calTransformation(v_i, v_j, r, adaptr=False, ratio=10):
length = np.linalg.norm(direc)
direc = direc/length
rotdir = np.cross(xaxis, direc)
if np.linalg.norm(rotdir) > 1e-3:
rotdir = rotdir/np.linalg.norm(rotdir)
rotdir = rotdir * np.arccos(np.dot(direc, xaxis))
rotmat, _ = cv2.Rodrigues(rotdir)
else:
rotmat = np.eye(3)
# set the minimal radius for the finger and face
shrink = min(max(length/ratio, 0.005), 0.05)
eigval = np.array([[length/2/r, 0, 0], [0, shrink, 0], [0, 0, shrink]])
@ -44,6 +47,7 @@ class SkelModel:
config = CONFIG[body_type]
self.nJoints = config['nJoints']
self.kintree = config['kintree']
self.device = 'none'
cur_dir = os.path.dirname(__file__)
faces = np.loadtxt(join(cur_dir, 'sphere_faces_20.txt'), dtype=np.int)
self.vertices = np.loadtxt(join(cur_dir, 'sphere_vertices_20.txt'))
@ -54,34 +58,54 @@ class SkelModel:
for nk in range(len(self.kintree)):
faces_all.append(faces + self.nJoints*self.vertices.shape[0] + nk*self.vertices.shape[0])
self.faces = np.vstack(faces_all)
self.color = None
self.nVertices = self.vertices.shape[0] * self.nJoints + self.vertices.shape[0] * len(self.kintree)
self.joint_radius = joint_radius
def __call__(self, keypoints3d, id=0, return_verts=True, return_tensor=False, **kwargs):
vertices_all = []
if len(keypoints3d.shape) == 2:
keypoints3d = keypoints3d[None]
if not return_verts:
return keypoints3d
if keypoints3d.shape[-1] == 3: # add confidence
keypoints3d = np.hstack((keypoints3d, np.ones((keypoints3d.shape[0], 1))))
r = self.joint_radius
# joints
min_conf = 0.1
verts_final = []
for nper in range(keypoints3d.shape[0]):
vertices_all = []
kpts3d = keypoints3d[nper]
for nj in range(self.nJoints):
if keypoints3d[nj, -1] < 0.01:
if kpts3d[nj, -1] < min_conf:
vertices_all.append(self.vertices*0.001)
continue
vertices_all.append(self.vertices*r + keypoints3d[nj:nj+1, :3])
vertices_all.append(self.vertices*r + kpts3d[nj:nj+1, :3])
# limb
for nk, (i, j) in enumerate(self.kintree):
if keypoints3d[i][-1] < 0.1 or keypoints3d[j][-1] < 0.1:
if kpts3d[i][-1] < min_conf or kpts3d[j][-1] < min_conf:
vertices_all.append(self.vertices*0.001)
continue
T, _, length = calTransformation(keypoints3d[i, :3], keypoints3d[j, :3], r=1)
T, _, length = calTransformation(kpts3d[i, :3], kpts3d[j, :3], r=1)
if length > 2: # 超过两米的
vertices_all.append(self.vertices*0.001)
continue
vertices = self.vertices @ T[:3, :3].T + T[:3, 3:].T
vertices_all.append(vertices)
vertices = np.vstack(vertices_all)
return vertices[None, :, :]
verts_final.append(vertices)
verts_final = np.stack(verts_final)
return verts_final
def to(self, none):
pass
def merge_params(self, params, share_shape=False):
keypoints = np.stack([param['keypoints3d'] for param in params])
return {'keypoints3d': keypoints, 'id': [0]}
def init_params(self, nFrames):
return np.zeros((self.nJoints, 4))
return {'keypoints3d': np.zeros((self.nJoints, 4))}
class SMPLSKEL:
def __init__(self, model_type, gender, body_type) -> None: