def __get_smpl_model(demo_type, smpl_type): smplx_model_path = './extra_data/smpl/SMPLX_NEUTRAL.pkl' smpl_model_path = './extra_data/smpl//basicModel_neutral_lbs_10_207_0_v1.0.0.pkl' if demo_type == 'hand': # use original smpl-x smpl = smplx.create(smplx_model_path, model_type="smplx", batch_size=1, gender='neutral', num_betas=10, use_pca=False, ext='pkl') else: if smpl_type == 'smplx': # use modified smpl-x from body module smpl = SMPLX(smplx_model_path, batch_size=1, num_betas=10, use_pca=False, create_transl=False) else: # use modified smpl from body module assert smpl_type == 'smpl' smpl = SMPL(smpl_model_path, batch_size=1, create_transl=False) return smpl
def main(model_folder, corr_fname, ext='npz', head_color=(0.3, 0.3, 0.6), gender='neutral'): head_idxs = np.load(corr_fname) model = smplx.create(model_folder, model_type='smplx', gender=gender, ext=ext) betas = torch.zeros([1, 10], dtype=torch.float32) expression = torch.zeros([1, 10], dtype=torch.float32) output = model(betas=betas, expression=expression, return_verts=True) vertices = output.vertices.detach().cpu().numpy().squeeze() joints = output.joints.detach().cpu().numpy().squeeze() print('Vertices shape =', vertices.shape) print('Joints shape =', joints.shape) mesh = o3d.geometry.TriangleMesh() mesh.vertices = o3d.utility.Vector3dVector(vertices) mesh.triangles = o3d.utility.Vector3iVector(model.faces) mesh.compute_vertex_normals() colors = np.ones_like(vertices) * [0.3, 0.3, 0.3] colors[head_idxs] = head_color mesh.vertex_colors = o3d.utility.Vector3dVector(colors) o3d.visualization.draw_geometries([mesh])
def get_all_smpl(pkl_data, json_data): gender = json_data['people'][0]['gender_gt'] all_meshes = [] trans = np.array([4, 0, 0]) for i, result in enumerate(pkl_data['all_results']): t = trans + [0, i * 3, 0] betas = torch.Tensor(result['betas']).unsqueeze(0) pose = torch.Tensor(result['body_pose']).unsqueeze(0) transl = torch.Tensor(result['transl']).unsqueeze(0) global_orient = torch.Tensor(result['global_orient']).unsqueeze(0) model = smplx.create('models', model_type='smpl', gender=gender) output = model(betas=betas, body_pose=pose, transl=transl, global_orient=global_orient, return_verts=True) smpl_vertices = output.vertices.detach().cpu().numpy().squeeze() smpl_o3d = o3d.TriangleMesh() smpl_o3d.triangles = o3d.Vector3iVector(model.faces) smpl_o3d.vertices = o3d.Vector3dVector(smpl_vertices) smpl_o3d.compute_vertex_normals() smpl_o3d.translate(t) for idx, key in enumerate(result['loss_dict'].keys()): lbl = '{} {:.2f}'.format(key, float(result['loss_dict'][key])) all_meshes.append(text_3d(lbl, t + [1, idx * 0.2 - 1, 2], direction=(0.01, 0, -1), degree=-90, font_size=150, density=0.2)) all_meshes.append(smpl_o3d) return all_meshes
def __init__(self, trainconfig, lossconfig): for key, val in trainconfig.items(): setattr(self, key, val) for key, val in lossconfig.items(): setattr(self, key, val) if not os.path.exists(self.save_dir): os.mkdir(self.save_dir) if len(self.ckp_dir) > 0: self.resume_training=True ### define model if self.use_cont_rot: n_dim_body=72+3 else: n_dim_body=72 self.model_h_latentD = 256 self.model_h = HumanCVAES2(latentD_g=self.model_h_latentD, latentD_l=self.model_h_latentD, scene_model_ckpt=self.scene_model_ckpt, n_dim_body=n_dim_body, n_dim_scene=self.model_h_latentD) self.optimizer_h = optim.Adam(self.model_h.parameters(), lr=self.init_lr_h) ### body mesh model self.vposer, _ = load_vposer(self.vposer_ckpt_path, vp_model='snapshot') self.body_mesh_model = smplx.create(self.human_model_path, model_type='smplx', gender='neutral', ext='npz', num_pca_comps=12, create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True, batch_size=self.batch_size ) self.smplx_face_idx = np.load(os.path.join(self.human_model_path, 'smplx/SMPLX_NEUTRAL.npz'), allow_pickle=True)['f'].reshape(-1,3) self.smplx_face_idx = torch.tensor(self.smplx_face_idx.astype(np.int64), device=self.device) print('--[INFO] device: '+str(torch.cuda.get_device_name(self.device)) )
def get_smpl(pkl_data, json_data): gender = json_data['people'][0]['gender_gt'] print('Target height {}, weight {}'.format(json_data['people'][0]['height'], json_data['people'][0]['weight'])) betas = torch.Tensor(pkl_data['betas']).unsqueeze(0) pose = torch.Tensor(pkl_data['body_pose']).unsqueeze(0) transl = torch.Tensor(pkl_data['transl']).unsqueeze(0) global_orient = torch.Tensor(pkl_data['global_orient']).unsqueeze(0) model = smplx.create('models', model_type='smpl', gender=gender) output = model(betas=betas, body_pose=pose, transl=transl, global_orient=global_orient, return_verts=True) smpl_vertices = output.vertices.detach().cpu().numpy().squeeze() smpl_joints = output.joints.detach().cpu().numpy().squeeze() output_unposed = model(betas=betas, body_pose=pose * 0, transl=transl, global_orient=global_orient, return_verts=True) smpl_vertices_unposed = output_unposed.vertices.detach().cpu().numpy().squeeze() for i, lbl in enumerate(['Wingspan', 'Height', 'Thickness']): print('Actual', lbl, smpl_vertices_unposed[:, i].max() - smpl_vertices_unposed[:, i].min(), end=' ') print() smpl_trimesh = trimesh.Trimesh(vertices=np.asarray(smpl_vertices_unposed), faces=model.faces) print('Est weight from volume', smpl_trimesh.volume * 1.03 * 1000) # print('Pose embedding', pkl_data['pose_embedding']) # print('Body pose', np.array2string(pkl_data['body_pose'], separator=', ')) smpl_o3d = o3d.TriangleMesh() smpl_o3d.triangles = o3d.Vector3iVector(model.faces) smpl_o3d.vertices = o3d.Vector3dVector(smpl_vertices) smpl_o3d.compute_vertex_normals() # smpl_o3d.paint_uniform_color([0.3, 0.3, 0.3]) smpl_o3d_2 = o3d.TriangleMesh() smpl_o3d_2.triangles = o3d.Vector3iVector(model.faces) smpl_o3d_2.vertices = o3d.Vector3dVector(smpl_vertices + np.array([1.5, 0, 0])) smpl_o3d_2.compute_vertex_normals() smpl_o3d_2.paint_uniform_color([0.7, 0.3, 0.3]) # Visualize SMPL joints - Patrick camera = PerspectiveCamera(rotation=torch.tensor(pkl_data['camera_rotation']).unsqueeze(0), translation=torch.tensor(pkl_data['camera_translation']).unsqueeze(0), center=torch.tensor(pkl_data['camera_center']), focal_length_x=torch.tensor(pkl_data['camera_focal_length_x']), focal_length_y=torch.tensor(pkl_data['camera_focal_length_y'])) gt_pos_3d = camera.inverse_camera_tform(torch.tensor(pkl_data['gt_joints']).unsqueeze(0), 1.8).detach().squeeze(0).cpu().numpy() all_markers = [] for i in range(25): color = cm.jet(i / 25.0)[:3] # smpl_marker = get_o3d_sphere(color=color, pos=smpl_joints[i, :]) # all_markers.append(smpl_marker) pred_marker = get_o3d_sphere(color=color, pos=gt_pos_3d[i, :], radius=0.03) all_markers.append(pred_marker) return smpl_vertices, model.faces, smpl_o3d, smpl_o3d_2, all_markers
def create_smpl_model(**args): assert os.path.exists( args['smpl_model_folder'] ), 'Path {} does not exist in argument smpl_model_folder!'.format( args['smpl_model_folder']) model_params['model_path'] = args['smpl_model_folder'] body_model = smplx.create(model_type='smpl', gender='male', **model_params) body_model = body_model.to(device=DEVICE) return body_model
def __init__(self, fittingconfig, lossconfig): for key, val in fittingconfig.items(): setattr(self, key, val) for key, val in lossconfig.items(): setattr(self, key, val) self.vposer, _ = load_vposer(self.vposer_ckpt_path, vp_model='snapshot') self.body_mesh_model = smplx.create(self.human_model_path, model_type='smplx', gender='neutral', ext='npz', num_pca_comps=12, create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True, batch_size=self.batch_size ) self.vposer.to(self.device) self.body_mesh_model.to(self.device) self.xhr_rec = Variable(torch.randn(1,75).to(self.device), requires_grad=True) self.optimizer = optim.Adam([self.xhr_rec], lr=self.init_lr_h) ## read scene sdf with open(self.scene_sdf_path+'.json') as f: sdf_data = json.load(f) grid_min = np.array(sdf_data['min']) grid_max = np.array(sdf_data['max']) grid_dim = sdf_data['dim'] sdf = np.load(self.scene_sdf_path + '_sdf.npy').reshape(grid_dim, grid_dim, grid_dim) self.s_grid_min_batch = torch.tensor(grid_min, dtype=torch.float32, device=self.device).unsqueeze(0) self.s_grid_max_batch = torch.tensor(grid_max, dtype=torch.float32, device=self.device).unsqueeze(0) self.s_sdf_batch = torch.tensor(sdf, dtype=torch.float32, device=self.device).unsqueeze(0) ## read scene vertices scene_o3d = o3d.io.read_triangle_mesh(self.scene_verts_path) scene_verts = np.asarray(scene_o3d.vertices) self.s_verts_batch = torch.tensor(scene_verts, dtype=torch.float32, device=self.device).unsqueeze(0)
def get_smpl_vertices(betas, expression, smpl_file_name: str = None, texture_file_name: str = None, uv_map_file_name: str = None, body_pose: torch.Tensor = None, return_betas_exps=False) -> np.array: """ Load SMPL model, texture file and uv-map. Set arm angles and convert to mesh. Parameters ---------- betas: np.array(1, 10) Betas for smpl expression: np.array(1, 10) Expression for smpl smpl_file_name : str file name of smpl model (.pkl). texture_file_name : str file name of texture for smpl (.jpg). uv_map_file_name : str file name of uv-map for smpl (.npy). right_arm_angle : float, optional desired right arm angle in radians. The default is 0.. left_arm_angle : float, optional desired left arm angle in radians. The default is 0. body_pose : torch.Tensor[1, 69] Body poses for SMPL Returns ------- vertices : np.array SMPL mesh with texture and desired body pose. """ if smpl_file_name is None: smpl_file_name = "SMPLs/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl" if texture_file_name is None: texture_file_name = "textures/female1.jpg" if uv_map_file_name is None: uv_map_file_name = "textures/smpl_uv_map.npy" model = smplx.create(smpl_file_name, model_type='smpl') # set betas and expression to fixed values betas = torch.tensor(betas).float() expression = torch.tensor(expression).float() output = model(betas=betas, expression=expression, return_verts=True, body_pose=body_pose) vertices = output.vertices.detach().cpu().numpy().squeeze() return vertices
def __init__(self, model_type, body_models_dp, device): super().__init__() self.device = device self.models = dict() for gender in ['female', 'male']: model = smplx.create( body_models_dp, model_type=model_type, gender=gender, batch_size=1, create_transl=False ).to(device=self.device) model.eval() self.models[gender] = model
def main(args, fitting_dir): recording_name = os.path.abspath(fitting_dir).split("/")[-1] fitting_dir = osp.join(fitting_dir, 'results') scene_name = recording_name.split("_")[0] print("scene_name") base_dir = args.base_dir cam2world_dir = osp.join(base_dir, 'cam2world') scene_dir = osp.join(base_dir, 'scenes') female_subjects_ids = [162, 3452, 159, 3403] subject_id = int(recording_name.split('_')[1]) if subject_id in female_subjects_ids: gender = 'female' else: gender = 'male' with open(os.path.join(cam2world_dir, scene_name + '.json'), 'r') as f: trans = np.array(json.load(f)) model = smplx.create(args.model_folder, model_type='smplx', gender=gender, ext='npz', num_pca_comps=args.num_pca_comps, create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True) lis = sorted(os.listdir(fitting_dir)) arr = torch.zeros((len(lis), 103)) for i, img_name in enumerate(lis): try: vec = get_vec(osp.join(fitting_dir, img_name, '000.pkl')) except: arr[i, :] = arr[i - 1, :] continue trans_vec = transform(vec, trans, model) # print(trans_vec.shape) arr[i, :] = trans_vec[0, :] arr = vec_to_cont(arr) np.save("../dataset/PROXD_cleaned/" + recording_name, arr.numpy()) print("recording_name compiled all dataset")
def vis_sequence(cfg,sequence, mv): seq_data = parse_npz(sequence) n_comps = seq_data['n_comps'] gender = seq_data['gender'] T = seq_data.n_frames sbj_mesh = os.path.join(grab_path, '..', seq_data.body.vtemp) sbj_vtemp = np.array(Mesh(filename=sbj_mesh).vertices) sbj_m = smplx.create(model_path=cfg.model_path, model_type='smplx', gender=gender, num_pca_comps=n_comps, v_template=sbj_vtemp, batch_size=T) sbj_parms = params2torch(seq_data.body.params) verts_sbj = to_cpu(sbj_m(**sbj_parms).vertices) obj_mesh = os.path.join(grab_path, '..', seq_data.object.object_mesh) obj_mesh = Mesh(filename=obj_mesh) obj_vtemp = np.array(obj_mesh.vertices) obj_m = ObjectModel(v_template=obj_vtemp, batch_size=T) obj_parms = params2torch(seq_data.object.params) verts_obj = to_cpu(obj_m(**obj_parms).vertices) table_mesh = os.path.join(grab_path, '..', seq_data.table.table_mesh) table_mesh = Mesh(filename=table_mesh) table_vtemp = np.array(table_mesh.vertices) table_m = ObjectModel(v_template=table_vtemp, batch_size=T) table_parms = params2torch(seq_data.table.params) verts_table = to_cpu(table_m(**table_parms).vertices) skip_frame = 4 for frame in range(0,T, skip_frame): o_mesh = Mesh(vertices=verts_obj[frame], faces=obj_mesh.faces, vc=colors['yellow']) o_mesh.set_vertex_colors(vc=colors['red'], vertex_ids=seq_data['contact']['object'][frame] > 0) s_mesh = Mesh(vertices=verts_sbj[frame], faces=sbj_m.faces, vc=colors['pink'], smooth=True) s_mesh.set_vertex_colors(vc=colors['red'], vertex_ids=seq_data['contact']['body'][frame] > 0) t_mesh = Mesh(vertices=verts_table[frame], faces=table_mesh.faces, vc=colors['white']) mv.set_static_meshes([o_mesh, s_mesh, t_mesh])
def get_smpl(joints, axes, amounts, translation=(0, 0, 0)): model = smplx.create('models', model_type='smpl', gender='male') body_pose = torch.zeros([1, 69]) for i in range(len(joints)): pose_index = int(joints[i] * 3 + axes[i]) body_pose[0, pose_index] = axang_mean[ pose_index] + axang_var[pose_index] * (amounts[i] * 2 - 1) output = model(body_pose=torch.Tensor(body_pose), return_verts=True) smpl_vertices = output.vertices.detach().cpu().numpy().squeeze() smpl_o3d = o3d.TriangleMesh() smpl_o3d.triangles = o3d.Vector3iVector(model.faces) smpl_o3d.vertices = o3d.Vector3dVector(smpl_vertices + np.array(translation)) smpl_o3d.compute_vertex_normals() smpl_o3d.paint_uniform_color([amounts[0] / 2 + 0.5, 0.3, 0.3]) return smpl_o3d
def __init__(self, body_models_dp, use_pca=True, num_pca_comps=45, flat_hand_mean=False, device=torch.device('cpu') ): super().__init__() self.models = dict() self.device = device for gender in ['female', 'male']: model = smplx.create( body_models_dp, model_type='smplx', gender=gender, batch_size=1, use_pca=use_pca, num_pca_comps=num_pca_comps, flat_hand_mean=flat_hand_mean, create_transl=False ).to(device=self.device) model.eval() self.models[gender] = model
def __init__(self, testconfig): for key, val in testconfig.items(): setattr(self, key, val) if not os.path.exists(self.ckpt_dir): print('--[ERROR] checkpoints do not exist') sys.exit() #define model if self.use_cont_rot: n_dim_body = 72 + 3 else: n_dim_body = 72 self.model_h_latentD = 256 self.model_h = HumanCVAES2(latentD_g=self.model_h_latentD, latentD_l=self.model_h_latentD, n_dim_body=n_dim_body, n_dim_scene=self.model_h_latentD, test=True) ### body mesh model self.vposer, _ = load_vposer(self.vposer_ckpt_path, vp_model='snapshot') self.body_mesh_model = smplx.create(self.human_model_path, model_type='smplx', gender='neutral', ext='npz', num_pca_comps=12, create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True, batch_size=self.n_samples)
def main(model_folder, corr_fname, ext='npz', hand_color=(0.3, 0.3, 0.6), gender='neutral', hand='right'): with open(corr_fname, 'rb') as f: idxs_data = pickle.load(f) if hand == 'both': hand_idxs = np.concatenate( [idxs_data['left_hand'], idxs_data['right_hand']]) else: hand_idxs = idxs_data[f'{hand}_hand'] model = smplx.create(model_folder, model_type='smplx', gender=gender, ext=ext) betas = torch.zeros([1, 10], dtype=torch.float32) expression = torch.zeros([1, 10], dtype=torch.float32) output = model(betas=betas, expression=expression, return_verts=True) vertices = output.vertices.detach().cpu().numpy().squeeze() joints = output.joints.detach().cpu().numpy().squeeze() print('Vertices shape =', vertices.shape) print('Joints shape =', joints.shape) mesh = o3d.geometry.TriangleMesh() mesh.vertices = o3d.utility.Vector3dVector(vertices) mesh.triangles = o3d.utility.Vector3iVector(model.faces) mesh.compute_vertex_normals() colors = np.ones_like(vertices) * [0.3, 0.3, 0.3] colors[hand_idxs] = hand_color mesh.vertex_colors = o3d.utility.Vector3dVector(colors) o3d.visualization.draw_geometries([mesh])
def get_smplx(model_root="misc/models"): model_params = dict( model_path=model_root, model_type="smplx", gender="neutral", create_body_pose=True, dtype=torch.float32, use_face=False, ) model = smplx.create(**model_params) print(model) # Initialize body parameters. betas = torch.zeros([1, 10], dtype=torch.float32) expression = torch.randn([1, 10], dtype=torch.float32) mean_pose = get_vposer_mean_pose() # Set them. model.expression.data.copy_(expression) model.betas.data.copy_(betas) model.body_pose.data.copy_(mean_pose) return model
def main(model_folder, model_type='smplx', ext='npz', gender='neutral', plot_joints=False, use_face_contour=False): model = smplx.create(model_folder, model_type=model_type, gender=gender, use_face_contour=use_face_contour, ext=ext) print(model) betas = torch.randn([1, 10], dtype=torch.float32) expression = torch.randn([1, 10], dtype=torch.float32) output = model(betas=betas, expression=expression, return_verts=True) vertices = output.vertices.detach().cpu().numpy().squeeze() joints = output.joints.detach().cpu().numpy().squeeze() print('Vertices shape =', vertices.shape) print('Joints shape =', joints.shape) vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8] tri_mesh = trimesh.Trimesh(vertices, model.faces, vertex_colors=vertex_colors) mesh = pyrender.Mesh.from_trimesh(tri_mesh) scene = pyrender.Scene() scene.add(mesh) if plot_joints: sm = trimesh.creation.uv_sphere(radius=0.005) sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0] tfs = np.tile(np.eye(4), (len(joints), 1, 1)) tfs[:, :3, 3] = joints joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs) scene.add(joints_pcl) pyrender.Viewer(scene, use_raymond_lighting=True)
def main(model_folder, model_type='smplx', ext='npz', gender='neutral', plot_joints=False, plotting_module='pyrender', use_face_contour=False): model = smplx.create(model_folder, model_type=model_type, gender=gender, use_face_contour=use_face_contour, ext=ext) print(model) betas = torch.randn([1, 10], dtype=torch.float32) expression = torch.randn([1, 10], dtype=torch.float32) output = model(betas=betas, expression=expression, return_verts=True) vertices = output.vertices.detach().cpu().numpy().squeeze() joints = output.joints.detach().cpu().numpy().squeeze() print('Vertices shape =', vertices.shape) print('Joints shape =', joints.shape) if plotting_module == 'pyrender': import pyrender import trimesh vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8] tri_mesh = trimesh.Trimesh(vertices, model.faces, vertex_colors=vertex_colors) mesh = pyrender.Mesh.from_trimesh(tri_mesh) scene = pyrender.Scene() scene.add(mesh) if plot_joints: sm = trimesh.creation.uv_sphere(radius=0.005) sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0] tfs = np.tile(np.eye(4), (len(joints), 1, 1)) tfs[:, :3, 3] = joints joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs) scene.add(joints_pcl) pyrender.Viewer(scene, use_raymond_lighting=True) elif plotting_module == 'matplotlib': from matplotlib import pyplot as plt from mpl_toolkits.mplot3d import Axes3D from mpl_toolkits.mplot3d.art3d import Poly3DCollection fig = plt.figure() ax = fig.add_subplot(111, projection='3d') mesh = Poly3DCollection(vertices[model.faces], alpha=0.1) face_color = (1.0, 1.0, 0.9) edge_color = (0, 0, 0) mesh.set_edgecolor(edge_color) mesh.set_facecolor(face_color) ax.add_collection3d(mesh) ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], color='r') if plot_joints: ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], alpha=0.1) plt.show() elif plotting_module == 'open3d': import open3d as o3d mesh = o3d.TriangleMesh() mesh.vertices = o3d.Vector3dVector(vertices) mesh.triangles = o3d.Vector3iVector(model.faces) mesh.compute_vertex_normals() mesh.paint_uniform_color([0.3, 0.3, 0.3]) o3d.visualization.draw_geometries([mesh]) else: raise ValueError('Unknown plotting_module: {}'.format(plotting_module))
def main(args): scene_name = os.path.abspath(args.gen_folder).split("/")[-1] outimg_dir = args.outimg_dir if not os.path.exists(outimg_dir): os.makedirs(outimg_dir) ### setup visualization window vis = o3d.visualization.Visualizer() vis.create_window(width=960, height=540, visible=True) render_opt = vis.get_render_option().mesh_show_back_face = True ### put the scene into the environment scene = o3d.io.read_triangle_mesh( osp.join(args.prox_dir, scene_name + '.ply')) vis.add_geometry(scene) vis.update_geometry() # put the body into the environment vposer_ckpt = osp.join(args.model_folder, 'vposer_v1_0') vposer, _ = load_vposer(vposer_ckpt, vp_model='snapshot') model = smplx.create(args.model_folder, model_type='smplx', gender=args.gender, ext='npz', num_pca_comps=args.num_pca_comps, create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True) ## create a corn at the camera location # mesh_corn = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1) # mesh_corn.transform(trans) # vis.add_geometry(mesh_corn) # vis.update_geometry() # print(trans) ## create a corn at the world origin # mesh_corn2 = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1) # vis.add_geometry(mesh_corn2) # vis.update_geometry() # print(trans) cv2.namedWindow("GUI") gen_file_list = glob.glob(os.path.join(args.gen_folder, '*')) body = o3d.geometry.TriangleMesh() vis.add_geometry(body) for idx, gen_file in enumerate(gen_file_list): with open(gen_file, 'rb') as f: param = pickle.load(f) cam_ext = param['cam_ext'][0] cam_int = param['cam_int'][0] body_pose = vposer.decode(torch.tensor(param['body_pose']), output_type='aa').view(1, -1) torch_param = {} for key in param.keys(): if key in ['body_pose', 'camera_rotation', 'camera_translation']: continue else: torch_param[key] = torch.tensor(param[key]) output = model(return_verts=True, body_pose=body_pose, **torch_param) vertices = output.vertices.detach().cpu().numpy().squeeze() body.vertices = o3d.utility.Vector3dVector(vertices) body.triangles = o3d.utility.Vector3iVector(model.faces) body.vertex_normals = o3d.utility.Vector3dVector([]) body.triangle_normals = o3d.utility.Vector3dVector([]) body.compute_vertex_normals() T_mat = np.eye(4) T_mat[1, :] = np.array([0, -1, 0, 0]) T_mat[2, :] = np.array([0, 0, -1, 0]) trans = np.dot(cam_ext, T_mat) body.transform(trans) vis.update_geometry() # while True: # vis.poll_events() # vis.update_renderer() # cv2.imshow("GUI", np.random.random([10,10,3])) # # ctr = vis.get_view_control() # # cam_param = ctr.convert_to_pinhole_camera_parameters() # # print(cam_param.extrinsic) # key = cv2.waitKey(15) # if key == 27: # break ctr = vis.get_view_control() cam_param = ctr.convert_to_pinhole_camera_parameters() cam_param = update_cam(cam_param, trans) ctr.convert_from_pinhole_camera_parameters(cam_param) vis.poll_events() vis.update_renderer() capture_image(vis, outfilename=os.path.join( outimg_dir, 'img_{:06d}_cam1.png'.format(idx))) # vis.run() # capture_image(vis, outfilename=os.path.join(outimg_dir, 'img_{:06d}_cam1.png'.format(idx))) ### setup rendering cam, depth capture, segmentation capture ctr = vis.get_view_control() cam_param = ctr.convert_to_pinhole_camera_parameters() cam_param.extrinsic = trans2_dict[scene_name] ctr.convert_from_pinhole_camera_parameters(cam_param) vis.poll_events() vis.update_renderer() capture_image(vis, outfilename=os.path.join( outimg_dir, 'img_{:06d}_cam2.png'.format(idx)))
def main(args): fitting_dir = args.fitting_dir recording_name = os.path.abspath(fitting_dir).split("/")[-1] fitting_dir = osp.join(fitting_dir, 'results') data_dir = args.data_dir cam2world_dir = osp.join(data_dir, 'cam2world') scene_dir = osp.join(data_dir, 'scenes_semantics') recording_dir = osp.join(data_dir, 'recordings', recording_name) scene_name = os.path.abspath(recording_dir).split("/")[-1].split("_")[0] ## setup the output folder output_folder = os.path.join('/mnt/hdd/PROX', 'snapshot_virtualcam_TNoise0.5') if not os.path.exists(output_folder): os.mkdir(output_folder) ### setup visualization window vis = o3d.visualization.Visualizer() vis.create_window(width=480, height=270, visible=True) render_opt = vis.get_render_option().mesh_show_back_face = True ### put the scene into the visualizer scene = o3d.io.read_triangle_mesh( osp.join(scene_dir, scene_name + '_withlabels.ply')) vis.add_geometry(scene) ## get scene 3D scene bounding box scene_o = o3d.io.read_triangle_mesh( osp.join(scene_dir, scene_name + '.ply')) scene_min = scene_o.get_min_bound() #[x_min, y_min, z_min] scene_max = scene_o.get_max_bound() #[x_max, y_max, z_max] # reduce the scene region furthermore, to avoid cams behind the window shift = 0.7 scene_min = scene_min + shift scene_max = scene_max - shift ### get the real camera config trans_calib = np.eye(4) with open(os.path.join(cam2world_dir, scene_name + '.json'), 'r') as f: trans_calib = np.array(json.load(f)) ## put the body into the environment vposer_ckpt = osp.join(args.model_folder, 'vposer_v1_0') vposer, _ = load_vposer(vposer_ckpt, vp_model='snapshot') model = smplx.create(args.model_folder, model_type='smplx', gender='neutral', ext='npz', num_pca_comps=12, create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=False, create_jaw_pose=False, create_leye_pose=False, create_reye_pose=False, create_transl=True) rec_count = -1 sample_rate = 15 # 0.5second for img_name in sorted(os.listdir(fitting_dir))[::sample_rate]: ## get humam body params filename = osp.join(fitting_dir, img_name, '000.pkl') print('frame: ' + filename) if not os.path.exists(filename): print('file does not exist. Continue') continue with open(osp.join(fitting_dir, img_name, '000.pkl'), 'rb') as f: body_dict = pickle.load(f) if np.sum(np.isnan(body_dict['body_pose'])) > 0: continue rec_count = rec_count + 1 ## save depth, semantics and render cam outname1 = os.path.join(output_folder, recording_name) if not os.path.exists(outname1): os.mkdir(outname1) ######################### then we obtain the virutal cam ################################ ## find world coordinate of the human body in the current frame body_params_W_list, dT = update_globalRT_for_smplx( body_dict, model, vposer, [trans_calib]) body_T_world = body_params_W_list[0]['transl'][0] + dT ## get virtual cams, and transform global_R and global_T to virtual cams new_cammat_ext_list0 = [] new_cammat_ext_list0 = get_new_cams(scene_name, s_min=scene_min, s_max=scene_max, body_T=body_T_world) random.shuffle(new_cammat_ext_list0) new_cammat_ext_list = new_cammat_ext_list0[:30] print('--obtain {:d} cams'.format(len(new_cammat_ext_list))) new_cammat_list = [invert_transform(x) for x in new_cammat_ext_list] body_params_new_list, _ = update_globalRT_for_smplx( body_params_W_list[0], model, vposer, new_cammat_list, delta_T=dT) #### capture depth and seg in new cams for idx_cam, cam_ext in enumerate(new_cammat_ext_list): ### save filename outname = os.path.join( outname1, 'rec_frame{:06d}_cam{:06d}.mat'.format(rec_count, idx_cam)) ## put the render cam to the real cam ctr = vis.get_view_control() cam_param = ctr.convert_to_pinhole_camera_parameters() cam_param = update_render_cam(cam_param, cam_ext) ctr.convert_from_pinhole_camera_parameters(cam_param) vis.poll_events() vis.update_renderer() ## get render cam parameters cam_dict = {} cam_dict['extrinsic'] = cam_param.extrinsic cam_dict['intrinsic'] = cam_param.intrinsic.intrinsic_matrix ## capture depth image depth = np.asarray(vis.capture_depth_float_buffer(do_render=True)) _h = depth.shape[0] _w = depth.shape[1] depth0 = depth depth_canvas, scaling_factor = data_preprocessing(depth, 'depth') ### skip settings when the human body is severely occluded. body_is_occluded = is_body_occluded(body_params_new_list[idx_cam], cam_dict, depth) if body_is_occluded: print( '-- body is occluded or not in the scene at current view.') continue ## capture semantics seg = np.asarray(vis.capture_screen_float_buffer(do_render=True)) verid = np.mean(seg * 255 / 5.0, axis=-1) #.astype(int) seg0 = verid # verid = cv2.resize(verid, (_w//factor, _h//factor)) seg_canvas, _ = data_preprocessing(verid, 'seg') # pdb.set_trace() ## save file to disk ot_dict = {} ot_dict['scaling_factor'] = scaling_factor ot_dict['depth'] = depth_canvas ot_dict['depth0'] = depth0 ot_dict['seg0'] = seg0 ot_dict['seg'] = seg_canvas ot_dict['cam'] = cam_dict ot_dict['body'] = body_params_new_list[idx_cam] sio.savemat(outname, ot_dict) vis.destroy_window()
def main(): description = 'Example script for untangling SMPL self intersections' parser = argparse.ArgumentParser(description=description, prog='Batch SMPL-Untangle') parser.add_argument('--param_fn', type=str, nargs='*', required=True, help='The pickle file with the model parameters') parser.add_argument('--interactive', default=True, type=lambda arg: arg.lower() in ['true', '1'], help='Display the mesh during the optimization' + ' process') parser.add_argument('--delay', type=int, default=50, help='The delay for the animation callback in ms') parser.add_argument('--model_folder', type=str, default='models', help='The path to the LBS model') parser.add_argument('--model_type', type=str, default='smpl', choices=['smpl', 'smplx', 'smplh'], help='The type of model to create') parser.add_argument('--point2plane', default=False, type=lambda arg: arg.lower() in ['true', '1'], help='Use point to distance') parser.add_argument('--optimize_pose', default=True, type=lambda arg: arg.lower() in ['true', '1'], help='Enable optimization over the joint pose') parser.add_argument('--optimize_shape', default=False, type=lambda arg: arg.lower() in ['true', '1'], help='Enable optimization over the shape of the model') parser.add_argument('--sigma', default=0.5, type=float, help='The height of the cone used to calculate the' + ' distance field loss') parser.add_argument('--lr', default=1, type=float, help='The learning rate for SGD') parser.add_argument('--coll_loss_weight', default=1e-4, type=float, help='The weight for the collision loss') parser.add_argument('--pose_reg_weight', default=0, type=float, help='The weight for the pose regularizer') parser.add_argument('--shape_reg_weight', default=0, type=float, help='The weight for the shape regularizer') parser.add_argument('--max_collisions', default=8, type=int, help='The maximum number of bounding box collisions') parser.add_argument('--part_segm_fn', default='', type=str, help='The file with the part segmentation for the' + ' faces of the model') parser.add_argument('--print_timings', default=False, type=lambda arg: arg.lower() in ['true', '1'], help='Print timings for all the operations') args = parser.parse_args() model_folder = args.model_folder model_type = args.model_type param_fn = args.param_fn interactive = args.interactive delay = args.delay point2plane = args.point2plane # optimize_shape = args.optimize_shape # optimize_pose = args.optimize_pose lr = args.lr coll_loss_weight = args.coll_loss_weight pose_reg_weight = args.pose_reg_weight shape_reg_weight = args.shape_reg_weight max_collisions = args.max_collisions sigma = args.sigma part_segm_fn = args.part_segm_fn print_timings = args.print_timings if interactive: import trimesh import pyrender device = torch.device('cuda') batch_size = len(param_fn) params_dict = defaultdict(lambda: []) for idx, fn in enumerate(param_fn): with open(fn, 'rb') as param_file: data = pickle.load(param_file, encoding='latin1') assert 'betas' in data, \ 'No key for shape parameter in provided pickle file' assert 'global_pose' in data, \ 'No key for the global pose in the given pickle file' assert 'pose' in data, \ 'No key for the pose of the joints in the given pickle file' for key, val in data.items(): params_dict[key].append(val) params = {} for key in params_dict: params[key] = np.stack(params_dict[key], axis=0).astype(np.float32) if len(params[key].shape) < 2: params[key] = params[key][np.newaxis] if 'global_pose' in params: params['global_orient'] = params['global_pose'] if 'pose' in params: params['body_pose'] = params['pose'] if part_segm_fn: # Read the part segmentation with open(part_segm_fn, 'rb') as faces_parents_file: data = pickle.load(faces_parents_file, encoding='latin1') faces_segm = data['segm'] faces_parents = data['parents'] # Create the module used to filter invalid collision pairs filter_faces = FilterFaces(faces_segm, faces_parents).to(device=device) # Create the body model body = create(model_folder, batch_size=batch_size, model_type=model_type).to(device=device) body.reset_params(**params) # Clone the given pose to use it as a target for regularization init_pose = body.body_pose.clone().detach() # Create the search tree search_tree = BVH(max_collisions=max_collisions) pen_distance = \ collisions_loss.DistanceFieldPenetrationLoss(sigma=sigma, point2plane=point2plane, vectorized=True) mse_loss = nn.MSELoss(reduction='sum').to(device=device) face_tensor = torch.tensor(body.faces.astype(np.int64), dtype=torch.long, device=device).unsqueeze_(0).repeat([batch_size, 1, 1]) with torch.no_grad(): output = body(get_skin=True) verts = output.vertices bs, nv = verts.shape[:2] bs, nf = face_tensor.shape[:2] faces_idx = face_tensor + \ (torch.arange(bs, dtype=torch.long).to(device) * nv)[:, None, None] optimizer = torch.optim.SGD([body.body_pose], lr=lr) if interactive: # Plot the initial mesh with torch.no_grad(): output = body(get_skin=True) verts = output.vertices np_verts = verts.detach().cpu().numpy() def create_mesh(vertices, faces, color=(0.3, 0.3, 0.3, 1.0), wireframe=False): tri_mesh = trimesh.Trimesh(vertices=vertices, faces=faces) rot = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0]) tri_mesh.apply_transform(rot) material = pyrender.MetallicRoughnessMaterial( metallicFactor=0.0, alphaMode='BLEND', baseColorFactor=color) return pyrender.Mesh.from_trimesh( tri_mesh, material=material) scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 1.0], ambient_light=(1.0, 1.0, 1.0)) for bidx in range(np_verts.shape[0]): curr_verts = np_verts[bidx].copy() body_mesh = create_mesh(curr_verts, body.faces, color=(0.3, 0.3, 0.3, 0.99), wireframe=True) pose = np.eye(4) pose[0, 3] = bidx * 2 scene.add(body_mesh, name='body_mesh_{:03d}'.format(bidx), pose=pose) viewer = pyrender.Viewer(scene, use_raymond_lighting=True, viewport_size=(1200, 800), cull_faces=False, run_in_thread=True) query_names = ['recv_mesh', 'intr_mesh', 'body_mesh'] step = 0 while True: optimizer.zero_grad() if print_timings: start = time.time() if print_timings: torch.cuda.synchronize() output = body(get_skin=True) verts = output.vertices if print_timings: torch.cuda.synchronize() print('Body model forward: {:5f}'.format(time.time() - start)) if print_timings: torch.cuda.synchronize() start = time.time() triangles = verts.view([-1, 3])[faces_idx] if print_timings: torch.cuda.synchronize() print('Triangle indexing: {:5f}'.format(time.time() - start)) with torch.no_grad(): if print_timings: start = time.time() collision_idxs = search_tree(triangles) if print_timings: torch.cuda.synchronize() print('Collision Detection: {:5f}'.format(time.time() - start)) if part_segm_fn: if print_timings: start = time.time() collision_idxs = filter_faces(collision_idxs) if print_timings: torch.cuda.synchronize() print('Collision filtering: {:5f}'.format(time.time() - start)) if print_timings: start = time.time() pen_loss = coll_loss_weight * \ pen_distance(triangles, collision_idxs) if print_timings: torch.cuda.synchronize() print('Penetration loss: {:5f}'.format(time.time() - start)) shape_reg_loss = torch.tensor(0, device=device, dtype=torch.float32) if shape_reg_weight > 0: shape_reg_loss = shape_reg_weight * torch.sum(body.betas ** 2) pose_reg_loss = torch.tensor(0, device=device, dtype=torch.float32) if pose_reg_weight > 0: pose_reg_loss = pose_reg_weight * \ mse_loss(body.pose, init_pose) loss = pen_loss + pose_reg_loss + shape_reg_loss np_loss = loss.detach().cpu().squeeze().tolist() if type(np_loss) != list: np_loss = [np_loss] msg = '{:.5f} ' * len(np_loss) print('Loss per model:', msg.format(*np_loss)) if print_timings: start = time.time() loss.backward(torch.ones_like(loss)) if print_timings: torch.cuda.synchronize() print('Backward pass: {:5f}'.format(time.time() - start)) if interactive: with torch.no_grad(): output = body(get_skin=True) verts = output.vertices np_verts = verts.detach().cpu().numpy() np_collision_idxs = collision_idxs.detach().cpu().numpy() np_receivers = np_collision_idxs[:, :, 0] np_intruders = np_collision_idxs[:, :, 1] viewer.render_lock.acquire() for node in scene.get_nodes(): if node.name is None: continue if any([query in node.name for query in query_names]): scene.remove_node(node) for bidx in range(batch_size): recv_faces_idxs = np_receivers[bidx][np_receivers[bidx] >= 0] intr_faces_idxs = np_intruders[bidx][np_intruders[bidx] >= 0] recv_faces = body.faces[recv_faces_idxs] intr_faces = body.faces[intr_faces_idxs] curr_verts = np_verts[bidx].copy() body_mesh = create_mesh(curr_verts, body.faces, color=(0.3, 0.3, 0.3, 0.99), wireframe=True) pose = np.eye(4) pose[0, 3] = bidx * 2 scene.add(body_mesh, name='body_mesh_{:03d}'.format(bidx), pose=pose) if len(intr_faces) > 0: intr_mesh = create_mesh(curr_verts, intr_faces, color=(0.9, 0.0, 0.0, 1.0)) scene.add(intr_mesh, name='intr_mesh_{:03d}'.format(bidx), pose=pose) if len(recv_faces) > 0: recv_mesh = create_mesh(curr_verts, recv_faces, color=(0.0, 0.9, 0.0, 1.0)) scene.add(recv_mesh, name='recv_mesh_{:03d}'.format(bidx), pose=pose) viewer.render_lock.release() if not viewer.is_active: break time.sleep(delay / 1000) optimizer.step() step += 1
def main(**args): data_folder = args.get('recording_dir') recording_name = osp.basename(args.get('recording_dir')) scene_name = recording_name.split("_")[0] base_dir = os.path.abspath( osp.join(args.get('recording_dir'), os.pardir, os.pardir)) keyp_dir = osp.join(base_dir, 'keypoints') keyp_folder = osp.join(keyp_dir, recording_name) cam2world_dir = osp.join(base_dir, 'cam2world') scene_dir = osp.join(base_dir, 'scenes') calib_dir = osp.join(base_dir, 'calibration') sdf_dir = osp.join(base_dir, 'sdf') body_segments_dir = osp.join(base_dir, 'body_segments') output_folder = args.get('output_folder') output_folder = osp.expandvars(output_folder) output_folder = osp.join(output_folder, recording_name) if not osp.exists(output_folder): os.makedirs(output_folder) # Store the arguments for the current experiment conf_fn = osp.join(output_folder, 'conf.yaml') with open(conf_fn, 'w') as conf_file: yaml.dump(args, conf_file) #remove 'output_folder' from args list args.pop('output_folder') result_folder = args.pop('result_folder', 'results') result_folder = osp.join(output_folder, result_folder) if not osp.exists(result_folder): os.makedirs(result_folder) mesh_folder = args.pop('mesh_folder', 'meshes') mesh_folder = osp.join(output_folder, mesh_folder) if not osp.exists(mesh_folder): os.makedirs(mesh_folder) out_img_folder = osp.join(output_folder, 'images') if not osp.exists(out_img_folder): os.makedirs(out_img_folder) body_scene_rendering_dir = os.path.join(output_folder, 'renderings') if not osp.exists(body_scene_rendering_dir): os.mkdir(body_scene_rendering_dir) float_dtype = args['float_dtype'] if float_dtype == 'float64': dtype = torch.float64 elif float_dtype == 'float32': dtype = torch.float64 else: print('Unknown float type {}, exiting!'.format(float_dtype)) sys.exit(-1) use_cuda = args.get('use_cuda', True) if use_cuda and not torch.cuda.is_available(): print('CUDA is not available, exiting!') sys.exit(-1) img_folder = args.pop('img_folder', 'Color') dataset_obj = create_dataset(img_folder=img_folder, data_folder=data_folder, keyp_folder=keyp_folder, calib_dir=calib_dir, **args) start = time.time() input_gender = args.pop('gender', 'neutral') gender_lbl_type = args.pop('gender_lbl_type', 'none') max_persons = args.pop('max_persons', -1) float_dtype = args.get('float_dtype', 'float32') if float_dtype == 'float64': dtype = torch.float64 elif float_dtype == 'float32': dtype = torch.float32 else: raise ValueError('Unknown float type {}, exiting!'.format(float_dtype)) joint_mapper = JointMapper(dataset_obj.get_model2data()) model_params = dict(model_path=args.get('model_folder'), joint_mapper=joint_mapper, create_global_orient=True, create_body_pose=not args.get('use_vposer'), create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True, dtype=dtype, **args) male_model = smplx.create(gender='male', **model_params) # SMPL-H has no gender-neutral model if args.get('model_type') != 'smplh': neutral_model = smplx.create(gender='neutral', **model_params) female_model = smplx.create(gender='female', **model_params) # Create the camera object camera_center = None \ if args.get('camera_center_x') is None or args.get('camera_center_y') is None \ else torch.tensor([args.get('camera_center_x'), args.get('camera_center_y')], dtype=dtype).view(-1, 2) camera = create_camera(focal_length_x=args.get('focal_length_x'), focal_length_y=args.get('focal_length_y'), center=camera_center, batch_size=args.get('batch_size'), dtype=dtype) if hasattr(camera, 'rotation'): camera.rotation.requires_grad = False use_hands = args.get('use_hands', True) use_face = args.get('use_face', True) body_pose_prior = create_prior(prior_type=args.get('body_prior_type'), dtype=dtype, **args) jaw_prior, expr_prior = None, None if use_face: jaw_prior = create_prior(prior_type=args.get('jaw_prior_type'), dtype=dtype, **args) expr_prior = create_prior(prior_type=args.get('expr_prior_type', 'l2'), dtype=dtype, **args) left_hand_prior, right_hand_prior = None, None if use_hands: lhand_args = args.copy() lhand_args['num_gaussians'] = args.get('num_pca_comps') left_hand_prior = create_prior( prior_type=args.get('left_hand_prior_type'), dtype=dtype, use_left_hand=True, **lhand_args) rhand_args = args.copy() rhand_args['num_gaussians'] = args.get('num_pca_comps') right_hand_prior = create_prior( prior_type=args.get('right_hand_prior_type'), dtype=dtype, use_right_hand=True, **rhand_args) shape_prior = create_prior(prior_type=args.get('shape_prior_type', 'l2'), dtype=dtype, **args) angle_prior = create_prior(prior_type='angle', dtype=dtype) if use_cuda and torch.cuda.is_available(): device = torch.device('cuda') camera = camera.to(device=device) female_model = female_model.to(device=device) male_model = male_model.to(device=device) if args.get('model_type') != 'smplh': neutral_model = neutral_model.to(device=device) body_pose_prior = body_pose_prior.to(device=device) angle_prior = angle_prior.to(device=device) shape_prior = shape_prior.to(device=device) if use_face: expr_prior = expr_prior.to(device=device) jaw_prior = jaw_prior.to(device=device) if use_hands: left_hand_prior = left_hand_prior.to(device=device) right_hand_prior = right_hand_prior.to(device=device) else: device = torch.device('cpu') # A weight for every joint of the model joint_weights = dataset_obj.get_joint_weights().to(device=device, dtype=dtype) # Add a fake batch dimension for broadcasting joint_weights.unsqueeze_(dim=0) for idx, data in enumerate(dataset_obj): img = data['img'] fn = data['fn'] keypoints = data['keypoints'] depth_im = data['depth_im'] mask = data['mask'] init_trans = None if data['init_trans'] is None else torch.tensor( data['init_trans'], dtype=dtype).view(-1, 3) scan = data['scan_dict'] print('Processing: {}'.format(data['img_path'])) curr_result_folder = osp.join(result_folder, fn) if not osp.exists(curr_result_folder): os.makedirs(curr_result_folder) curr_mesh_folder = osp.join(mesh_folder, fn) if not osp.exists(curr_mesh_folder): os.makedirs(curr_mesh_folder) #TODO: SMPLifyD and PROX won't work for multiple persons for person_id in range(keypoints.shape[0]): if person_id >= max_persons and max_persons > 0: continue curr_result_fn = osp.join(curr_result_folder, '{:03d}.pkl'.format(person_id)) curr_mesh_fn = osp.join(curr_mesh_folder, '{:03d}.ply'.format(person_id)) curr_body_scene_rendering_fn = osp.join(body_scene_rendering_dir, fn + '.png') curr_img_folder = osp.join(output_folder, 'images', fn, '{:03d}'.format(person_id)) if not osp.exists(curr_img_folder): os.makedirs(curr_img_folder) if gender_lbl_type != 'none': if gender_lbl_type == 'pd' and 'gender_pd' in data: gender = data['gender_pd'][person_id] if gender_lbl_type == 'gt' and 'gender_gt' in data: gender = data['gender_gt'][person_id] else: gender = input_gender if gender == 'neutral': body_model = neutral_model elif gender == 'female': body_model = female_model elif gender == 'male': body_model = male_model out_img_fn = osp.join(curr_img_folder, 'output.png') fit_single_frame( img, keypoints[[person_id]], init_trans, scan, cam2world_dir=cam2world_dir, scene_dir=scene_dir, sdf_dir=sdf_dir, body_segments_dir=body_segments_dir, scene_name=scene_name, body_model=body_model, camera=camera, joint_weights=joint_weights, dtype=dtype, output_folder=output_folder, result_folder=curr_result_folder, out_img_fn=out_img_fn, result_fn=curr_result_fn, mesh_fn=curr_mesh_fn, body_scene_rendering_fn=curr_body_scene_rendering_fn, shape_prior=shape_prior, expr_prior=expr_prior, body_pose_prior=body_pose_prior, left_hand_prior=left_hand_prior, right_hand_prior=right_hand_prior, jaw_prior=jaw_prior, angle_prior=angle_prior, **args) elapsed = time.time() - start time_msg = time.strftime('%H hours, %M minutes, %S seconds', time.gmtime(elapsed)) print('Processing the data took: {}'.format(time_msg))
def main(args): fitting_dir = args.fitting_dir recording_name = os.path.abspath(fitting_dir).split("/")[-1] fitting_dir = osp.join(fitting_dir, 'results') scene_name = recording_name.split("_")[0] base_dir = args.base_dir cam2world_dir = osp.join(base_dir, 'cam2world') scene_dir = osp.join(base_dir, 'scenes') recording_dir = osp.join(base_dir, 'recordings', recording_name) color_dir = os.path.join(recording_dir, 'Color') female_subjects_ids = [162, 3452, 159, 3403] subject_id = int(recording_name.split('_')[1]) if subject_id in female_subjects_ids: gender = 'female' else: gender = 'male' cv2.namedWindow('frame', cv2.WINDOW_NORMAL) vis = o3d.Visualizer() vis.create_window() scene = o3d.io.read_triangle_mesh(osp.join(scene_dir, scene_name + '.ply')) with open(os.path.join(cam2world_dir, scene_name + '.json'), 'r') as f: trans = np.array(json.load(f)) vis.add_geometry(scene) model = smplx.create(args.model_folder, model_type='smplx', gender=gender, ext='npz', num_pca_comps=args.num_pca_comps, create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True ) count = 0 for img_name in sorted(os.listdir(fitting_dir))[args.start::args.step]: print('viz frame {}'.format(img_name)) with open(osp.join(fitting_dir, img_name, '000.pkl'), 'rb') as f: param = pickle.load(f) torch_param = {} for key in param.keys(): if key in ['pose_embedding', 'camera_rotation', 'camera_translation']: continue else: torch_param[key] = torch.tensor(param[key]) output = model(return_verts=True, **torch_param) vertices = output.vertices.detach().cpu().numpy().squeeze() if count == 0: body = o3d.TriangleMesh() vis.add_geometry(body) body.vertices = o3d.Vector3dVector(vertices) body.triangles = o3d.Vector3iVector(model.faces) body.vertex_normals = o3d.Vector3dVector([]) body.triangle_normals = o3d.Vector3dVector([]) body.compute_vertex_normals() body.transform(trans) color_img = cv2.imread(os.path.join(color_dir, img_name + '.jpg')) color_img = cv2.flip(color_img, 1) vis.update_geometry() while True: cv2.imshow('frame', color_img) vis.poll_events() vis.update_renderer() key = cv2.waitKey(30) if key == 27: break count += 1
# coordinate masking for error calculation ih26m_joint_from_mesh = ih26m_joint_from_mesh[np.tile( ih26m_joint_valid == 1, (1, 3))].reshape(-1, 3) ih26m_joint_cam = ih26m_joint_cam[np.tile(ih26m_joint_valid == 1, (1, 3))].reshape(-1, 3) error = np.sqrt(np.sum((ih26m_joint_from_mesh - ih26m_joint_cam)**2, 1)).mean() return error # mano layer smplx_path = 'SMPLX_PATH' mano_layer = { 'right': smplx.create(smplx_path, 'mano', use_pca=False, is_rhand=True), 'left': smplx.create(smplx_path, 'mano', use_pca=False, is_rhand=False) } ih26m_joint_regressor = np.load('J_regressor_mano_ih26m.npy') # fix MANO shapedirs of the left hand bug (https://github.com/vchoutas/smplx/issues/48) if torch.sum( torch.abs(mano_layer['left'].shapedirs[:, 0, :] - mano_layer['right'].shapedirs[:, 0, :])) < 1: print('Fix shapedirs bug of MANO') mano_layer['left'].shapedirs[:, 0, :] *= -1 root_path = '../../data/InterHand2.6M/data/' img_root_path = osp.join(root_path, 'images') annot_root_path = osp.join(root_path, 'annotations') subset = 'all'
def train(): parser = config_parser() args = parser.parse_args() device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") args.default_device = device if args.model_type not in [ "nerf", "smpl_nerf", "append_to_nerf", "smpl", "warp", 'vertex_sphere', "smpl_estimator", "original_nerf", 'dummy_dynamic', 'image_wise_dynamic', "append_vertex_locations_to_nerf", 'append_smpl_params' ]: raise Exception("The model type ", args.model_type, " does not exist.") transform = transforms.Compose([ NormalizeRGB(), CoarseSampling(args.near, args.far, args.number_coarse_samples), ToTensor() ]) train_dir = os.path.join(args.dataset_dir, 'train') val_dir = os.path.join(args.dataset_dir, 'val') if args.model_type == "nerf": train_data = RaysFromImagesDataset( train_dir, os.path.join(train_dir, 'transforms.json'), transform) val_data = RaysFromImagesDataset( val_dir, os.path.join(val_dir, 'transforms.json'), transform) elif args.model_type == "smpl" or args.model_type == "warp": train_data = SmplDataset(train_dir, os.path.join(train_dir, 'transforms.json'), args, transform=NormalizeRGB()) val_data = SmplDataset(val_dir, os.path.join(val_dir, 'transforms.json'), args, transform=NormalizeRGB()) elif args.model_type == "smpl_nerf" or args.model_type == "append_to_nerf" or args.model_type == "append_smpl_params": train_data = SmplNerfDataset( train_dir, os.path.join(train_dir, 'transforms.json'), transform) val_data = SmplNerfDataset(val_dir, os.path.join(val_dir, 'transforms.json'), transform) elif args.model_type == "vertex_sphere": train_data = VertexSphereDataset( train_dir, os.path.join(train_dir, 'transforms.json'), args) val_data = VertexSphereDataset( val_dir, os.path.join(val_dir, 'transforms.json'), args) elif args.model_type == "smpl_estimator": transform = NormalizeRGBImage() train_data = SmplEstimatorDataset( train_dir, os.path.join(train_dir, 'transforms.json'), args.vertex_sphere_radius, transform) val_data = SmplEstimatorDataset( val_dir, os.path.join(val_dir, 'transforms.json'), args.vertex_sphere_radius, transform) elif args.model_type == "original_nerf": train_data = OriginalNerfDataset( args.dataset_dir, os.path.join(args.dataset_dir, 'transforms_train.json'), transform) val_data = OriginalNerfDataset( args.dataset_dir, os.path.join(args.dataset_dir, 'transforms_val.json'), transform) elif args.model_type == "dummy_dynamic": train_data = DummyDynamicDataset( train_dir, os.path.join(train_dir, 'transforms.json'), transform) val_data = DummyDynamicDataset( val_dir, os.path.join(val_dir, 'transforms.json'), transform) elif args.model_type == "append_vertex_locations_to_nerf": train_data = DummyDynamicDataset( train_dir, os.path.join(train_dir, 'transforms.json'), transform) val_data = DummyDynamicDataset( val_dir, os.path.join(val_dir, 'transforms.json'), transform) elif args.model_type == 'image_wise_dynamic': canonical_pose1 = torch.zeros(38).view(1, -1) canonical_pose2 = torch.zeros(2).view(1, -1) canonical_pose3 = torch.zeros(27).view(1, -1) arm_angle_l = torch.tensor([np.deg2rad(10)]).float().view(1, -1) arm_angle_r = torch.tensor([np.deg2rad(10)]).float().view(1, -1) smpl_estimator = DummyImageWiseEstimator(canonical_pose1, canonical_pose2, canonical_pose3, arm_angle_l, arm_angle_r, torch.zeros(10).view(1, -1), torch.zeros(69).view(1, -1)) train_data = ImageWiseDataset( train_dir, os.path.join(train_dir, 'transforms.json'), smpl_estimator, transform, args) val_data = ImageWiseDataset(val_dir, os.path.join(val_dir, 'transforms.json'), smpl_estimator, transform, args) train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batchsize, shuffle=True, num_workers=0) val_loader = torch.utils.data.DataLoader(val_data, batch_size=args.batchsize_val, shuffle=False, num_workers=0) position_encoder = PositionalEncoder(args.number_frequencies_postitional, args.use_identity_positional) direction_encoder = PositionalEncoder(args.number_frequencies_directional, args.use_identity_directional) model_coarse = RenderRayNet(args.netdepth, args.netwidth, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, skips=args.skips) model_fine = RenderRayNet(args.netdepth_fine, args.netwidth_fine, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, skips=args.skips_fine) if args.model_type == "smpl_nerf": human_pose_encoder = PositionalEncoder(args.number_frequencies_pose, args.use_identity_pose) positions_dim = position_encoder.output_dim if args.human_pose_encoding else 1 human_pose_dim = human_pose_encoder.output_dim if args.human_pose_encoding else 1 model_warp_field = WarpFieldNet(args.netdepth_warp, args.netwidth_warp, positions_dim * 3, human_pose_dim * 2) solver = SmplNerfSolver(model_coarse, model_fine, model_warp_field, position_encoder, direction_encoder, human_pose_encoder, train_data.canonical_smpl, args, torch.optim.Adam, torch.nn.MSELoss()) solver.train(train_loader, val_loader, train_data.h, train_data.w) save_run(solver.writer.log_dir, [model_coarse, model_fine, model_warp_field], ['model_coarse.pt', 'model_fine.pt', 'model_warp_field.pt'], parser) elif args.model_type == 'smpl': solver = SmplSolver(model_coarse, model_fine, position_encoder, direction_encoder, args, torch.optim.Adam, torch.nn.MSELoss()) solver.train(train_loader, val_loader, train_data.h, train_data.w, parser) save_run(solver.writer.log_dir, [model_coarse, model_fine], ['model_coarse.pt', 'model_fine.pt'], parser) elif args.model_type == 'nerf' or args.model_type == "original_nerf": solver = NerfSolver(model_coarse, model_fine, position_encoder, direction_encoder, args, torch.optim.Adam, torch.nn.MSELoss()) solver.train(train_loader, val_loader, train_data.h, train_data.w, parser) save_run(solver.writer.log_dir, [model_coarse, model_fine], ['model_coarse.pt', 'model_fine.pt'], parser) elif args.model_type == 'warp': human_pose_encoder = PositionalEncoder(args.number_frequencies_pose, args.use_identity_pose) positions_dim = position_encoder.output_dim if args.human_pose_encoding else 1 human_pose_dim = human_pose_encoder.output_dim if args.human_pose_encoding else 1 model_warp_field = WarpFieldNet(args.netdepth_warp, args.netwidth_warp, positions_dim * 3, human_pose_dim * 2) human_pose_encoder = PositionalEncoder(args.number_frequencies_pose, args.use_identity_pose) solver = WarpSolver(model_warp_field, position_encoder, direction_encoder, human_pose_encoder, args) solver.train(train_loader, val_loader, train_data.h, train_data.w) save_run(solver.writer.log_dir, [model_warp_field], ['model_warp_field.pt'], parser) elif args.model_type == 'append_smpl_params': human_pose_encoder = PositionalEncoder(args.number_frequencies_pose, args.use_identity_pose) human_pose_dim = human_pose_encoder.output_dim if args.human_pose_encoding else 1 model_coarse = RenderRayNet( args.netdepth, args.netwidth, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, human_pose_dim * 69, skips=args.skips, use_directional_input=args.use_directional_input) model_fine = RenderRayNet( args.netdepth_fine, args.netwidth_fine, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, human_pose_dim * 69, skips=args.skips_fine, use_directional_input=args.use_directional_input) if args.load_run is not None: model_coarse.load_state_dict( torch.load(os.path.join(args.load_run, 'model_coarse.pt'), map_location=torch.device(device))) model_fine.load_state_dict( torch.load(os.path.join(args.load_run, 'model_fine.pt'), map_location=torch.device(device))) print("Models loaded from ", args.load_run) if args.siren: model_coarse = SirenRenderRayNet( args.netdepth, args.netwidth, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, human_pose_dim * 69, skips=args.skips, use_directional_input=args.use_directional_input) model_fine = SirenRenderRayNet( args.netdepth_fine, args.netwidth_fine, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, human_pose_dim * 69, skips=args.skips_fine, use_directional_input=args.use_directional_input) solver = AppendSmplParamsSolver(model_coarse, model_fine, position_encoder, direction_encoder, human_pose_encoder, args, torch.optim.Adam, torch.nn.MSELoss()) solver.train(train_loader, val_loader, train_data.h, train_data.w, parser) save_run(solver.writer.log_dir, [model_coarse, model_fine], ['model_coarse.pt', 'model_fine.pt'], parser) model_dependent = [human_pose_encoder, human_pose_dim] inference_gif(solver.writer.log_dir, args.model_type, args, train_data, val_data, position_encoder, direction_encoder, model_coarse, model_fine, model_dependent) elif args.model_type == 'append_to_nerf': human_pose_encoder = PositionalEncoder(args.number_frequencies_pose, args.use_identity_pose) human_pose_dim = human_pose_encoder.output_dim if args.human_pose_encoding else 1 model_coarse = RenderRayNet( args.netdepth, args.netwidth, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, human_pose_dim * 2, skips=args.skips, use_directional_input=args.use_directional_input) model_fine = RenderRayNet( args.netdepth_fine, args.netwidth_fine, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, human_pose_dim * 2, skips=args.skips_fine, use_directional_input=args.use_directional_input) solver = AppendToNerfSolver(model_coarse, model_fine, position_encoder, direction_encoder, human_pose_encoder, args, torch.optim.Adam, torch.nn.MSELoss()) solver.train(train_loader, val_loader, train_data.h, train_data.w, parser) save_run(solver.writer.log_dir, [model_coarse, model_fine], ['model_coarse.pt', 'model_fine.pt'], parser) model_dependent = [human_pose_encoder, human_pose_dim] inference_gif(solver.writer.log_dir, args.model_type, args, train_data, val_data, position_encoder, direction_encoder, model_coarse, model_fine, model_dependent) elif args.model_type == 'append_vertex_locations_to_nerf': model_coarse = AppendVerticesNet(args.netdepth, args.netwidth, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, 6890, additional_input_layers=1, skips=args.skips) model_fine = AppendVerticesNet(args.netdepth_fine, args.netwidth_fine, position_encoder.output_dim * 3, direction_encoder.output_dim * 3, 6890, additional_input_layers=1, skips=args.skips_fine) smpl_estimator = DummySmplEstimatorModel(train_data.goal_poses, train_data.betas) smpl_file_name = "SMPLs/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl" smpl_model = smplx.create(smpl_file_name, model_type='smpl') smpl_model.batchsize = args.batchsize solver = AppendVerticesSolver(model_coarse, model_fine, smpl_estimator, smpl_model, position_encoder, direction_encoder, args, torch.optim.Adam, torch.nn.MSELoss()) solver.train(train_loader, val_loader, train_data.h, train_data.w) save_run(solver.writer.log_dir, [model_coarse, model_fine], ['model_coarse.pt', 'model_fine.pt'], parser) elif args.model_type == 'vertex_sphere': solver = VertexSphereSolver(model_coarse, model_fine, position_encoder, direction_encoder, args, torch.optim.Adam, torch.nn.MSELoss()) solver.train(train_loader, val_loader, train_data.h, train_data.w) save_run(solver.writer.log_dir, [model_coarse, model_fine], ['model_coarse.pt', 'model_fine.pt'], parser) elif args.model_type == 'smpl_estimator': model = SmplEstimator(human_size=len(args.human_joints)) solver = SmplEstimatorSolver(model, args, torch.optim.Adam, torch.nn.MSELoss()) solver.train(train_loader, val_loader) save_run(solver.writer.log_dir, [model], ['model_smpl_estimator.pt'], parser) elif args.model_type == "dummy_dynamic": smpl_file_name = "SMPLs/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl" smpl_model = smplx.create(smpl_file_name, model_type='smpl') smpl_model.batchsize = args.batchsize smpl_estimator = DummySmplEstimatorModel(train_data.goal_poses, train_data.betas) solver = DynamicSolver(model_fine, model_coarse, smpl_estimator, smpl_model, position_encoder, direction_encoder, args) solver.train(train_loader, val_loader, train_data.h, train_data.w) save_run(solver.writer.log_dir, [model_coarse, model_fine, smpl_estimator], ['model_coarse.pt', 'model_fine.pt', 'smpl_estimator.pt'], parser) elif args.model_type == "image_wise_dynamic": if args.load_coarse_model != None: print("Load model..") model_coarse.load_state_dict( torch.load(args.load_coarse_model, map_location=torch.device(device))) for params in model_coarse.parameters(): params.requires_grad = False model_coarse.eval() train_loader = torch.utils.data.DataLoader(train_data, batch_size=1, shuffle=True, num_workers=0) val_loader = torch.utils.data.DataLoader(val_data, batch_size=1, shuffle=False, num_workers=0) smpl_file_name = "SMPLs/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl" smpl_model = smplx.create(smpl_file_name, model_type='smpl') smpl_model.batchsize = args.batchsize solver = ImageWiseSolver(model_coarse, model_fine, smpl_estimator, smpl_model, position_encoder, direction_encoder, args) solver.train(train_loader, val_loader, train_data.h, train_data.w) save_run(solver.writer.log_dir, [model_coarse, model_fine, smpl_estimator], ['model_coarse.pt', 'model_fine.pt', 'smpl_estimator.pt'], parser)
model_params = dict(model_path=args.get('model_folder'), # joint_mapper=joint_mapper, create_global_orient=True, create_body_pose=not args.get('use_vposer'), create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=False, dtype=dtype, **args) model = smplx.create(**model_params) model = model.to(device=device) batch_size = args.get('batch_size', 1) use_vposer = args.get('use_vposer', True) vposer, pose_embedding = [None, ] * 2 vposer_ckpt = args.get('vposer_ckpt', '') if use_vposer: pose_embedding = torch.zeros([batch_size, 32], dtype=dtype, device=device, requires_grad=True) vposer_ckpt = osp.expandvars(vposer_ckpt) vposer, _ = load_vposer(vposer_ckpt, vp_model='snapshot') vposer = vposer.to(device=device) vposer.eval()
def optimize_visulize(): # read scene mesh, scene sdf scene, cur_scene_verts, s_grid_min_batch, s_grid_max_batch, s_sdf_batch = read_mesh_sdf( args.dataset_path, args.dataset, args.scene_name) smplx_model = smplx.create(args.smplx_model_path, model_type='smplx', gender='neutral', ext='npz', num_pca_comps=12, create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True, batch_size=1).to(device) print('[INFO] smplx model loaded.') vposer_model, _ = load_vposer(args.vposer_model_path, vp_model='snapshot') vposer_model = vposer_model.to(device) print('[INFO] vposer model loaded') ##################### load optimization results ################## shift_list = np.load('{}/{}/shift_list.npy'.format( args.optimize_result_dir, args.scene_name)) rot_angle_list_1 = np.load('{}/{}/rot_angle_list_1.npy'.format( args.optimize_result_dir, args.scene_name)) if args.optimize: body_params_opt_list_s1 = np.load( '{}/{}/body_params_opt_list_s1.npy'.format( args.optimize_result_dir, args.scene_name)) body_params_opt_list_s2 = np.load( '{}/{}/body_params_opt_list_s2.npy'.format( args.optimize_result_dir, args.scene_name)) body_verts_sample_list = np.load('{}/{}/body_verts_sample_list.npy'.format( args.optimize_result_dir, args.scene_name)) n_sample = len(body_verts_sample_list) ########################## evaluation (contact/collision score) ######################### loss_non_collision_sample, loss_contact_sample = 0, 0 loss_non_collision_opt_s1, loss_contact_opt_s1 = 0, 0 loss_non_collision_opt_s2, loss_contact_opt_s2 = 0, 0 body_params_prox_list_s1, body_params_prox_list_s2 = [], [] body_verts_opt_prox_s2_list = [] for cnt in tqdm(range(0, n_sample)): body_verts_sample = body_verts_sample_list[cnt] # [10475, 3] # smplx params --> body mesh body_params_opt_s1 = torch.from_numpy( body_params_opt_list_s1[cnt]).float().unsqueeze(0).to( device) # [1,75] body_params_opt_s1 = convert_to_3D_rot( body_params_opt_s1) # tensor, [bs=1, 72] body_pose_joint = vposer_model.decode(body_params_opt_s1[:, 16:48], output_type='aa').view( 1, -1) # [1, 63] body_verts_opt_s1 = gen_body_mesh(body_params_opt_s1, body_pose_joint, smplx_model)[0] # [n_body_vert, 3] body_verts_opt_s1 = body_verts_opt_s1.detach().cpu().numpy() body_params_opt_s2 = torch.from_numpy( body_params_opt_list_s2[cnt]).float().unsqueeze(0).to(device) body_params_opt_s2 = convert_to_3D_rot( body_params_opt_s2) # tensor, [bs=1, 72] body_pose_joint = vposer_model.decode(body_params_opt_s2[:, 16:48], output_type='aa').view(1, -1) body_verts_opt_s2 = gen_body_mesh(body_params_opt_s2, body_pose_joint, smplx_model)[0] body_verts_opt_s2 = body_verts_opt_s2.detach().cpu().numpy() ####################### transfrom local body verts to prox coodinate system #################### # generated body verts from cvae, before optimization body_verts_sample_prox = np.zeros( body_verts_sample.shape) # [10475, 3] temp = body_verts_sample - shift_list[cnt] body_verts_sample_prox[:, 0] = temp[:, 0] * math.cos(-rot_angle_list_1[cnt]) - \ temp[:, 1] * math.sin(-rot_angle_list_1[cnt]) body_verts_sample_prox[:, 1] = temp[:, 0] * math.sin(-rot_angle_list_1[cnt]) + \ temp[:, 1] * math.cos(-rot_angle_list_1[cnt]) body_verts_sample_prox[:, 2] = temp[:, 2] ######### optimized body verts trans_matrix_1 = np.array([[ math.cos(-rot_angle_list_1[cnt]), -math.sin(-rot_angle_list_1[cnt]), 0, 0 ], [ math.sin(-rot_angle_list_1[cnt]), math.cos(-rot_angle_list_1[cnt]), 0, 0 ], [0, 0, 1, 0], [0, 0, 0, 1]]) trans_matrix_2 = np.array([[1, 0, 0, -shift_list[cnt][0]], [0, 1, 0, -shift_list[cnt][1]], [0, 0, 1, -shift_list[cnt][2]], [0, 0, 0, 1]]) ### stage 1: simple optimization results body_verts_opt_prox_s1 = np.zeros( body_verts_opt_s1.shape) # [10475, 3] temp = body_verts_opt_s1 - shift_list[cnt] body_verts_opt_prox_s1[:, 0] = temp[:, 0] * math.cos(-rot_angle_list_1[cnt]) - \ temp[:, 1] * math.sin(-rot_angle_list_1[cnt]) body_verts_opt_prox_s1[:, 1] = temp[:, 0] * math.sin(-rot_angle_list_1[cnt]) + \ temp[:, 1] * math.cos(-rot_angle_list_1[cnt]) body_verts_opt_prox_s1[:, 2] = temp[:, 2] # transfrom local params to prox coordinate system body_params_prox_s1 = update_globalRT_for_smplx( body_params_opt_s1[0].cpu().numpy(), smplx_model, trans_matrix_2) # [72] body_params_prox_s1 = update_globalRT_for_smplx( body_params_prox_s1, smplx_model, trans_matrix_1) # [72] body_params_prox_list_s1.append(body_params_prox_s1) ### stage 2: advanced optimiation results body_verts_opt_prox_s2 = np.zeros( body_verts_opt_s2.shape) # [10475, 3] temp = body_verts_opt_s2 - shift_list[cnt] body_verts_opt_prox_s2[:, 0] = temp[:, 0] * math.cos(-rot_angle_list_1[cnt]) - \ temp[:, 1] * math.sin(-rot_angle_list_1[cnt]) body_verts_opt_prox_s2[:, 1] = temp[:, 0] * math.sin(-rot_angle_list_1[cnt]) + \ temp[:, 1] * math.cos(-rot_angle_list_1[cnt]) body_verts_opt_prox_s2[:, 2] = temp[:, 2] # transfrom local params to prox coordinate system body_params_prox_s2 = update_globalRT_for_smplx( body_params_opt_s2[0].cpu().numpy(), smplx_model, trans_matrix_2) # [72] body_params_prox_s2 = update_globalRT_for_smplx( body_params_prox_s2, smplx_model, trans_matrix_1) # [72] body_params_prox_list_s2.append(body_params_prox_s2) body_verts_opt_prox_s2_list.append(body_verts_opt_prox_s2) ########################### visualization ########################## if args.visualize: body_mesh_sample = o3d.geometry.TriangleMesh() body_mesh_sample.vertices = o3d.utility.Vector3dVector( body_verts_sample_prox) body_mesh_sample.triangles = o3d.utility.Vector3iVector( smplx_model.faces) body_mesh_sample.compute_vertex_normals() body_mesh_opt_s1 = o3d.geometry.TriangleMesh() body_mesh_opt_s1.vertices = o3d.utility.Vector3dVector( body_verts_opt_prox_s1) body_mesh_opt_s1.triangles = o3d.utility.Vector3iVector( smplx_model.faces) body_mesh_opt_s1.compute_vertex_normals() body_mesh_opt_s2 = o3d.geometry.TriangleMesh() body_mesh_opt_s2.vertices = o3d.utility.Vector3dVector( body_verts_opt_prox_s2) body_mesh_opt_s2.triangles = o3d.utility.Vector3iVector( smplx_model.faces) body_mesh_opt_s2.compute_vertex_normals() o3d.visualization.draw_geometries( [scene, body_mesh_sample]) # generated body mesh by cvae o3d.visualization.draw_geometries([scene, body_mesh_opt_s1 ]) # simple-optimized body mesh o3d.visualization.draw_geometries([scene, body_mesh_opt_s2 ]) # adv-optimizaed body mesh ##################### compute non-collision/contact score ############## # body verts before optimization body_verts_sample_prox_tensor = torch.from_numpy( body_verts_sample_prox).float().unsqueeze(0).to( device) # [1, 10475, 3] norm_verts_batch = (body_verts_sample_prox_tensor - s_grid_min_batch ) / (s_grid_max_batch - s_grid_min_batch) * 2 - 1 body_sdf_batch = F.grid_sample(s_sdf_batch.unsqueeze(1), norm_verts_batch[:, :, [2, 1, 0]].view( -1, 10475, 1, 1, 3), padding_mode='border') if body_sdf_batch.lt(0).sum().item( ) < 1: # if no interpenetration: negative sdf entries is less than one loss_non_collision_sample += 1.0 loss_contact_sample += 0.0 else: loss_non_collision_sample += (body_sdf_batch > 0).sum().float().item() / 10475.0 loss_contact_sample += 1.0 # stage 1: simple optimization results body_verts_opt_prox_tensor = torch.from_numpy( body_verts_opt_prox_s1).float().unsqueeze(0).to( device) # [1, 10475, 3] norm_verts_batch = (body_verts_opt_prox_tensor - s_grid_min_batch) / ( s_grid_max_batch - s_grid_min_batch) * 2 - 1 body_sdf_batch = F.grid_sample(s_sdf_batch.unsqueeze(1), norm_verts_batch[:, :, [2, 1, 0]].view( -1, 10475, 1, 1, 3), padding_mode='border') if body_sdf_batch.lt(0).sum().item( ) < 1: # if no interpenetration: negative sdf entries is less than one loss_non_collision_opt_s1 += 1.0 loss_contact_opt_s1 += 0.0 else: loss_non_collision_opt_s1 += (body_sdf_batch > 0).sum().float().item() / 10475.0 loss_contact_opt_s1 += 1.0 # stage 2: advanced optimization results body_verts_opt_prox_tensor = torch.from_numpy( body_verts_opt_prox_s2).float().unsqueeze(0).to( device) # [1, 10475, 3] norm_verts_batch = (body_verts_opt_prox_tensor - s_grid_min_batch) / ( s_grid_max_batch - s_grid_min_batch) * 2 - 1 body_sdf_batch = F.grid_sample(s_sdf_batch.unsqueeze(1), norm_verts_batch[:, :, [2, 1, 0]].view( -1, 10475, 1, 1, 3), padding_mode='border') if body_sdf_batch.lt(0).sum().item( ) < 1: # if no interpenetration: negative sdf entries is less than one loss_non_collision_opt_s2 += 1.0 loss_contact_opt_s2 += 0.0 else: loss_non_collision_opt_s2 += (body_sdf_batch > 0).sum().float().item() / 10475.0 loss_contact_opt_s2 += 1.0 print('scene:', args.scene_name) loss_non_collision_sample = loss_non_collision_sample / n_sample loss_contact_sample = loss_contact_sample / n_sample print('w/o optimization body: non_collision score:', loss_non_collision_sample) print('w/o optimization body: contact score:', loss_contact_sample) loss_non_collision_opt_s1 = loss_non_collision_opt_s1 / n_sample loss_contact_opt_s1 = loss_contact_opt_s1 / n_sample print('optimized body s1: non_collision score:', loss_non_collision_opt_s1) print('optimized body s1: contact score:', loss_contact_opt_s1) loss_non_collision_opt_s2 = loss_non_collision_opt_s2 / n_sample loss_contact_opt_s2 = loss_contact_opt_s2 / n_sample print('optimized body s2: non_collision score:', loss_non_collision_opt_s2) print('optimized body s2: contact score:', loss_contact_opt_s2)
def __init__(self, opt): BaseModel.initialize(self, opt) # set params self.inputSize = opt.inputSize self.single_branch = opt.single_branch self.two_branch = opt.two_branch self.aux_as_main = opt.aux_as_main assert (not self.single_branch and self.two_branch) or ( self.single_branch and not self.two_branch) if self.aux_as_main: assert self.single_branch if opt.isTrain and opt.process_rank <= 0: if self.two_branch: print("!!!!!!!!!!!! Attention, use two branch framework") time.sleep(10) else: print("!!!!!!!!!!!! Attention, use one branch framework") self.total_params_dim = opt.total_params_dim self.cam_params_dim = opt.cam_params_dim self.pose_params_dim = opt.pose_params_dim self.shape_params_dim = opt.shape_params_dim self.top_finger_joints_type = opt.top_finger_joints_type assert(self.total_params_dim == self.cam_params_dim+self.pose_params_dim+self.shape_params_dim) if opt.dist: self.batch_size = opt.batchSize // torch.distributed.get_world_size() else: self.batch_size = opt.batchSize nb = self.batch_size # set input image and 2d keypoints self.input_img = self.Tensor( nb, opt.input_nc, self.inputSize, self.inputSize) # joints 2d self.keypoints = self.Tensor(nb, opt.num_joints, 2) self.keypoints_weights = self.Tensor(nb, opt.num_joints) # mano pose params self.gt_pose_params = self.Tensor(nb, opt.pose_params_dim) self.mano_params_weight = self.Tensor(nb, 1) # joints 3d self.joints_3d = self.Tensor(nb, opt.num_joints, 3) self.joints_3d_weight = self.Tensor(nb, opt.num_joints, 1) # load mean params, the mean params are from HMR self.mean_param_file = osp.join( opt.model_root, opt.mean_param_file) self.load_params() # set differential SMPL (implemented with pytorch) and smpl_renderer # smplx_model_path = osp.join(opt.model_root, opt.smplx_model_file) smplx_model_path = opt.smplx_model_file self.smplx = smplx.create( smplx_model_path, model_type = "smplx", batch_size = self.batch_size, gender = 'neutral', num_betas = 10, use_pca = False, ext='pkl').cuda() # set encoder and optimizer self.encoder = H3DWEncoder(opt, self.mean_params).cuda() if opt.dist: self.encoder = DistributedDataParallel( self.encoder, device_ids=[torch.cuda.current_device()]) if self.isTrain: self.optimizer_E = torch.optim.Adam( self.encoder.parameters(), lr=opt.lr_e) # load pretrained / trained weights for encoder if self.isTrain: assert False, "Not implemented Yet" pass else: # load trained model for testing which_epoch = opt.which_epoch if which_epoch == 'latest' or int(which_epoch)>0: self.success_load = self.load_network(self.encoder, 'encoder', which_epoch) else: checkpoint_path = opt.checkpoint_path if not osp.exists(checkpoint_path): print(f"Error: {checkpoint_path} does not exists") self.success_load = False else: if self.opt.dist: self.encoder.module.load_state_dict(torch.load( checkpoint_path, map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()))) else: saved_weights = torch.load(checkpoint_path) self.encoder.load_state_dict(saved_weights) self.success_load = True
proxd_path = os.path.join(dataset_path, 'PROXD') cam2world_path = os.path.join(dataset_path, 'cam2world') vposer_model_path = os.path.join(dataset_path, 'body_models/vposer_v1_0') smplx_model_path = os.path.join(dataset_path, 'body_models/smplx_model') scene_mesh_path = os.path.join(dataset_path, 'scenes_downsampled') scene_sdf_path = os.path.join(dataset_path, 'scenes_sdf') smplx_model = smplx.create(smplx_model_path, model_type='smplx', gender='neutral', ext='npz', num_pca_comps=12, create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True, batch_size=batch_size ).to(device) print('[INFO] smplx model loaded.') vposer_model, _ = load_vposer(vposer_model_path, vp_model='snapshot') vposer_model = vposer_model.to(device) print('[INFO] vposer model loaded') ######## set body mesh gen dataloader ###########
def save_grab_vertices(cfg, logger=None, **params): grab_path = cfg.grab_path out_path = cfg.out_path makepath(out_path) if logger is None: logger = makelogger(log_dir=os.path.join(out_path, 'grab_preprocessing.log'), mode='a').info else: logger = logger logger('Starting to get vertices for GRAB!') all_seqs = glob.glob(grab_path + '/*/*.npz') logger('Total sequences: %d' % len(all_seqs)) # stime = datetime.now().replace(microsecond=0) # shutil.copy2(sys.argv[0], # os.path.join(out_path, # os.path.basename(sys.argv[0]).replace('.py','_%s.py' % datetime.strftime(stime,'%Y%m%d_%H%M')))) if out_path is None: out_path = grab_path for sequence in tqdm(all_seqs): outfname = makepath(sequence.replace(grab_path, out_path).replace( '.npz', '_verts_body.npz'), isfile=True) action_name = os.path.basename(sequence) if os.path.exists(outfname): logger('Results for %s split already exist.' % (action_name)) continue else: logger('Processing data for %s split.' % (action_name)) seq_data = parse_npz(sequence) n_comps = seq_data['n_comps'] gender = seq_data['gender'] T = seq_data.n_frames if cfg.save_body_verts: sbj_mesh = os.path.join(grab_path, '..', seq_data.body.vtemp) sbj_vtemp = np.array(Mesh(filename=sbj_mesh).vertices) sbj_m = smplx.create(model_path=cfg.model_path, model_type='smplx', gender=gender, num_pca_comps=n_comps, v_template=sbj_vtemp, batch_size=T) sbj_parms = params2torch(seq_data.body.params) verts_sbj = to_cpu(sbj_m(**sbj_parms).vertices) np.savez_compressed(outfname, verts_body=verts_sbj) if cfg.save_lhand_verts: lh_mesh = os.path.join(grab_path, '..', seq_data.lhand.vtemp) lh_vtemp = np.array(Mesh(filename=lh_mesh).vertices) lh_m = smplx.create(model_path=cfg.model_path, model_type='mano', is_rhand=False, v_template=lh_vtemp, num_pca_comps=n_comps, flat_hand_mean=True, batch_size=T) lh_parms = params2torch(seq_data.lhand.params) verts_lh = to_cpu(lh_m(**lh_parms).vertices) np.savez_compressed(outfname.replace('_verts_body.npz', '_verts_lhand.npz'), verts_body=verts_lh) if cfg.save_rhand_verts: rh_mesh = os.path.join(grab_path, '..', seq_data.rhand.vtemp) rh_vtemp = np.array(Mesh(filename=rh_mesh).vertices) rh_m = smplx.create(model_path=cfg.model_path, model_type='mano', is_rhand=True, v_template=rh_vtemp, num_pca_comps=n_comps, flat_hand_mean=True, batch_size=T) rh_parms = params2torch(seq_data.body.params) verts_rh = to_cpu(rh_m(**rh_parms).vertices) np.savez_compressed(outfname.replace('_verts_body.npz', '_verts_rhand.npz'), verts_body=verts_rh) if cfg.save_object_verts: obj_mesh = os.path.join(grab_path, '..', seq_data.object.object_mesh) obj_vtemp = np.array(Mesh(filename=obj_mesh).vertices) sample_id = np.random.choice(obj_vtemp.shape[0], cfg.n_verts_sample, replace=False) obj_m = ObjectModel(v_template=obj_vtemp[sample_id], batch_size=T) obj_parms = params2torch(seq_data.object.params) verts_obj = to_cpu(obj_m(**obj_parms).vertices) np.savez_compressed(outfname.replace('_verts_body.npz', '_verts_object.npz'), verts_object=verts_obj) logger('Processing finished')