|
- 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
|