import numpy as np import tqdm import os, glob import json import argparse from render_utils.lib.networks.smpl_torch import SmplTorch from render_utils.lib.utils.gaussian_np_utils import load_gaussians_from_ply from render_utils.stitch_body_and_head import load_body_params, load_face_params, get_smpl_verts_and_head_transformation, calc_livehead2livebody def load_rendering_camera(camera_fpath): with open(camera_fpath, 'r') as fp: camera_data = json.load(fp) camera_data = camera_data[0] image_size = [camera_data['width'], camera_data['height']] cam_f = [camera_data['fx'], camera_data['fy']] cam_pos = np.array(camera_data['position']) cam_rot = np.array(camera_data['rotation']).reshape(3, 3) c2w = np.eye(4) c2w[:3, :3] = cam_rot c2w[:3, 3] = cam_pos cam_extr = np.linalg.inv(c2w) cam_intr = np.eye(3) cam_intr[0, 0] = cam_f[0] cam_intr[1, 1] = cam_f[1] cam_intr[0, 2] = image_size[0] / 2 cam_intr[1, 2] = image_size[1] / 2 return cam_extr, cam_intr, image_size def load_camera_list(camera_fpath): with open(camera_fpath, 'r') as fp: camera_data = json.load(fp) image_size = [camera_data[0]['width'], camera_data[0]['height']] cam_list = [] for cam in camera_data: cam_f = [cam['fx'], cam['fy']] cam_pos = np.array(cam['position']) cam_rot = np.array(cam['rotation']).reshape(3, 3) c2w = np.eye(4) c2w[:3, :3] = cam_rot c2w[:3, 3] = cam_pos cam_extr = np.linalg.inv(c2w) cam_intr = np.eye(3) cam_intr[0, 0] = cam_f[0] cam_intr[1, 1] = cam_f[1] cam_intr[0, 2] = image_size[0] / 2 cam_intr[1, 2] = image_size[1] / 2 cam_list.append((cam_extr, cam_intr)) return cam_list, image_size def load_camera_data(cam): image_size = [cam['width'], cam['height']] cam_f = [cam['fx'], cam['fy']] cam_pos = np.array(cam['position']) cam_rot = np.array(cam['rotation']).reshape(3, 3) c2w = np.eye(4) c2w[:3, :3] = cam_rot c2w[:3, 3] = cam_pos cam_extr = np.linalg.inv(c2w) cam_intr = np.eye(3) cam_intr[0, 0] = cam_f[0] cam_intr[1, 1] = cam_f[1] cam_intr[0, 2] = image_size[0] / 2 cam_intr[1, 2] = image_size[1] / 2 return (cam_extr, cam_intr), image_size def calc_offline_rendering_param( body_gaussian_root_dir, ref_head_gaussian_path, ref_head_param_path, render_cam_fpath, body_head_blending_param_path): body_param_flist = sorted(glob.glob(os.path.join(body_gaussian_root_dir, 'posed_params/*.npz'))) head_gaussians = load_gaussians_from_ply(ref_head_gaussian_path) head_pose, head_scale, id_coeff, exp_coeff = load_face_params(ref_head_param_path) # cam_extr_body, cam_intr_body, image_size = load_rendering_camera(render_cam_fpath) cam_list, image_size = load_camera_list(render_cam_fpath) body_head_blending_params = np.load(body_head_blending_param_path) smplx_to_faceverse = body_head_blending_params['smplx_to_faceverse'] residual_transf = body_head_blending_params['residual_transf'] body_nonface_mask = body_head_blending_params['body_nonface_mask'] head_nonface_mask = body_head_blending_params['head_nonface_mask'] head_facial_idx = body_head_blending_params['head_facial_idx'] body_facial_idx = body_head_blending_params['body_facial_idx'] head_body_corr_idx = body_head_blending_params['head_body_corr_idx'] head_color_bw = body_head_blending_params['head_color_bw'] color_transfer = body_head_blending_params['color_transfer'] smpl = SmplTorch(model_file='./AnimatableGaussians/smpl_files/smplx/SMPLX_NEUTRAL.npz') head_cam_extr = [] head_cam_intr = [] head_cam_intr_zoom = [] head_zoom_center = [] head_zoom_scale = [] for i, body_param_fpath in enumerate(tqdm.tqdm(body_param_flist)): global_orient, transl, body_pose, betas = load_body_params(body_param_fpath) # body_gaussians = load_gaussians_from_ply(body_gaussian_fpath) smpl_verts, head_joint_transfmat = get_smpl_verts_and_head_transformation( smpl, global_orient, body_pose, transl, betas) livehead2livebody = calc_livehead2livebody(head_pose, smplx_to_faceverse, head_joint_transfmat) total_transf = np.matmul(livehead2livebody, residual_transf) cam_extr = np.matmul(cam_list[i][0], total_transf) cam_intr = np.copy(cam_list[i][1]) head_cam_extr.append(cam_extr) head_cam_intr.append(cam_intr) pts = np.copy(head_gaussians.xyz) pts_proj = np.matmul(pts, cam_extr[:3, :3].transpose()) + cam_extr[:3, 3] pts_proj = np.matmul(pts_proj, cam_intr.transpose()) pts_proj = pts_proj / pts_proj[:, 2:] # pts_proj = np.int32(np.round(pts_proj[:, :2])) # img = np.zeros([image_size[1], image_size[0], 3], dtype=np.uint8) # for p in pts_proj[::50]: # p = np.clip(p, 0, image_size[0] - 1) # cv.circle(img, (int(p[0]), int(p[1])), 2, (0, 255, 0), -1) # cv.imshow('img', img) pts_min, pts_max = np.min(pts_proj, axis=0), np.max(pts_proj, axis=0) pts_center = (pts_min + pts_max) // 2 pts_size = np.max(pts_max - pts_min) tgt_pts_size = 350 tgt_image_size = 512 zoom_scale = tgt_pts_size / pts_size cam_intr_zoom = np.copy(cam_intr) cam_intr_zoom[:2] *= zoom_scale cam_intr_zoom[0, 2] = cam_intr_zoom[0, 2] - (pts_center[0]*zoom_scale - tgt_image_size/2) cam_intr_zoom[1, 2] = cam_intr_zoom[1, 2] - (pts_center[1]*zoom_scale - tgt_image_size/2) head_cam_intr_zoom.append(cam_intr_zoom) head_zoom_center.append(pts_center) head_zoom_scale.append(zoom_scale) # pts_proj = np.matmul(pts, cam_extr[:3, :3].transpose()) + cam_extr[:3, 3] # pts_proj = np.matmul(pts_proj, cam_intr_zoom.transpose()) # pts_proj = pts_proj / pts_proj[:, 2:] # pts_proj = np.int32(np.round(pts_proj[:, :2])) # img = np.zeros([512, 512, 3], dtype=np.uint8) # for p in pts_proj[::50]: # p = np.clip(p, 0, image_size[0] - 1) # cv.circle(img, (int(p[0]), int(p[1])), 2, (0, 255, 0), -1) # cv.imshow('img_zoom', img) # cv.waitKey() np.savez(os.path.join(os.path.dirname(body_head_blending_param_path), 'head_zoomin_render_param.npz'), cam_extr=head_cam_extr, cam_intr=head_cam_intr, image_size=image_size, cam_intr_zoom=head_cam_intr_zoom, zoom_image_size=[tgt_image_size, tgt_image_size], zoom_center=head_zoom_center, zoom_scale=head_zoom_scale, head_pose=head_pose, head_scale=head_scale, head_color_bw=head_color_bw) if __name__ == '__main__': parser = argparse.ArgumentParser() """ body_gaussian_root_dir, ref_head_gaussian_path, ref_head_param_path, render_cam_fpath, body_head_blending_param_path """ parser.add_argument('--body_gaussian_root_dir', type=str) parser.add_argument('--ref_head_gaussian_path', type=str) parser.add_argument('--ref_head_param_path', type=str) parser.add_argument('--render_cam_fpath', type=str) parser.add_argument('--body_head_blending_param_path', type=str) args = parser.parse_args() calc_offline_rendering_param( args.body_gaussian_root_dir, args.ref_head_gaussian_path, args.ref_head_param_path, args.render_cam_fpath, args.body_head_blending_param_path ) """ python calc_offline_rendering_param.py ^ --body_gaussian_root_dir ./AnimatableGaussians/test_results/huawei0425/checkpoints/AMASS__test_poses_ours_front_view/batch_750000/pca_20_sigma_2.00/ ^ --ref_head_gaussian_path ./Gaussian-Head-Avatar/results/reenactment/huawei0425_self/posed_gaussians/000000.ply ^ --ref_head_param_path ./Gaussian-Head-Avatar/results/reenactment/huawei0425_self/params/000000_param.npz ^ --render_cam_fpath ./AnimatableGaussians/test_results/huawei0425/checkpoints/AMASS__test_poses_ours_front_view/batch_750000/pca_20_sigma_2.00/cameras.json ^ --body_head_blending_param_path ./data/body_face_stitching_sr/body_head_blending_param.npz """