diff --git a/config/vis3d/o3d_scene_smpl.yml b/config/vis3d/o3d_scene_smpl.yml new file mode 100644 index 0000000..aa8b529 --- /dev/null +++ b/config/vis3d/o3d_scene_smpl.yml @@ -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" \ No newline at end of file diff --git a/easymocap/config/vis_socket.py b/easymocap/config/vis_socket.py index fd521ce..b3dcbd9 100644 --- a/easymocap/config/vis_socket.py +++ b/easymocap/config/vis_socket.py @@ -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 diff --git a/easymocap/config/yacs.py b/easymocap/config/yacs.py index 930f87b..39d3989 100644 --- a/easymocap/config/yacs.py +++ b/easymocap/config/yacs.py @@ -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]) diff --git a/easymocap/smplmodel/body_model.py b/easymocap/smplmodel/body_model.py index f96dc9d..bc9d7db 100644 --- a/easymocap/smplmodel/body_model.py +++ b/easymocap/smplmodel/body_model.py @@ -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,54 +39,91 @@ def load_regressor(regressor_path): import ipdb; ipdb.set_trace() return X_regressor +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) + else: + smpl_path = model_path + assert osp.exists(smpl_path), 'Path {} does not exist!'.format( + smpl_path) + + 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) -> 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 - if osp.isdir(model_path): - model_fn = 'SMPL_{}.{ext}'.format(gender.upper(), ext='pkl') - smpl_path = osp.join(model_path, model_fn) + data = load_bodydata(model_path, gender) + if with_color: + self.color = data['vertex_colors'] else: - smpl_path = model_path - assert osp.exists(smpl_path), 'Path {} does not exist!'.format( - smpl_path) - - with open(smpl_path, 'rb') as smpl_file: - data = pickle.load(smpl_file, encoding='latin1') + self.color = None self.faces = data['f'] self.register_buffer('faces_tensor', to_tensor(to_np(self.faces, dtype=np.int64), dtype=torch.long)) - # 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 - - for key in ['J_regressor', 'v_template', 'weights', 'posedirs', 'shapedirs']: + 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) + + 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.model_type == 'smplx': - # shape - self.num_expression_coeffs = 10 - self.num_shapes = 10 - self.shapedirs = self.shapedirs[:, :, :self.num_shapes+self.num_expression_coeffs] - elif self.model_type in ['smpl', 'smplh']: - self.shapedirs = self.shapedirs[:, :, :NUM_SHAPES] + + if self.use_shape_blending: + if self.model_type == 'smplx': + # shape + self.num_expression_coeffs = 10 + self.num_shapes = 10 + self.shapedirs = self.shapedirs[:, :, :self.num_shapes+self.num_expression_coeffs] + 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 - 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) - self.register_buffer('j_shapedirs', j_shapedirs) + 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) - vertices = vertices[:, self.J_regressor.shape[0]:, :] + 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): diff --git a/easymocap/smplmodel/lbs.py b/easymocap/smplmodel/lbs.py index 0ddfe99..c544911 100644 --- a/easymocap/smplmodel/lbs.py +++ b/easymocap/smplmodel/lbs.py @@ -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 - v_shaped = v_template + blend_shapes(betas, shapedirs) + 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) - # Get the joints - # NxJx3 array - J = vertices2joints(J_regressor, v_shaped) 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) - - v_posed = pose_offsets + v_shaped + 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) diff --git a/easymocap/socket/base_client.py b/easymocap/socket/base_client.py index f4983cf..cbd1719 100644 --- a/easymocap/socket/base_client.py +++ b/easymocap/socket/base_client.py @@ -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() \ No newline at end of file diff --git a/easymocap/socket/o3d.py b/easymocap/socket/o3d.py index 919d4ef..ddc2d74 100644 --- a/easymocap/socket/o3d.py +++ b/easymocap/socket/o3d.py @@ -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)) - datas = json.loads(datas) - datas = self.filter_human(datas) + 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() \ No newline at end of file diff --git a/easymocap/socket/utils.py b/easymocap/socket/utils.py index af020ae..cde700c 100644 --- a/easymocap/socket/utils.py +++ b/easymocap/socket/utils.py @@ -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 diff --git a/easymocap/visualize/geometry.py b/easymocap/visualize/geometry.py index afa84a4..e2ad0e8 100644 --- a/easymocap/visualize/geometry.py +++ b/easymocap/visualize/geometry.py @@ -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], # 位置 diff --git a/easymocap/visualize/o3dwrapper.py b/easymocap/visualize/o3dwrapper.py index 48f80f9..3cb51cd 100644 --- a/easymocap/visualize/o3dwrapper.py +++ b/easymocap/visualize/o3dwrapper.py @@ -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): diff --git a/easymocap/visualize/skelmodel.py b/easymocap/visualize/skelmodel.py index de4090f..e2683e6 100644 --- a/easymocap/visualize/skelmodel.py +++ b/easymocap/visualize/skelmodel.py @@ -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) - rotdir = rotdir/np.linalg.norm(rotdir) - rotdir = rotdir * np.arccos(np.dot(direc, xaxis)) - rotmat, _ = cv2.Rodrigues(rotdir) + 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 - for nj in range(self.nJoints): - if keypoints3d[nj, -1] < 0.01: - vertices_all.append(self.vertices*0.001) - continue - vertices_all.append(self.vertices*r + keypoints3d[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: - vertices_all.append(self.vertices*0.001) - continue - T, _, length = calTransformation(keypoints3d[i, :3], keypoints3d[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, :, :] + 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 kpts3d[nj, -1] < min_conf: + vertices_all.append(self.vertices*0.001) + continue + vertices_all.append(self.vertices*r + kpts3d[nj:nj+1, :3]) + # limb + for nk, (i, j) in enumerate(self.kintree): + if kpts3d[i][-1] < min_conf or kpts3d[j][-1] < min_conf: + vertices_all.append(self.vertices*0.001) + continue + 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) + 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: