| import os |
| import tyro |
| import mediapy |
| import torch |
| import numpy as np |
| import pyvista as pv |
| import trimesh |
| from PIL import Image |
|
|
| from dreifus.matrix import Intrinsics, Pose, CameraCoordinateConvention, PoseType |
| from dreifus.pyvista import add_camera_frustum, render_from_camera |
|
|
| from pixel3dmm.utils.utils_3d import rotation_6d_to_matrix |
| from pixel3dmm.env_paths import PREPROCESSED_DATA, TRACKING_OUTPUT |
|
|
|
|
| def main(vid_name : str, |
| HEAD_CENTRIC : bool = True, |
| DO_PROJECTION_TEST : bool = False, |
| ): |
| tracking_dir = f'{TRACKING_OUTPUT}/{vid_name}_nV1_noPho_uv2000.0_n1000.0' |
|
|
| meshes = [f for f in os.listdir(f'{tracking_dir}/mesh/') if f.endswith('.ply') and not 'canonical' in f] |
| meshes.sort() |
|
|
| ckpts = [f for f in os.listdir(f'{tracking_dir}/checkpoint/') if f.endswith('.frame')] |
| ckpts.sort() |
|
|
| N_STEPS = len(meshes) |
|
|
| pl = pv.Plotter() |
| vid_frames = [] |
| for i in range(N_STEPS): |
| ckpt = torch.load(f'{tracking_dir}/checkpoint/{ckpts[i]}', weights_only=False) |
|
|
| mesh = trimesh.load(f'{tracking_dir}/mesh/{meshes[i]}', process=False) |
|
|
| head_rot = rotation_6d_to_matrix(torch.from_numpy(ckpt['flame']['R'])).numpy()[0] |
|
|
| if not HEAD_CENTRIC: |
| |
| mesh.vertices = mesh.vertices @ head_rot.T + (ckpt['flame']['t']) |
| else: |
| |
| verts_hom = np.concatenate([mesh.vertices, np.ones_like(mesh.vertices[..., :1])], axis=-1) |
| verts_hom = verts_hom @ np.linalg.inv(ckpt['joint_transforms'][0, 1, :, :]).T |
| mesh.vertices = verts_hom[..., :3] |
|
|
|
|
|
|
| extr_open_gl_world_to_cam = np.eye(4) |
| extr_open_gl_world_to_cam[:3, :3] = ckpt['camera']['R_base_0'][0] |
| extr_open_gl_world_to_cam[:3, 3] = ckpt['camera']['t_base_0'][0] |
| if HEAD_CENTRIC: |
| flame2world = np.eye(4) |
| flame2world[:3, :3] = head_rot |
| flame2world[:3, 3] = np.squeeze(ckpt['flame']['t']) |
| |
| extr_open_gl_world_to_cam = extr_open_gl_world_to_cam @ flame2world @ ckpt['joint_transforms'][0, 1, :, :] |
|
|
|
|
|
|
|
|
| extr_open_gl_world_to_cam = Pose(extr_open_gl_world_to_cam, |
| camera_coordinate_convention=CameraCoordinateConvention.OPEN_GL, |
| pose_type=PoseType.WORLD_2_CAM) |
|
|
| intr = np.eye(3) |
| intr[0, 0] = ckpt['camera']['fl'][0, 0] * 256 |
| intr[1, 1] = ckpt['camera']['fl'][0, 0] * 256 |
| intr[:2, 2] = ckpt['camera']['pp'][0] * (256/2+0.5) + 256/2 + 0.5 |
|
|
| intr = Intrinsics(intr) |
|
|
|
|
|
|
| pl.add_mesh(mesh, color=[(i/N_STEPS), 0, ((N_STEPS-i)/N_STEPS)]) |
| add_camera_frustum(pl, extr_open_gl_world_to_cam, intr, color=[(i/N_STEPS), 0, ((N_STEPS-i)/N_STEPS)]) |
|
|
| if DO_PROJECTION_TEST: |
| pll = pv.Plotter(off_screen=True, window_size=(256, 256)) |
| pll.add_mesh(mesh) |
| img = render_from_camera(pll, extr_open_gl_world_to_cam, intr) |
|
|
| gt_img = np.array(Image.open(f'{PREPROCESSED_DATA}/{vid_name}/cropped/{i:05d}.jpg').resize((256, 256))) |
|
|
| alpha = img[..., 3] |
|
|
| overlay = (gt_img *0.5 + img[..., :3]*0.5).astype(np.uint8) |
| vid_frames.append(overlay) |
|
|
|
|
|
|
|
|
| pl.show() |
|
|
| if DO_PROJECTION_TEST: |
| mediapy.write_video(f'{tracking_dir}/projection_test.mp4', images=vid_frames) |
|
|
|
|
|
|
| if __name__ == '__main__': |
| tyro.cli(main) |
