Spaces:
Configuration error
Configuration error
| import pickle | |
| import os | |
| import h5py | |
| import numpy as np | |
| import open3d as o3d | |
| from snapshot_smpl.renderer import Renderer | |
| import cv2 | |
| import tqdm | |
| def read_pickle(pkl_path): | |
| with open(pkl_path, 'rb') as f: | |
| u = pickle._Unpickler(f) | |
| u.encoding = 'latin1' | |
| return u.load() | |
| def get_KRTD(camera): | |
| K = np.zeros([3, 3]) | |
| K[0, 0] = camera['camera_f'][0] | |
| K[1, 1] = camera['camera_f'][1] | |
| K[:2, 2] = camera['camera_c'] | |
| K[2, 2] = 1 | |
| R = np.eye(3) | |
| T = np.zeros([3]) | |
| D = camera['camera_k'] | |
| return K, R, T, D | |
| def get_o3d_mesh(vertices, faces): | |
| mesh = o3d.geometry.TriangleMesh() | |
| mesh.vertices = o3d.utility.Vector3dVector(vertices) | |
| mesh.triangles = o3d.utility.Vector3iVector(faces) | |
| mesh.compute_vertex_normals() | |
| return mesh | |
| def get_smpl(base_smpl, betas, poses, trans): | |
| base_smpl.betas = betas | |
| base_smpl.pose = poses | |
| base_smpl.trans = trans | |
| vertices = np.array(base_smpl) | |
| faces = base_smpl.f | |
| mesh = get_o3d_mesh(vertices, faces) | |
| return mesh | |
| def render_smpl(vertices, img, K, R, T): | |
| rendered_img = renderer.render_multiview(vertices, K[None], R[None], | |
| T[None, None], [img])[0] | |
| return rendered_img | |
| data_root = 'data/people_snapshot' | |
| video = 'female-3-casual' | |
| # if you do not have these smpl models, you could download them from https://zjueducn-my.sharepoint.com/:u:/g/personal/pengsida_zju_edu_cn/Eb_JIyA74O1Cnfhvn1ddrG4BC9TMK31022TykVxGdRenUQ?e=JU8pPt | |
| model_paths = [ | |
| 'basicModel_f_lbs_10_207_0_v1.0.0.pkl', | |
| 'basicmodel_m_lbs_10_207_0_v1.0.0.pkl' | |
| ] | |
| camera_path = os.path.join(data_root, video, 'camera.pkl') | |
| camera = read_pickle(camera_path) | |
| K, R, T, D = get_KRTD(camera) | |
| mask_path = os.path.join(data_root, video, 'masks.hdf5') | |
| masks = h5py.File(mask_path)['masks'] | |
| smpl_path = os.path.join(data_root, video, 'reconstructed_poses.hdf5') | |
| smpl = h5py.File(smpl_path) | |
| betas = smpl['betas'] | |
| pose = smpl['pose'] | |
| trans = smpl['trans'] | |
| if 'female' in video: | |
| model_path = model_paths[0] | |
| else: | |
| model_path = model_paths[1] | |
| model_data = read_pickle(model_path) | |
| faces = model_data['f'] | |
| renderer = Renderer(height=1080, width=1080, faces=faces) | |
| img_dir = os.path.join(data_root, video, 'image') | |
| vertices_dir = os.path.join(data_root, video, 'vertices') | |
| num_img = len(os.listdir(img_dir)) | |
| for i in tqdm.tqdm(range(num_img)): | |
| img = cv2.imread(os.path.join(img_dir, '{}.jpg'.format(i))) | |
| img = cv2.undistort(img, K, D) | |
| vertices = np.load(os.path.join(vertices_dir, '{}.npy'.format(i))) | |
| rendered_img = render_smpl(vertices, img, K, R, T) | |
| cv2.imshow('main', rendered_img) | |
| cv2.waitKey(50) & 0xFF | |