calc_offline_rendering_param.py

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

"""