upload_part1

render_utils/calc_smplx2faceverse.py

@@ -0,0 +1,238 @@
+import copy
+import math
+import time, random
+
+import torch
+import numpy as np
+import cv2
+import smplx
+import pickle
+import trimesh
+import os, glob
+import argparse
+
+from .lib.networks.faceverse_torch import FaceVerseModel
+
+
+def estimate_rigid_transformation(src, tgt):
+    src = src.transpose()
+    tgt = tgt.transpose()
+    mu1, mu2 = src.mean(axis=1, keepdims=True), tgt.mean(axis=1, keepdims=True)
+    X1, X2 = src - mu1, tgt - mu2
+
+    K = X1.dot(X2.T)
+    U, s, Vh = np.linalg.svd(K)
+    V = Vh.T
+    Z = np.eye(U.shape[0])
+    Z[-1, -1] *= np.sign(np.linalg.det(U.dot(V.T)))
+    R = V.dot(Z.dot(U.T))
+    t = mu2 - R.dot(mu1)
+
+    orient, _ = cv2.Rodrigues(R)
+    orient = orient.reshape([-1])
+    t = t.reshape([-1])
+    return orient, t
+
+
+def load_smpl_beta(body_fitting_result_fpath):
+    if os.path.isdir(body_fitting_result_fpath):
+        body_fitting_result = torch.load(
+            os.path.join(body_fitting_result_fpath, 'checkpoints/latest.pt'), map_location='cpu')
+        betas = body_fitting_result['betas']['weight']
+    elif body_fitting_result_fpath.endswith('.pt'):
+        body_fitting_result = torch.load(body_fitting_result_fpath, map_location='cpu')
+        betas = body_fitting_result['beta'].reshape(1, -1)
+    elif body_fitting_result_fpath.endswith('.npz'):
+        body_fitting_result = np.load(body_fitting_result_fpath)
+        betas = body_fitting_result['betas'].reshape(1, -1)
+        betas = torch.from_numpy(betas.astype(np.float32))
+    else:
+        raise ValueError('Unknown body fitting result file format: {}'.format(body_fitting_result_fpath))
+    return betas
+
+
+def load_face_id_scale_param(face_fitting_result_fpath):
+    if os.path.isfile(face_fitting_result_fpath):
+        face_fitting_result = dict(np.load(face_fitting_result_fpath))
+        id_tensor = face_fitting_result['id_coeff'].astype(np.float32)
+        scale_tensor = face_fitting_result['scale'].astype(np.float32)
+        id_tensor = torch.from_numpy(id_tensor).reshape(1, -1)
+        scale_tensor = torch.from_numpy(scale_tensor).reshape(1, -1)
+    else:
+        param_paths = sorted(glob.glob(os.path.join(face_fitting_result_fpath, '*', 'params.npz')))
+        param = np.load(param_paths[0])
+        id_tensor = torch.from_numpy(param['id_coeff']).reshape(1, -1)
+        scale_tensor = torch.from_numpy(param['scale']).reshape(1, 1)
+
+    return id_tensor, scale_tensor
+
+
+def calc_smplx2faceverse(body_fitting_result_fpath, face_fitting_result_fpath, output_dir):
+    device = torch.device('cuda')
+    os.makedirs(output_dir, exist_ok=True)
+
+    betas = load_smpl_beta(body_fitting_result_fpath)
+    id_tensor, scale_tensor = load_face_id_scale_param(face_fitting_result_fpath)
+
+    smpl = smplx.SMPLX(model_path='./AnimatableGaussians/smpl_files/smplx', gender='neutral',
+                       use_pca=True, num_pca_comps=45, flat_hand_mean=True, batch_size=1)
+    flame = smplx.FLAME(model_path='./AnimatableGaussians/smpl_files/FLAME2019', gender='neutral', batch_size=1)
+
+    pose = np.zeros([63], dtype=np.float32)
+
+    pose = torch.from_numpy(pose).unsqueeze(0)
+    smpl_out = smpl(body_pose=pose, betas=betas)
+    verts = smpl_out.vertices.detach().cpu().squeeze(0).numpy()
+    flame_out = flame()
+    verts_flame = flame_out.vertices.detach().cpu().squeeze(0).numpy()
+
+    smplx2flame_data = np.load('./data/smpl_models/smplx_mano_flame_correspondences/SMPL-X__FLAME_vertex_ids.npy')
+    verts_flame_on_smplx = verts[smplx2flame_data]
+
+    orient, t = estimate_rigid_transformation(verts_flame, verts_flame_on_smplx)
+    R, _ = cv2.Rodrigues(orient)
+
+    rel_transf = np.eye(4)
+    rel_transf[:3, :3] = R
+    rel_transf[:3, 3] = t.reshape(-1)
+    np.save('%s/flame_to_smplx.npy' % (output_dir), rel_transf.astype(np.float32))
+
+    # TODO: DELETE ME
+    with open('./debug/debug_verts_smplx.obj', 'w') as fp:
+        for v in verts:
+            fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+    with open('./debug/debug_verts_flame_in_smplx.obj', 'w') as fp:
+        for v in np.matmul(verts_flame, R.transpose()) + t.reshape([1, 3]):
+            fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+
+    # align Faceverse to T-pose FLAME on SMPL-X
+    faceverse_mesh = trimesh.load_mesh('./data/smpl_models/faceverse2flame/faceverse_icp.obj', process=False)
+    verts_faceverse_ref = np.matmul(np.asarray(faceverse_mesh.vertices), R.transpose()) + t.reshape(1, 3)
+
+    # TODO: DELETE ME
+    with open('./debug/debug_verts_faceverse_ref.obj', 'w') as fp:
+        for v in verts_faceverse_ref:
+            fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+
+    model_dict = np.load('./data/faceverse_models/faceverse_simple_v2.npy', allow_pickle=True).item()
+    faceverse_model = FaceVerseModel(model_dict, batch_size=1)
+    faceverse_model.init_coeff_tensors()
+    coeffs = faceverse_model.get_packed_tensors()
+    fv_out = faceverse_model.forward(coeffs=coeffs)
+    verts_faceverse = fv_out['v'].squeeze(0).detach().cpu().numpy()
+
+    # calculate the relative transformation between FLAME in canonical pose and its position on SMPL-X
+    orient, t = estimate_rigid_transformation(verts_faceverse, verts_faceverse_ref)
+    orient = torch.from_numpy(orient.astype(np.float32)).unsqueeze(0).to(device)
+    t = torch.from_numpy(t.astype(np.float32)).unsqueeze(0).to(device)
+
+    # optimize the Faceverse to fit SMPL-X
+    faceverse_model.init_coeff_tensors(
+        rot_coeff=orient, trans_coeff=t, id_coeff=id_tensor.to(device), scale_coeff=scale_tensor.to(device))
+    nonrigid_optim_params = [
+        faceverse_model.get_exp_tensor(), faceverse_model.get_rot_tensor(), faceverse_model.get_trans_tensor(),
+        # faceverse_model.get_scale_tensor(), faceverse_model.get_id_tensor()
+    ]
+    nonrigid_optimizer = torch.optim.Adam(nonrigid_optim_params, lr=1e-1)
+    verts_faceverse_ref = torch.from_numpy(verts_faceverse_ref.astype(np.float32)).to(device).unsqueeze(0)
+    for iter in range(1000):
+        coeffs = faceverse_model.get_packed_tensors()
+        fv_out = faceverse_model.forward(coeffs=coeffs)
+        verts_pred = fv_out['v']
+        loss = torch.mean(torch.square(verts_pred - verts_faceverse_ref))
+        if iter % 10 == 0:
+            print(loss.item())
+        nonrigid_optimizer.zero_grad()
+        loss.backward()
+        nonrigid_optimizer.step()
+
+    np.savez('%s/faceverse_param_to_smplx.npz' % (output_dir), {
+        'id': faceverse_model.get_id_tensor().detach().cpu().numpy(),
+        'exp': faceverse_model.get_exp_tensor().detach().cpu().numpy(),
+        'rot': faceverse_model.get_rot_tensor().detach().cpu().numpy(),
+        'transl': faceverse_model.get_trans_tensor().detach().cpu().numpy(),
+        'scale': faceverse_model.get_scale_tensor().detach().cpu().numpy(),
+    })
+
+    # calculate SMPLX to faceverse space transformation (without scale)
+    orient = faceverse_model.get_rot_tensor().detach().cpu().numpy()
+    transl = faceverse_model.get_trans_tensor().detach().cpu().numpy()
+    rotmat, _ = cv2.Rodrigues(orient)
+    transform_mat = np.eye(4)
+    transform_mat[:3, :3] = rotmat
+    transform_mat[:3, 3] = transl
+    transform_mat = np.linalg.inv(transform_mat)
+    np.save('%s/smplx_to_faceverse_space.npy' % (output_dir), transform_mat.astype(np.float32))
+
+    # calculate SMPLX to faceverse distance
+    dists = []
+    verts_faceverse_ref = verts_faceverse_ref.detach().cpu().numpy()
+    for v in verts:
+        dist = np.linalg.norm(v.reshape(1, 3) - verts_faceverse_ref, axis=-1)
+        dist = np.min(dist)
+        dists.append(dist)
+    dists = np.asarray(dists)
+    np.save('%s/smplx_verts_to_faceverse_dist.npy' % (output_dir), dists.astype(np.float32))
+
+    # sample nodes on facial area
+    dists_ = np.ones_like(dists)
+    smplx_topo_new = np.load('./data/smpl_models/smplx_topo_new.npy')
+    valid_vids = set(smplx_topo_new.reshape([-1]).tolist())
+    dists_[np.asarray(list(valid_vids))] = dists[np.asarray(list(valid_vids))]
+
+    vids_on_face = np.where(dists_ < 0.01)[0]
+    verts_on_face = verts[vids_on_face]
+    geod_dist_mat = np.linalg.norm(np.expand_dims(verts_on_face, axis=0) - np.expand_dims(verts_on_face, axis=1),
+                                   axis=2)
+    nodes = [0]  # nose
+    dist_nodes_to_rest_points = geod_dist_mat[nodes[0]]
+    for i in range(1, 256):
+        idx = np.argmax(dist_nodes_to_rest_points)
+        nodes.append(idx)
+        new_dist = geod_dist_mat[idx]
+        update_flag = new_dist < dist_nodes_to_rest_points
+        dist_nodes_to_rest_points[update_flag] = new_dist[update_flag]
+
+    # with open('./debug/debug_face_nodes.obj', 'w') as fp:
+    #     for n in verts_on_face[np.asarray(nodes)]:
+    #         fp.write('v %f %f %f\n' % (n[0], n[1], n[2]))
+    vids_on_faces_sampled = vids_on_face[np.asarray(nodes)]
+    vids_on_faces_sampled = np.ascontiguousarray(vids_on_faces_sampled).astype(np.int32)
+    np.save('%s/vids_on_faces_sampled.npy' % (output_dir), vids_on_faces_sampled)
+
+    # test SMPLX-to-faceverse space transformation (without scale)
+    verts_smpl_in_faceverse = np.matmul(verts, transform_mat[:3, :3].transpose()) + \
+        transform_mat[:3, 3].reshape(1, 3)
+    with open('./debug/debug_verts_smpl_in_faceverse.obj', 'w') as fp:
+        for v in verts_smpl_in_faceverse:
+            fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+
+    # save personalized, canonical faceverse model
+    faceverse_model.init_coeff_tensors(id_coeff=id_tensor.to(device), scale_coeff=scale_tensor.to(device))
+    coeffs = faceverse_model.get_packed_tensors()
+    fv_out = faceverse_model.forward(coeffs=coeffs)
+    verts_faceverse = fv_out['v'].squeeze(0).detach().cpu().numpy()
+    with open('./debug/debug_verts_faceverse.obj', 'w') as fp:
+        for v in verts_faceverse:
+            fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
+
+
+if __name__ == '__main__':
+    # body_fitting_result_dir = 'D:/UpperBodyAvatar/code/smplfitting_multiview_sequence_smplx/results/Shuangqing/zzr_fullbody_20221130_01_2k/whole.pt'
+    # face_fitting_result_dir = 'D:\\UpperBodyAvatar\\data\\Shuangqing\\zzr_face_20221130_01_2k\\faceverse_params'
+    #
+    # output_dir = './data/faceverse/'
+    # result_suffix = 'shuangqing_zzr'
+
+    # body_fitting_result_dir = 'D:/Product/FullAppRelease/smplfitting_multiview_sequence_smplx/results/body_data_model_20231224/whole.pt'
+    # face_fitting_result_dir = 'D:/Product/data/HuiyingCenter/20231224_model/model_20231224_face_data/faceverse_params'
+    #
+    # output_dir = './data/faceverse/'
+    # result_suffix = 'huiyin_model20231224'
+
+    body_fitting_result_dir = 'D:/Product/FullAppRelease/smplfitting_multiview_sequence_smplx/results/body_data_male20230530_betterhand/whole.pt'
+    face_fitting_result_dir = 'D:/Product/data/HuiyingCenter/20230531_models/male20230530_face_data/faceverse_params'
+
+    output_dir = './data/body_face_stitching/huiyin_male20230530'
+
+    calc_smplx2faceverse(body_fitting_result_dir, face_fitting_result_dir, output_dir)

render_utils/camera_dir.py

@@ -0,0 +1,171 @@
+import numpy as np
+
+
+def get_camera_dir(idx):
+    img_scale = self.body['test'].get('img_scale', 1.0)
+    view_setting = self.body['test'].get('view_setting', 'free')
+    if view_setting == 'camera':
+        # training view setting
+        cam_id = self.body['test']['render_view_idx']
+        intr = self.dataset.intr_mats[cam_id].copy()
+        intr[:2] *= img_scale
+        extr = self.dataset.extr_mats[cam_id].copy()
+        img_h, img_w = int(self.dataset.img_heights[cam_id] * img_scale), int(self.dataset.img_widths[cam_id] * img_scale)
+    elif view_setting.startswith('free'):
+        # free view setting
+        # frame_num_per_circle = 360
+        # print(self.opt['test'].get('global_orient', False))
+        frame_num_per_circle = 360
+        rot_Y = (idx % frame_num_per_circle) / float(frame_num_per_circle) * 2 * np.pi
+
+        extr = visualize_util.calc_free_mv(object_center,
+                                           tar_pos = np.array([0, 0, 2.5]),
+                                           rot_Y = rot_Y,
+                                           rot_X = 0.3 if view_setting.endswith('bird') else 0.,
+                                           global_orient = global_orient if self.body['test'].get('global_orient', False) else None)
+        intr = np.array([[1100, 0, 512], [0, 1100, 512], [0, 0, 1]], np.float32)
+        intr[:2] *= img_scale
+        img_h = int(1024 * img_scale)
+        img_w = int(1024 * img_scale)
+        
+        extr_list.append(extr)
+        intr_list.append(intr)
+        img_h_list.append(img_h)
+        img_w_list.append(img_w)
+        
+    elif view_setting.startswith('degree120'):
+        print('we render 120 degree')
+        # +- 60 degree
+        frame_per_cycle = 480
+        max_degree = 60
+        frame_half_cycle = frame_per_cycle // 2
+        if idx%frame_per_cycle < frame_per_cycle/2:
+            rot_Y = -max_degree + (2 * max_degree / frame_half_cycle) * (idx%frame_half_cycle)
+        # rot_Y = (idx % frame_per_60) / float(frame_per_60) * 2 * np.pi
+        else:
+            rot_Y = max_degree - (2 * max_degree / frame_half_cycle) * (idx%frame_half_cycle)
+        
+        # to radian
+        rot_Y = rot_Y * np.pi / 180
+        if rot_Y<0:
+            rot_Y = rot_Y + 2 * np.pi
+        # print('rot_Y: ', rot_Y)
+        extr = visualize_util.calc_free_mv(object_center,
+                                           tar_pos = np.array([0, 0, 2.5]),
+                                           rot_Y = rot_Y,
+                                           rot_X = 0.3 if view_setting.endswith('bird') else 0.,
+                                           global_orient = global_orient if self.body['test'].get('global_orient', False) else None)
+        intr = np.array([[1100, 0, 512], [0, 1100, 512], [0, 0, 1]], np.float32)
+        intr[:2] *= img_scale
+        img_h = int(1024 * img_scale)
+        img_w = int(1024 * img_scale)
+        
+        extr_list.append(extr)
+        intr_list.append(intr)
+        img_h_list.append(img_h)
+        img_w_list.append(img_w)
+    
+    elif view_setting.startswith('degree90'):
+        print('we render 90 degree')
+        # +- 60 degree
+        frame_per_cycle = 360
+        max_degree = 45
+        frame_half_cycle = frame_per_cycle // 2
+        if idx%frame_per_cycle < frame_per_cycle/2:
+            rot_Y = -max_degree + (2 * max_degree / frame_half_cycle) * (idx%frame_half_cycle)
+        # rot_Y = (idx % frame_per_60) / float(frame_per_60) * 2 * np.pi
+        else:
+            rot_Y = max_degree - (2 * max_degree / frame_half_cycle) * (idx%frame_half_cycle)
+        
+        # to radian
+        rot_Y = rot_Y * np.pi / 180
+        if rot_Y<0:
+            rot_Y = rot_Y + 2 * np.pi
+        # print('rot_Y: ', rot_Y)
+        extr = visualize_util.calc_free_mv(object_center,
+                                           tar_pos = np.array([0, 0, 2.5]),
+                                           rot_Y = rot_Y,
+                                           rot_X = 0.3 if view_setting.endswith('bird') else 0.,
+                                           global_orient = global_orient if self.body['test'].get('global_orient', False) else None)
+        intr = np.array([[1100, 0, 512], [0, 1100, 512], [0, 0, 1]], np.float32)
+        intr[:2] *= img_scale
+        img_h = int(1024 * img_scale)
+        img_w = int(1024 * img_scale)
+        
+        extr_list.append(extr)
+        intr_list.append(intr)
+        img_h_list.append(img_h)
+        img_w_list.append(img_w)
+        
+        
+    elif view_setting.startswith('front'):
+        # front view setting
+        extr = visualize_util.calc_free_mv(object_center,
+                                           tar_pos = np.array([0, 0, 2.5]),
+                                           rot_Y = 0.,
+                                           rot_X = 0.3 if view_setting.endswith('bird') else 0.,
+                                           global_orient = global_orient if self.body['test'].get('global_orient', False) else None)
+        intr = np.array([[1100, 0, 512], [0, 1100, 512], [0, 0, 1]], np.float32)
+        intr[:2] *= img_scale
+        img_h = int(1024 * img_scale)
+        img_w = int(1024 * img_scale)
+        
+        extr_list.append(extr)
+        intr_list.append(intr)
+        img_h_list.append(img_h)
+        img_w_list.append(img_w)
+        
+        # print('extr: ', extr)
+        # print('intr: ', intr)
+        # print('img_h: ', img_h)
+        # print('img_w: ', img_w)
+        # exit()
+        
+        
+        
+    elif view_setting.startswith('back'):
+        # back view setting
+        extr = visualize_util.calc_free_mv(object_center,
+                                           tar_pos = np.array([0, 0, 2.5]),
+                                           rot_Y = np.pi,
+                                           rot_X = 0.5 * np.pi / 4. if view_setting.endswith('bird') else 0.,
+                                           global_orient = global_orient if self.body['test'].get('global_orient', False) else None)
+        intr = np.array([[1100, 0, 512], [0, 1100, 512], [0, 0, 1]], np.float32)
+        intr[:2] *= img_scale
+        img_h = int(1024 * img_scale)
+        img_w = int(1024 * img_scale)
+    elif view_setting.startswith('moving'):
+        # moving camera setting
+        extr = visualize_util.calc_free_mv(object_center,
+                                           # tar_pos = np.array([0, 0, 3.0]),
+                                           # rot_Y = -0.3,
+                                           tar_pos = np.array([0, 0, 2.5]),
+                                           rot_Y = 0.,
+                                           rot_X = 0.3 if view_setting.endswith('bird') else 0.,
+                                           global_orient = global_orient if self.body['test'].get('global_orient', False) else None)
+        intr = np.array([[1100, 0, 512], [0, 1100, 512], [0, 0, 1]], np.float32)
+        intr[:2] *= img_scale
+        img_h = int(1024 * img_scale)
+        img_w = int(1024 * img_scale)
+    elif view_setting.startswith('cano'):
+        cano_center = self.dataset.cano_bounds.mean(0)
+        extr = np.identity(4, np.float32)
+        extr[:3, 3] = -cano_center
+        rot_x = np.identity(4, np.float32)
+        rot_x[:3, :3] = cv.Rodrigues(np.array([np.pi, 0, 0], np.float32))[0]
+        extr = rot_x @ extr
+        f_len = 5000
+        extr[2, 3] += f_len / 512
+        intr = np.array([[f_len, 0, 512], [0, f_len, 512], [0, 0, 1]], np.float32)
+        # item = self.dataset.getitem(idx,
+        #                             training = False,
+        #                             extr = extr,
+        #                             intr = intr,
+        #                             img_w = 1024,
+        #                             img_h = 1024)
+        img_w, img_h = 1024, 1024
+        # item['live_smpl_v'] = item['cano_smpl_v']
+        # item['cano2live_jnt_mats'] = torch.eye(4, dtype = torch.float32)[None].expand(item['cano2live_jnt_mats'].shape[0], -1, -1)
+        # item['live_bounds'] = item['cano_bounds']
+    else:
+        raise ValueError('Invalid view setting for animation!')

render_utils/lib/networks/faceverse_torch.py

@@ -0,0 +1,292 @@
+import torch
+from torch import nn
+import numpy as np
+from ..utils.rotation_conversions import axis_angle_to_matrix
+
+
+class FaceVerseModel(nn.Module):
+    def __init__(self, model_dict, batch_size=1, device='cuda:0', **kargs):
+        super(FaceVerseModel, self).__init__()
+
+        self.batch_size = batch_size
+        self.device = torch.device(device)
+
+        self.rotXYZ = torch.eye(3).view(1, 3, 3).repeat(3, 1, 1).view(3, 1, 3, 3).to(self.device)
+
+        self.skinmask = torch.tensor(model_dict['skinmask'], requires_grad=False, device=self.device)
+
+        self.kp_inds = torch.tensor(model_dict['keypoints'].reshape(-1, 1), requires_grad=False).squeeze().long().to(self.device)
+
+        self.meanshape = torch.tensor(model_dict['meanshape'].reshape(1, -1), dtype=torch.float32, requires_grad=False, device=self.device)
+        self.meantex = torch.tensor(model_dict['meantex'].reshape(1, -1), dtype=torch.float32, requires_grad=False, device=self.device)
+
+        self.idBase = torch.tensor(model_dict['idBase'], dtype=torch.float32, requires_grad=False, device=self.device)
+        self.expBase = torch.tensor(model_dict['exBase'], dtype=torch.float32, requires_grad=False, device=self.device)
+        self.texBase = torch.tensor(model_dict['texBase'], dtype=torch.float32, requires_grad=False, device=self.device)
+        # self.expBase[:, :8] *= 4    # 控制眼睛
+        # self.expBase[:, 8:] *= 2    # 控制其它
+
+        self.tri = torch.tensor(model_dict['tri'], dtype=torch.int64, requires_grad=False, device=self.device)
+        # self.tri = self.tri[:, [0, 2, 1]] # TODO
+        self.point_buf = torch.tensor(model_dict['point_buf'], dtype=torch.int64, requires_grad=False, device=self.device)
+
+        self.num_vertex = self.meanshape.shape[0] // 3
+        self.id_dims = self.idBase.shape[1]
+        self.tex_dims = self.texBase.shape[1]
+        self.exp_dims = self.expBase.shape[1]
+        self.all_dims = self.id_dims + self.tex_dims + self.exp_dims
+        self.init_coeff_tensors()
+
+        # for tracking by landmarks
+        self.kp_inds_view = torch.cat([self.kp_inds[:, None] * 3, self.kp_inds[:, None] * 3 + 1, self.kp_inds[:, None] * 3 + 2], dim=1).flatten()
+        self.idBase_view = self.idBase[self.kp_inds_view, :].detach().clone()
+        self.expBase_view = self.expBase[self.kp_inds_view, :].detach().clone()
+        self.meanshape_view = self.meanshape[:, self.kp_inds_view].detach().clone()
+
+        # zxc
+        self.identity = torch.eye(3, dtype=torch.float32, device=self.device)
+
+    def init_coeff_tensors(self, id_coeff=None, tex_coeff=None, exp_coeff=None, gamma_coeff=None, trans_coeff=None, rot_coeff=None, scale_coeff=None):
+        if id_coeff is None:
+            self.id_tensor = torch.zeros(
+                (1, self.id_dims), dtype=torch.float32,
+                requires_grad=True, device=self.device)
+        else:
+            assert id_coeff.shape == (1, self.id_dims)
+            # self.id_tensor = torch.tensor(id_coeff, dtype=torch.float32, requires_grad=True, device=self.device)
+            self.id_tensor = id_coeff.clone().detach().requires_grad_(True)
+
+        if tex_coeff is None:
+            self.tex_tensor = torch.zeros(
+                (1, self.tex_dims), dtype=torch.float32,
+                requires_grad=True, device=self.device)
+        else:
+            assert tex_coeff.shape == (1, self.tex_dims)
+            # self.tex_tensor = torch.tensor(tex_coeff, dtype=torch.float32, requires_grad=True, device=self.device)
+            self.tex_tensor = tex_coeff.clone().detach().requires_grad_(True)
+
+        if exp_coeff is None:
+            self.exp_tensor = torch.zeros(
+                (self.batch_size, self.exp_dims), dtype=torch.float32,
+                requires_grad=True, device=self.device)
+        else:
+            assert exp_coeff.shape == (1, self.exp_dims)
+            # self.exp_tensor = torch.tensor(exp_coeff, dtype=torch.float32, requires_grad=True, device=self.device)
+            self.exp_tensor = exp_coeff.clone().detach().requires_grad_(True)
+
+        if gamma_coeff is None:
+            self.gamma_tensor = torch.zeros(
+                (self.batch_size, 27), dtype=torch.float32,
+                requires_grad=True, device=self.device)
+        else:
+            # self.gamma_tensor = torch.tensor(
+            #     gamma_coeff, dtype=torch.float32,
+            #     requires_grad=True, device=self.device)
+            self.gamma_tensor = gamma_coeff.clone().detach().requires_grad_(True)
+
+        if trans_coeff is None:
+            self.trans_tensor = torch.zeros(
+                (self.batch_size, 3), dtype=torch.float32,
+                requires_grad=True, device=self.device)
+        else:
+            # self.trans_tensor = torch.tensor(
+            #     trans_coeff, dtype=torch.float32,
+            #     requires_grad=True, device=self.device)
+            self.trans_tensor = trans_coeff.clone().detach().requires_grad_(True)
+
+        if scale_coeff is None:
+            self.scale_tensor = 0.18 * torch.ones((self.batch_size, 1), dtype=torch.float32, device=self.device)
+            # self.scale_tensor = torch.ones((self.batch_size, 1), dtype=torch.float32, device=self.device)
+            self.scale_tensor.requires_grad_(True)
+            # self.scale_tensor = torch.ones(
+            #     (self.batch_size, 1), dtype=torch.float32,
+            #     requires_grad=True, device=self.device)
+        else:
+            # self.scale_tensor = torch.tensor(
+            #     scale_coeff, dtype=torch.float32,
+            #     requires_grad=True, device=self.device)
+            self.scale_tensor = scale_coeff.clone().detach().requires_grad_(True)
+
+        if rot_coeff is None:
+            self.rot_tensor = torch.zeros(
+                (self.batch_size, 3), dtype=torch.float32,
+                requires_grad=True, device=self.device)
+        else:
+            # self.rot_tensor = torch.tensor(
+            #     rot_coeff, dtype=torch.float32,
+            #     requires_grad=True, device=self.device)
+            self.rot_tensor = rot_coeff.clone().detach().requires_grad_(True)
+
+    def get_lms(self, vs):
+        lms = vs[:, self.kp_inds, :]
+        return lms
+
+    def split_coeffs(self, coeffs):
+        id_coeff = coeffs[:, :self.id_dims]  # identity(shape) coeff
+        exp_coeff = coeffs[:, self.id_dims:self.id_dims + self.exp_dims]  # expression coeff
+        tex_coeff = coeffs[:, self.id_dims + self.exp_dims:self.all_dims]  # texture(albedo) coeff
+        angles = coeffs[:, self.all_dims:self.all_dims + 3]  # ruler angles(x,y,z) for rotation of dim 3
+        gamma = coeffs[:, self.all_dims + 3:self.all_dims + 30]  # lighting coeff for 3 channel SH function of dim 27
+        translation = coeffs[:, self.all_dims + 30:-1]  # translation coeff of dim 3
+        scale = coeffs[:, -1:]
+        return id_coeff, exp_coeff, tex_coeff, angles, gamma, translation, scale
+
+    def merge_coeffs(self, id_coeff, exp_coeff, tex_coeff, angles, gamma, translation, scale):
+        coeffs = torch.cat([id_coeff, exp_coeff, tex_coeff, angles, gamma, translation, scale], dim=1)
+        return coeffs
+
+    def get_packed_tensors(self):
+        return self.merge_coeffs(self.id_tensor.repeat(self.batch_size, 1),
+                                 self.exp_tensor,
+                                 self.tex_tensor.repeat(self.batch_size, 1),
+                                 self.rot_tensor, self.gamma_tensor,
+                                 self.trans_tensor, self.scale_tensor)
+
+    def forward(self, coeffs=None, camT=None):
+        if coeffs is None:
+            id_coeff = self.id_tensor.repeat(self.batch_size, 1)
+            tex_coeff = self.tex_tensor.repeat(self.batch_size, 1)
+            exp_coeff, angles, gamma = self.exp_tensor, self.rot_tensor, self.gamma_tensor
+            translation, scale = self.rot_tensor, self.scale_tensor
+        else:
+            id_coeff, exp_coeff, tex_coeff, angles, gamma, translation, scale = self.split_coeffs(coeffs)
+        rotation = axis_angle_to_matrix(angles)
+
+        if camT is not None:
+            rotation2 = camT[:3, :3].reshape(1, 3, 3)
+            translation2 = camT[:3, 3:].reshape(1, 3, 1)
+            if torch.allclose(rotation2, self.identity):
+                translation = translation + translation2
+            else:
+                rotation = torch.matmul(rotation2, rotation)
+                translation = torch.matmul(rotation2, translation) + translation2
+
+        vs = self.get_vs(id_coeff, exp_coeff)
+        vs_t = self.rigid_transform(vs, rotation, translation, torch.abs(scale))
+
+        lms_t = self.get_lms(vs_t)
+
+        face_texture = self.get_color(tex_coeff)
+        face_norm = self.compute_norm(vs, self.tri, self.point_buf)
+        face_norm_r = face_norm.bmm(rotation)
+        face_color = self.add_illumination(face_texture, face_norm_r, gamma)
+
+        return {'v': vs_t, 'lm': lms_t, 'face_texture': face_texture, 'face_color': face_color}
+
+    def forward_landmarks(self, coeffs=None, camT=None):
+        if coeffs is None:
+            id_coeff = self.id_tensor.repeat(self.batch_size, 1)
+            tex_coeff = self.tex_tensor.repeat(self.batch_size, 1)
+            exp_coeff, angles, gamma = self.exp_tensor, self.rot_tensor, self.gamma_tensor
+            translation, scale = self.trans_tensor, self.scale_tensor
+        else:
+            id_coeff, exp_coeff, tex_coeff, angles, gamma, translation, scale = self.split_coeffs(coeffs)
+        rotation = axis_angle_to_matrix(angles)
+
+        if camT is not None:
+            rotation2 = camT[:3, :3].reshape(1, 3, 3)
+            translation2 = camT[:3, 3:].reshape(1, 3, 1)
+            if torch.allclose(rotation2, self.identity):
+                translation = translation + translation2
+            else:
+                rotation = torch.matmul(rotation2, rotation)
+                translation = torch.matmul(rotation2, translation) + translation2
+        lms = self.get_vs_lms(id_coeff, exp_coeff)
+        lms_t = self.rigid_transform(lms, rotation, translation, torch.abs(scale))
+        return lms_t
+
+    def get_vs(self, id_coeff, exp_coeff):
+        face_shape = torch.einsum('ij,aj->ai', self.idBase, id_coeff) + \
+                     torch.einsum('ij,aj->ai', self.expBase, exp_coeff) + self.meanshape
+        face_shape = face_shape.view(self.batch_size, -1, 3)
+        return face_shape
+
+    def get_vs_lms(self, id_coeff, exp_coeff):
+        face_shape = torch.einsum('ij,aj->ai', self.idBase_view, id_coeff) + \
+                     torch.einsum('ij,aj->ai', self.expBase_view, exp_coeff) + self.meanshape_view
+        face_shape = face_shape.view(self.batch_size, -1, 3)
+        return face_shape
+
+    def get_color(self, tex_coeff):
+        face_texture = torch.einsum('ij,aj->ai', self.texBase, tex_coeff) + self.meantex
+        face_texture = face_texture.view(self.batch_size, -1, 3)
+        return face_texture
+
+    def get_skinmask(self):
+        return self.skinmask
+
+    def compute_norm(self, vs, tri, point_buf):
+        face_id = tri
+        point_id = point_buf
+        v1 = vs[:, face_id[:, 0], :]
+        v2 = vs[:, face_id[:, 1], :]
+        v3 = vs[:, face_id[:, 2], :]
+        e1 = v1 - v2
+        e2 = v2 - v3
+        face_norm = e1.cross(e2)
+
+        v_norm = face_norm[:, point_id, :].sum(2)
+        v_norm = v_norm / (v_norm.norm(dim=2).unsqueeze(2) + 1e-9)
+
+        return v_norm
+
+    def add_illumination(self, face_texture, norm, gamma):
+        gamma = gamma.view(-1, 3, 9).clone()
+        gamma[:, :, 0] += 0.8
+        gamma = gamma.permute(0, 2, 1)
+
+        a0 = np.pi
+        a1 = 2 * np.pi / np.sqrt(3.0)
+        a2 = 2 * np.pi / np.sqrt(8.0)
+        c0 = 1 / np.sqrt(4 * np.pi)
+        c1 = np.sqrt(3.0) / np.sqrt(4 * np.pi)
+        c2 = 3 * np.sqrt(5.0) / np.sqrt(12 * np.pi)
+        d0 = 0.5 / np.sqrt(3.0)
+
+        norm = norm.view(-1, 3)
+        nx, ny, nz = norm[:, 0], norm[:, 1], norm[:, 2]
+        arrH = []
+
+        arrH.append(a0 * c0 * (nx * 0 + 1))
+        arrH.append(-a1 * c1 * ny)
+        arrH.append(a1 * c1 * nz)
+        arrH.append(-a1 * c1 * nx)
+        arrH.append(a2 * c2 * nx * ny)
+        arrH.append(-a2 * c2 * ny * nz)
+        arrH.append(a2 * c2 * d0 * (3 * nz.pow(2) - 1))
+        arrH.append(-a2 * c2 * nx * nz)
+        arrH.append(a2 * c2 * 0.5 * (nx.pow(2) - ny.pow(2)))
+
+        H = torch.stack(arrH, 1)
+        Y = H.view(self.batch_size, face_texture.shape[1], 9)
+        lighting = Y.bmm(gamma)
+
+        face_color = face_texture * lighting
+        return face_color
+
+    def rigid_transform(self, vs, rot, trans, scale):
+        scale = scale.reshape(-1, 1, 1)
+        vs_r = torch.matmul(vs * scale, rot.permute(0, 2, 1))
+        vs_t = vs_r + trans.reshape(-1, 1, 3)
+        return vs_t
+
+    def get_rot_tensor(self):
+        return self.rot_tensor
+
+    def get_trans_tensor(self):
+        return self.trans_tensor
+
+    def get_exp_tensor(self):
+        return self.exp_tensor
+
+    def get_tex_tensor(self):
+        return self.tex_tensor
+
+    def get_id_tensor(self):
+        return self.id_tensor
+
+    def get_gamma_tensor(self):
+        return self.gamma_tensor
+
+    def get_scale_tensor(self):
+        return self.scale_tensor

render_utils/lib/networks/smpl_torch.py

@@ -0,0 +1,341 @@
+import torch
+import torch.nn as nn
+import pickle
+import numpy as np
+import scipy.sparse
+import os
+
+import smplx
+import smplx.lbs
+
+from ..utils.geometry import rodrigues
+def _convert_to_sparse(mat):
+    if isinstance(mat, np.ndarray):
+        row_ind, col_ind = np.asarray(mat > 0).nonzero()
+        data = [mat[r, c] for (r, c) in zip(row_ind, col_ind)]
+        mat_sparse = scipy.sparse.csr_matrix((data, (row_ind, col_ind)), shape=mat.shape)
+        return mat_sparse
+    else:
+        return mat
+
+
+class SmplTorch(nn.Module):
+    def __init__(self, model_file='./data/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl'):
+        super(SmplTorch, self).__init__()
+        if model_file.endswith('.pkl'):
+            with open(model_file, 'rb') as f:
+                smpl_model = pickle.load(f, encoding='iso-8859-1')
+            # print('SmplTorch: Loading SMPL data from %s' % model_file)
+        elif model_file.endswith('.npz'):
+            smpl_model = np.load(model_file)
+            # print('SmplTorch: Loading SMPL data from %s' % model_file)
+        else:
+            raise ValueError('Unknown file extension: %s' % model_file)
+
+        J_regressor = _convert_to_sparse(smpl_model['J_regressor']).tocoo()
+        row = J_regressor.row
+        col = J_regressor.col
+        data = J_regressor.data
+        i = torch.LongTensor([row, col])
+        v = torch.FloatTensor(data)
+        J_regressor_shape = J_regressor.shape
+        self.v_num = len(smpl_model['v_template'])
+        self.joint_num = smpl_model['weights'].shape[1]
+        self.beta_dim = 10
+
+        self.register_buffer('J_regressor', torch.sparse.FloatTensor(i, v, J_regressor_shape).to_dense())
+        self.register_buffer('weights', torch.FloatTensor(smpl_model['weights']))
+        self.register_buffer('posedirs', torch.FloatTensor(smpl_model['posedirs']))
+        self.register_buffer('v_template', torch.FloatTensor(smpl_model['v_template']))
+        self.register_buffer('shapedirs', torch.FloatTensor(np.array(smpl_model['shapedirs'][:, :, :self.beta_dim])))
+        self.register_buffer('faces', torch.from_numpy(smpl_model['f'].astype(np.int64)))
+        self.register_buffer('kintree_table', torch.from_numpy(smpl_model['kintree_table'].astype(np.int64)))
+        id_to_col = {self.kintree_table[1, i].item(): i for i in range(self.kintree_table.shape[1])}
+        self.register_buffer('parent', torch.LongTensor([id_to_col[self.kintree_table[0, it].item()] for it in range(1, self.kintree_table.shape[1])]))
+
+        self.pose_shape = [self.joint_num, 3]
+        self.beta_shape = [self.beta_dim]
+        self.translation_shape = [3]
+
+        self.pose = torch.zeros(self.pose_shape)
+        self.beta = torch.zeros(self.beta_shape)
+        self.translation = torch.zeros(self.translation_shape)
+
+        self.verts = None
+
+        # # BODY25 joint regressor
+        # with open(os.path.join(os.path.dirname(model_file), 'J_regressor_body25.pkl'), 'rb') as f:
+        #     jreg_data = pickle.load(f, encoding='iso-8859-1')
+        #     J_regressor25 = jreg_data['J_regressor_body25'].tocoo()
+        # row = J_regressor25.row
+        # col = J_regressor25.col
+        # data = J_regressor25.data
+        # i = torch.LongTensor([row, col])
+        # v = torch.FloatTensor(data)
+        # J_regressor_shape = J_regressor25.shape
+        # self.register_buffer('J_regressor_body25', torch.sparse.FloatTensor(i, v, J_regressor_shape).to_dense())
+
+    def forward(self, pose, beta, apply_pose_blend_shape=True):
+        device = pose.device
+        batch_size = pose.shape[0]
+        v_template = self.v_template[None, :]
+        shapedirs = self.shapedirs.view(-1, self.beta_dim)[None, :].expand(batch_size, -1, -1)
+        beta = beta[:, :, None]
+        v_shaped = torch.matmul(shapedirs, beta).view(-1, self.v_num, 3) + v_template
+        # batched sparse matmul not supported in pytorch
+        J = []
+        for i in range(batch_size):
+            J.append(torch.matmul(self.J_regressor, v_shaped[i]))
+        J = torch.stack(J, dim=0)
+        # input it rotmat: (bs,24,3,3)
+        if pose.ndimension() == 4:
+            R = pose
+        # input it rotmat: (bs,72)
+        elif pose.ndimension() == 2:
+            pose_cube = pose.view(-1, 3)  # (batch_size * 24, 1, 3)
+            R = rodrigues(pose_cube).view(batch_size, self.joint_num, 3, 3)
+            R = R.view(batch_size, self.joint_num, 3, 3)
+
+        if apply_pose_blend_shape:
+            I_cube = torch.eye(3)[None, None, :].to(device)
+            # I_cube = torch.eye(3)[None, None, :].expand(theta.shape[0], R.shape[1]-1, -1, -1)
+            lrotmin = (R[:, 1:, :] - I_cube).view(batch_size, -1)
+            posedirs = self.posedirs.view(-1, 9*(self.joint_num-1))[None, :].expand(batch_size, -1, -1)
+            v_posed = v_shaped + torch.matmul(posedirs, lrotmin[:, :, None]).view(-1, self.v_num, 3)
+        else:
+            v_posed = v_shaped
+
+        J_ = J.clone()
+        J_[:, 1:, :] = J[:, 1:, :] - J[:, self.parent, :]
+        G_ = torch.cat([R, J_[:, :, :, None]], dim=-1)
+        pad_row = torch.FloatTensor([0, 0, 0, 1]).to(device).view(1, 1, 1, 4).expand(batch_size, self.joint_num, -1, -1)
+        G_ = torch.cat([G_, pad_row], dim=2)
+        G = [G_[:, 0].clone()]
+        for i in range(1, self.joint_num):
+            G.append(torch.matmul(G[self.parent[i - 1]], G_[:, i, :, :]))
+        G = torch.stack(G, dim=1)
+
+        rest = torch.cat([J, torch.zeros(batch_size, self.joint_num, 1).to(device)], dim=2).view(batch_size, self.joint_num, 4, 1)
+        zeros = torch.zeros(batch_size, self.joint_num, 4, 3).to(device)
+        rest = torch.cat([zeros, rest], dim=-1)
+        rest = torch.matmul(G, rest)
+        G = G - rest
+        T = torch.matmul(self.weights, G.permute(1, 0, 2, 3).contiguous().view(self.joint_num, -1)).view(self.v_num, batch_size, 4, 4).transpose(0, 1)
+        rest_shape_h = torch.cat([v_posed, torch.ones_like(v_posed)[:, :, [0]]], dim=-1)
+        v = torch.matmul(T, rest_shape_h[:, :, :, None])[:, :, :3, 0]
+
+        return v, {'J': J, 'G': G, 'T': T, 'R': R, 'v_shaped': v_shaped, 'v_posed': v_posed}
+
+    def get_joints(self, vertices):
+        """
+        This method is used to get the joint locations from the SMPL mesh
+        Input:
+            vertices: size = (B, self.v_num, 3)
+        Output:
+            3D joints: size = (B, 38, 3)
+        """
+        joints = torch.einsum('bik,ji->bjk', vertices, self.J_regressor)
+        return joints
+
+    def get_root(self, vertices):
+        """
+        This method is used to get the root locations from the SMPL mesh
+        Input:
+            vertices: size = (B, self.v_num, 3)
+        Output:
+            3D joints: size = (B, 1, 3)
+        """
+        joints = torch.einsum('bik,ji->bjk', vertices, self.J_regressor)
+        return joints[:, 0:1, :]
+
+    # def get_joints_body25(self, vertices):
+    #     joints = torch.einsum('bik,ji->bjk', vertices, self.J_regressor_body25)
+    #     return joints
+
+
+class SmplProjector(nn.Module):
+    def __init__(self, verts_cano=None, face_triangles=None):
+        super(SmplProjector, self).__init__()
+        self.verts_cano = verts_cano
+        self.face_triangles = face_triangles
+
+        face_node_idx = np.loadtxt('./data/smplx/face_coarse_level_ids.txt').astype(np.int64)
+        node_to_face_table = np.loadtxt('./data/smplx/face_coarse_level_to_all_table.txt').astype(np.int64)
+        self.register_buffer('face_node_idx', torch.from_numpy(face_node_idx), persistent=False)
+        self.register_buffer('node_to_face_table', torch.from_numpy(node_to_face_table), persistent=False)
+
+    @staticmethod
+    def batch_select(tensor, index, dim):
+        """
+        Perform index_select for batched inputs, where the index tensors of different items are different
+        :param tensor:  [B, G, ....]
+        :param index: [B, N], which is the key difference from torch.index_select()
+        :param dim: int
+        :return:
+        """
+        res = []
+        for bi in range(tensor.shape[0]):
+            res.append(
+                torch.index_select(tensor[bi], dim=dim-1, index=index[bi])
+            )
+        res = torch.stack(res, dim=0)
+        return res
+
+    def get_current_beta(self):
+        return self.density_func.get_beta()
+
+    def calculate_nearest_barycentric(self, verts, pts):
+        nearest_face_id, nearest_dist, nearest_point_barycentric = calc_nearest_barycentric_coord(
+            verts, self.face_triangles, pts, self.face_node_idx, self.node_to_face_table)
+        return nearest_face_id, nearest_dist, nearest_point_barycentric
+
+    def interpolate_vert_feat(self, vert_feat, pts, nearest_face_id, nearest_dist, nearest_point_barycentric):
+        fvfa = torch.index_select(vert_feat, dim=1, index=self.face_triangles[:, 0])  # [B, G, 3]
+        fvfb = torch.index_select(vert_feat, dim=1, index=self.face_triangles[:, 1])  # [B, G, 3]
+        fvfc = torch.index_select(vert_feat, dim=1, index=self.face_triangles[:, 2])  # [B, G, 3]
+        fvfs = torch.cat([fvfa.unsqueeze(2), fvfb.unsqueeze(2), fvfc.unsqueeze(2)], dim=2)  # [B, G, 3, 3]
+        nearest_fvfs = self.batch_select(fvfs, nearest_face_id, dim=1)
+        nearest_vf = torch.sum(nearest_point_barycentric.unsqueeze(-1)*nearest_fvfs, dim=-2)      # [B, N, 3]
+        return nearest_vf
+
+    def calculate_nearest_point(self, verts, vert_nmls, pts, nearest_face_id, nearest_dist, nearest_point_barycentric):
+        verts_cano = self.verts_cano.unsqueeze(0).expand(verts.shape[0], -1, -1)
+        nearest_v = self.interpolate_vert_feat(verts, pts, nearest_face_id, nearest_dist, nearest_point_barycentric)
+        nearest_nml = self.interpolate_vert_feat(vert_nmls, pts, nearest_face_id, nearest_dist, nearest_point_barycentric)
+        nearest_vc = self.interpolate_vert_feat(verts_cano, pts, nearest_face_id, nearest_dist, nearest_point_barycentric)
+        return nearest_v, nearest_nml, nearest_vc
+
+    def forward(self, verts, vert_nmls, pts, extra_vert_feat=None):
+        nearest_face_id, nearest_dist, nearest_point_barycentric = self.calculate_nearest_barycentric(verts, pts)
+        nearest_v, nearest_nml, nearest_vc = self.calculate_nearest_point(verts, vert_nmls, pts, nearest_face_id, nearest_dist, nearest_point_barycentric)
+        nearest_dist = nearest_dist.unsqueeze(-1)
+        nearest_h = torch.sum((pts-nearest_v)*nearest_nml, dim=-1, keepdim=True)    # [B, N, 1]
+        sdf = nearest_h.sign() * nearest_dist
+        if extra_vert_feat is None:
+            return sdf, nearest_v, nearest_nml, nearest_vc
+        else:
+            if len(extra_vert_feat.shape) == 2:
+                extra_vert_feat = extra_vert_feat.unsqueeze(0).expand(verts.shape[0], -1, -1)
+            nearest_feat = self.interpolate_vert_feat(extra_vert_feat, pts, nearest_face_id, nearest_dist, nearest_point_barycentric)
+            return sdf, nearest_v, nearest_nml, nearest_vc, nearest_feat
+
+
+class SmplProjectorBF(nn.Module):
+    """
+    Calculate barycentric projection in a brute-force manner
+    """
+    def __init__(self, verts_cano=None, face_triangles=None):
+        super(SmplProjectorBF, self).__init__()
+        self.verts_cano = verts_cano
+        self.face_triangles = face_triangles
+
+    @staticmethod
+    def batch_select(tensor, index, dim):
+        """
+        Perform index_select for batched inputs, where the index tensors of different items are different
+        :param tensor:  [B, G, ....]
+        :param index: [B, N], which is the key difference from torch.index_select()
+        :param dim: int
+        :return:
+        """
+        res = []
+        for bi in range(tensor.shape[0]):
+            res.append(
+                torch.index_select(tensor[bi], dim=dim-1, index=index[bi])
+            )
+        res = torch.stack(res, dim=0)
+        return res
+
+    def get_current_beta(self):
+        return self.density_func.get_beta()
+
+    def calculate_nearest_barycentric(self, verts, pts):
+        nearest_face_id, nearest_dist, nearest_point_barycentric = calc_nearest_barycentric_coord_bf(
+            verts, self.face_triangles, pts)
+        return nearest_face_id, nearest_dist, nearest_point_barycentric
+
+    def interpolate_vert_feat(self, vert_feat, pts, nearest_face_id, nearest_dist, nearest_point_barycentric):
+        fvfa = torch.index_select(vert_feat, dim=1, index=self.face_triangles[:, 0])  # [B, G, 3]
+        fvfb = torch.index_select(vert_feat, dim=1, index=self.face_triangles[:, 1])  # [B, G, 3]
+        fvfc = torch.index_select(vert_feat, dim=1, index=self.face_triangles[:, 2])  # [B, G, 3]
+        fvfs = torch.cat([fvfa.unsqueeze(2), fvfb.unsqueeze(2), fvfc.unsqueeze(2)], dim=2)  # [B, G, 3, 3]
+        nearest_fvfs = self.batch_select(fvfs, nearest_face_id, dim=1)
+        nearest_vf = torch.sum(nearest_point_barycentric.unsqueeze(-1)*nearest_fvfs, dim=-2)      # [B, N, 3]
+        return nearest_vf
+
+    def calculate_nearest_point(self, verts, vert_nmls, pts, nearest_face_id, nearest_dist, nearest_point_barycentric):
+        verts_cano = self.verts_cano.unsqueeze(0).expand(verts.shape[0], -1, -1)
+        nearest_v = self.interpolate_vert_feat(verts, pts, nearest_face_id, nearest_dist, nearest_point_barycentric)
+        nearest_nml = self.interpolate_vert_feat(vert_nmls, pts, nearest_face_id, nearest_dist, nearest_point_barycentric)
+        nearest_vc = self.interpolate_vert_feat(verts_cano, pts, nearest_face_id, nearest_dist, nearest_point_barycentric)
+        return nearest_v, nearest_nml, nearest_vc
+
+    def forward(self, verts, vert_nmls, pts):
+        nearest_face_id, nearest_dist, nearest_point_barycentric = self.calculate_nearest_barycentric(verts, pts)
+        nearest_v, nearest_nml, nearest_vc = self.calculate_nearest_point(verts, vert_nmls, pts, nearest_face_id, nearest_dist, nearest_point_barycentric)
+        nearest_dist = nearest_dist.unsqueeze(-1)
+        nearest_h = torch.sum((pts-nearest_v)*nearest_nml, dim=-1, keepdim=True)    # [B, N, 1]
+        sdf = nearest_h.sign() * nearest_dist
+        return sdf, nearest_v, nearest_nml, nearest_vc
+
+
+class SmplJointRegressor(nn.Module):
+    def __init__(self, model_file, extra_regressor_file):
+        super().__init__()
+        # '../smpl_models/J_regressor_extra_' + 'smplx' + '.npy'
+
+        self.smpl = smplx.SMPLX(model_path=model_file, gender='neutral', use_pca=False,
+                                num_pca_comps=45, flat_hand_mean=True, batch_size=1)
+        self.extra_jregressor = torch.tensor(np.load(extra_regressor_file), dtype=torch.float32)
+        self.smpl2coco = {
+            'body': torch.LongTensor([55, 57, 56, 59, 58, 16, 17, 18, 19, 20, 21, 128, 127, 4, 5, 7, 8]),
+            'foot': torch.LongTensor(np.arange(60, 66, dtype=int)),
+            'face': torch.LongTensor(np.arange(76, 127, dtype=int)),
+            'lhand': torch.LongTensor(
+                [20, 37, 38, 39, 66, 25, 26, 27, 67, 28, 29, 30, 68, 34, 35, 36, 69, 31, 32, 33, 70]),
+            'rhand': torch.LongTensor(
+                [21, 52, 53, 54, 71, 40, 41, 42, 72, 43, 44, 45, 73, 49, 50, 51, 74, 46, 47, 48, 75])
+        }
+
+    def to(self, device):
+        super().to(device)
+        self.smpl = self.smpl.to(device)
+        self.smpl2coco = {k: v.to(device) for k, v in self.smpl2coco.items()}
+        self.extra_jregressor = self.extra_jregressor.to(device)
+        return self
+
+    def get_all_joints(self, verts):
+        lmk_faces_idx = self.smpl.lmk_faces_idx.unsqueeze(dim=0).expand(verts.shape[0], -1).contiguous()
+        lmk_bary_coords = self.smpl.lmk_bary_coords.unsqueeze(dim=0).repeat(verts.shape[0], 1, 1)
+        landmarks = smplx.lbs.vertices2landmarks(verts, self.smpl.faces_tensor, lmk_faces_idx, lmk_bary_coords)
+
+        joints = torch.einsum('bik,ji->bjk', verts, self.smpl.J_regressor)
+        joints = self.smpl.vertex_joint_selector(verts, joints)
+        joints = torch.cat([joints, landmarks], dim=1)
+        if self.smpl.joint_mapper is not None:
+            joints = self.joint_mapper(joints=joints, vertices=verts)
+
+        extra_joints = torch.einsum('bik,ji->bjk', verts, self.extra_jregressor)
+        joints = torch.cat([joints, extra_joints], dim=1)
+
+        return joints
+
+    def forward(self, verts, cam_param=None):
+        joints = self.get_all_joints(verts)
+
+        if cam_param:
+            if len(cam_param['K'].shape) == 3:
+                smpl_j2d = torch.einsum("ljk,lik->lij", cam_param['K'], torch.einsum(
+                    "ijk,lk->ilj", cam_param['R'], joints.reshape(-1, 3)) + cam_param['T'].reshape(-1, 1, 3))
+            else:
+                smpl_j2d = torch.einsum("jk,lik->lij", cam_param['K'], torch.einsum(
+                    "ijk,lk->ilj", cam_param['R'], joints.reshape(-1, 3)) + cam_param['T'].reshape(-1, 1, 3))
+            valid = smpl_j2d[:, :, 2] > 1e-3
+            smpl_j2d[:, :, :2][valid] = smpl_j2d[:, :, :2][valid] / smpl_j2d[:, :, 2:3][valid]
+            smpl_j2d[:, :, 2] = valid.float()
+            skel_2d = {k: torch.index_select(smpl_j2d, dim=1, index=v) for k, v in self.smpl2coco.items()}
+        else:
+            skel_2d = None
+
+        skel = {k: torch.index_select(joints, dim=1, index=v) for k, v in self.smpl2coco.items()}
+        return skel, skel_2d

render_utils/lib/utils/gaussian_np_utils.py

@@ -0,0 +1,162 @@
+import numpy as np
+import torch
+from plyfile import PlyData, PlyElement
+from typing import NamedTuple
+import smplx
+import tqdm
+import cv2 as cv
+import os
+
+from scipy.spatial.transform import Rotation as R
+
+
+class GaussianAttributes(NamedTuple):
+    xyz: np.ndarray
+    opacities: np.ndarray
+    features_dc: np.ndarray
+    features_extra: np.ndarray
+    scales: np.ndarray
+    rot: np.ndarray
+
+def load_gaussians_from_ply(path):
+    max_sh_degree = 3
+    plydata = PlyData.read(path)
+
+    xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
+                    np.asarray(plydata.elements[0]["y"]),
+                    np.asarray(plydata.elements[0]["z"])), axis=1)
+    opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
+
+    features_dc = np.zeros((xyz.shape[0], 3, 1))
+    features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
+    features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
+    features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
+
+    extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
+    extra_f_names = sorted(extra_f_names, key=lambda x: int(x.split('_')[-1]))
+    assert len(extra_f_names) == 3 * (max_sh_degree + 1) ** 2 - 3
+    features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
+    for idx, attr_name in enumerate(extra_f_names):
+        features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
+    # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
+    features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1))
+
+    scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
+    scale_names = sorted(scale_names, key=lambda x: int(x.split('_')[-1]))
+    scales = np.zeros((xyz.shape[0], len(scale_names)))
+    for idx, attr_name in enumerate(scale_names):
+        scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+    rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
+    rot_names = sorted(rot_names, key=lambda x: int(x.split('_')[-1]))
+    rots = np.zeros((xyz.shape[0], len(rot_names)))
+    for idx, attr_name in enumerate(rot_names):
+        rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+    return GaussianAttributes(xyz, opacities, features_dc, features_extra, scales, rots)
+
+
+def construct_list_of_attributes(_features_dc, _features_rest, _scaling, _rotation):
+    l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+    # All channels except the 3 DC
+    for i in range(_features_dc.shape[1] * _features_dc.shape[2]):
+        l.append('f_dc_{}'.format(i))
+    for i in range(_features_rest.shape[1] * _features_rest.shape[2]):
+        l.append('f_rest_{}'.format(i))
+    l.append('opacity')
+    for i in range(_scaling.shape[1]):
+        l.append('scale_{}'.format(i))
+    for i in range(_rotation.shape[1]):
+        l.append('rot_{}'.format(i))
+    return l
+
+
+def select_gaussians(gaussian_attrs, select_mask_or_idx):
+    return GaussianAttributes(
+        xyz=gaussian_attrs.xyz[select_mask_or_idx],
+        opacities=gaussian_attrs.opacities[select_mask_or_idx],
+        features_dc=gaussian_attrs.features_dc[select_mask_or_idx],
+        features_extra=gaussian_attrs.features_extra[select_mask_or_idx],
+        scales=gaussian_attrs.scales[select_mask_or_idx],
+        rot=gaussian_attrs.rot[select_mask_or_idx]
+    )
+
+
+def combine_gaussians(gaussian_attrs_list):
+    return GaussianAttributes(
+        xyz=np.concatenate([gau.xyz for gau in gaussian_attrs_list], axis=0),
+        opacities=np.concatenate([gau.opacities for gau in gaussian_attrs_list], axis=0),
+        features_dc=np.concatenate([gau.features_dc for gau in gaussian_attrs_list], axis=0),
+        features_extra=np.concatenate([gau.features_extra for gau in gaussian_attrs_list], axis=0),
+        scales=np.concatenate([gau.scales for gau in gaussian_attrs_list], axis=0),
+        rot=np.concatenate([gau.rot for gau in gaussian_attrs_list], axis=0),
+    )
+
+
+def apply_transformation_to_gaussians(gaussian_attrs, spatial_transformation, color_transformation=None):
+    xyzs = np.copy(gaussian_attrs.xyz)
+    xyzs = np.matmul(xyzs, spatial_transformation[:3, :3].transpose()) + spatial_transformation[:3, 3].reshape([1, 3])
+
+    gaussian_rotmats = R.from_quat(gaussian_attrs.rot[:, (1, 2, 3, 0)]).as_matrix()
+    new_rots = []
+    for rotmat in gaussian_rotmats:
+        rotmat = np.matmul(spatial_transformation[:3, :3], rotmat)
+        rotq = R.from_matrix(rotmat).as_quat()
+        rotq = np.array([rotq[3], rotq[0], rotq[1], rotq[2]])
+        new_rots.append(rotq)
+    new_rots = np.stack(new_rots, axis=0)
+    if color_transformation is not None:
+        if color_transformation.shape[0] == 3 and color_transformation.shape[1] == 3:
+            new_clrs = np.matmul(gaussian_attrs.features_dc[:, :, 0], color_transformation)[:, :, np.newaxis]
+        elif color_transformation.shape[0] == 4 and color_transformation.shape[1] == 4:
+            clrs = gaussian_attrs.features_dc[:, :, 0]
+            clrs = np.concatenate([clrs, np.ones_like(clrs[:, :1])], axis=1)
+            new_clrs = np.matmul(clrs, color_transformation)
+            new_clrs = new_clrs[:, :3, np.newaxis]
+    else:
+        new_clrs = gaussian_attrs.features_dc
+
+    return GaussianAttributes(
+        xyz=xyzs,
+        opacities=gaussian_attrs.opacities,
+        features_dc=new_clrs,
+        features_extra=gaussian_attrs.features_extra,
+        scales=gaussian_attrs.scales,
+        rot=new_rots,
+    )
+
+
+def update_gaussian_attributes(
+        orig_gaussian,
+        new_xyz=None, new_rgb=None, new_rot=None, new_opacity=None, new_scale=None):
+    return GaussianAttributes(
+        xyz=orig_gaussian.xyz if new_xyz is None else new_xyz,
+        opacities=orig_gaussian.opacities if new_opacity is None else new_opacity,
+        features_dc=orig_gaussian.features_dc if new_rgb is None else new_rgb,
+        features_extra=orig_gaussian.features_extra,
+        scales=orig_gaussian.scales if new_scale is None else new_scale,
+        rot=orig_gaussian.rot if new_rot is None else new_rot,
+    )
+
+
+def save_gaussians_as_ply(path, gaussian_attrs: GaussianAttributes):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    xyz = gaussian_attrs.xyz
+    normals = np.zeros_like(xyz)
+    features_dc = gaussian_attrs.features_dc
+    features_rest = gaussian_attrs.features_extra
+    opacities = gaussian_attrs.opacities
+    scale = gaussian_attrs.scales
+    rotation = gaussian_attrs.rot
+
+    dtype_full = [(attribute, 'f4') for attribute in construct_list_of_attributes(features_dc, features_rest, scale, rotation)]
+
+    elements = np.empty(xyz.shape[0], dtype = dtype_full)
+    attributes = np.concatenate((xyz, normals, features_dc.reshape(features_dc.shape[0], -1),
+                                 features_rest.reshape(features_rest.shape[0], -1),
+                                 opacities, scale, rotation), axis=1)
+    elements[:] = list(map(tuple, attributes))
+    el = PlyElement.describe(elements, 'vertex')
+    PlyData([el]).write(path)
+    return

render_utils/lib/utils/geometry.py

@@ -0,0 +1,517 @@
+import glob
+import os
+import numpy as np
+import torch
+import torch.nn.functional as F
+import trimesh
+import scipy.sparse as sp
+import collections
+
+
+def rodrigues(theta):
+    """Convert axis-angle representation to rotation matrix.
+    Args:
+        theta: size = [B, 3]
+    Returns:
+        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
+    """
+    l1norm = torch.norm(theta + 1e-8, p=2, dim=1)
+    angle = torch.unsqueeze(l1norm, -1)
+    normalized = torch.div(theta, angle)
+    angle = angle * 0.5
+    v_cos = torch.cos(angle)
+    v_sin = torch.sin(angle)
+    quat = torch.cat([v_cos, v_sin * normalized], dim=1)
+    return quat2mat(quat)
+
+
+def quat2mat(quat):
+    """Convert quaternion coefficients to rotation matrix.
+    Args:
+        quat: size = [B, 4] 4 <===>(w, x, y, z)
+    Returns:
+        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
+    """
+    norm_quat = quat
+    norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True)
+    w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, 2], norm_quat[:, 3]
+
+    B = quat.size(0)
+
+    w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
+    wx, wy, wz = w * x, w * y, w * z
+    xy, xz, yz = x * y, x * z, y * z
+
+    rotMat = torch.stack([w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz,
+                          2 * wz + 2 * xy, w2 - x2 + y2 - z2, 2 * yz - 2 * wx,
+                          2 * xz - 2 * wy, 2 * wx + 2 * yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)
+    return rotMat
+
+
+def inv_4x4(mats):
+    """Calculate the inverse of homogeneous transformations
+    :param mats: [B, 4, 4]
+    :return:
+    """
+    Rs = mats[:, :3, :3]
+    ts = mats[:, :3, 3:]
+    # R_invs = torch.transpose(Rs, 1, 2)
+    R_invs = torch.inverse(Rs)
+    t_invs = -torch.matmul(R_invs, ts)
+    Rt_invs = torch.cat([R_invs, t_invs], dim=-1)   # [B, 3, 4]
+
+    device = R_invs.device
+    pad_row = torch.FloatTensor([0, 0, 0, 1]).to(device).view(1, 1, 4).expand(Rs.shape[0], -1, -1)
+    mat_invs = torch.cat([Rt_invs, pad_row], dim=1)
+    return mat_invs
+
+
+def as_mesh(scene_or_mesh):
+    """
+    Convert a possible scene to a mesh.
+
+    If conversion occurs, the returned mesh has only vertex and face data.
+    """
+    if isinstance(scene_or_mesh, trimesh.Scene):
+        if len(scene_or_mesh.geometry) == 0:
+            mesh = None  # empty scene
+        else:
+            # we lose texture information here
+            mesh = trimesh.util.concatenate(
+                tuple(trimesh.Trimesh(vertices=g.vertices, faces=g.faces)
+                    for g in scene_or_mesh.geometry.values()))
+    else:
+        assert(isinstance(scene_or_mesh, trimesh.Trimesh))
+        mesh = scene_or_mesh
+    return mesh
+
+
+def get_edge_unique(faces):
+    """
+    Parameters
+    ------------
+    faces: n x 3 int array
+      Should be from a watertight mesh without degenerated triangles and intersection
+    """
+    faces = np.asanyarray(faces)
+
+    # each face has three edges
+    edges = faces[:, [0, 1, 1, 2, 2, 0]].reshape((-1, 2))
+    flags = edges[:, 0] < edges[:, 1]
+    edges = edges[flags]
+    return edges
+
+
+def get_neighbors(edges):
+    neighbors = collections.defaultdict(set)
+    [(neighbors[edge[0]].add(edge[1]),
+      neighbors[edge[1]].add(edge[0]))
+     for edge in edges]
+
+    max_index = edges.max() + 1
+    array = [list(neighbors[i]) for i in range(max_index)]
+
+    return array
+
+
+def construct_degree_matrix(vnum, faces):
+    row = col = list(range(vnum))
+    value = [0] * vnum
+    es = get_edge_unique(faces)
+    for e in es:
+        if e[0] < e[1]:
+            value[e[0]] += 1
+            value[e[0]] += 1
+
+    dm = sp.coo_matrix((value, (row, col)), shape=(vnum, vnum), dtype=np.float32)
+    return dm
+
+
+def construct_neighborhood_matrix(vnum, faces):
+    row = list()
+    col = list()
+    value = list()
+    es = get_edge_unique(faces)
+    for e in es:
+        if e[0] < e[1]:
+            row.append(e[0])
+            col.append(e[1])
+            value.append(1)
+            row.append(e[1])
+            col.append(e[0])
+            value.append(1)
+
+    nm = sp.coo_matrix((value, (row, col)), shape=(vnum, vnum), dtype=np.float32)
+    return nm
+
+
+def construct_laplacian_matrix(vnum, faces, normalized=False):
+    edges = get_edge_unique(faces)
+    neighbors = get_neighbors(edges)
+
+    col = np.concatenate(neighbors)
+    row = np.concatenate([[i] * len(n)
+                          for i, n in enumerate(neighbors)])
+    col = np.concatenate([col, np.arange(0, vnum)])
+    row = np.concatenate([row, np.arange(0, vnum)])
+
+    if normalized:
+        data = [[1.0 / len(n)] * len(n) for n in neighbors]
+        data += [[-1.0] * vnum]
+    else:
+        data = [[1.0] * len(n) for n in neighbors]
+        data += [[-len(n) for n in neighbors]]
+
+    data = np.concatenate(data)
+    # create the sparse matrix
+    matrix = sp.coo_matrix((data, (row, col)), shape=[vnum] * 2)
+    return matrix
+
+
+def rotationx_4x4(theta):
+    return np.array([
+        [1.0, 0.0, 0.0, 0.0],
+        [0.0, np.cos(theta / 180 * np.pi), np.sin(theta / 180 * np.pi), 0.0],
+        [0.0, -np.sin(theta / 180 * np.pi), np.cos(theta / 180 * np.pi), 0.0],
+        [0.0, 0.0, 0.0, 1.0]
+    ])
+
+
+def rotationy_4x4(theta):
+    return np.array([
+        [np.cos(theta / 180 * np.pi), 0.0, np.sin(theta / 180 * np.pi), 0.0],
+        [0.0, 1.0, 0.0, 0.0],
+        [-np.sin(theta / 180 * np.pi), 0.0, np.cos(theta / 180 * np.pi), 0.0],
+        [0.0, 0.0, 0.0, 1.0]
+    ])
+
+
+def rotationz_4x4(theta):
+    return np.array([
+        [np.cos(theta / 180 * np.pi), np.sin(theta / 180 * np.pi), 0.0, 0.0],
+        [-np.sin(theta / 180 * np.pi), np.cos(theta / 180 * np.pi), 0.0, 0.0],
+        [0.0, 0.0, 1.0, 0.0],
+        [0.0, 0.0, 0.0, 1.0]
+    ])
+
+
+def rotationx_3x3(theta):
+    return np.array([
+        [1.0, 0.0, 0.0],
+        [0.0, np.cos(theta / 180 * np.pi), np.sin(theta / 180 * np.pi)],
+        [0.0, -np.sin(theta / 180 * np.pi), np.cos(theta / 180 * np.pi)],
+    ])
+
+
+def rotationy_3x3(theta):
+    return np.array([
+        [np.cos(theta / 180 * np.pi), 0.0, np.sin(theta / 180 * np.pi)],
+        [0.0, 1.0, 0.0],
+        [-np.sin(theta / 180 * np.pi), 0.0, np.cos(theta / 180 * np.pi)],
+    ])
+
+
+def rotationz_3x3(theta):
+    return np.array([
+        [np.cos(theta / 180 * np.pi), np.sin(theta / 180 * np.pi), 0.0],
+        [-np.sin(theta / 180 * np.pi), np.cos(theta / 180 * np.pi), 0.0],
+        [0.0, 0.0, 1.0],
+    ])
+
+
+def generate_point_grids(vol_res):
+    x_coords = np.array(range(0, vol_res), dtype=np.float32)
+    y_coords = np.array(range(0, vol_res), dtype=np.float32)
+    z_coords = np.array(range(0, vol_res), dtype=np.float32)
+
+    yv, xv, zv = np.meshgrid(x_coords, y_coords, z_coords)
+    xv = np.reshape(xv, (-1, 1))
+    yv = np.reshape(yv, (-1, 1))
+    zv = np.reshape(zv, (-1, 1))
+    pts = np.concatenate([xv, yv, zv], axis=-1)
+    pts = pts.astype(np.float32)
+    return pts
+
+
+def infer_occupancy_value_grid_octree(test_res, pts, query_fn, init_res=64, ignore_thres=0.05):
+    pts = np.reshape(pts, (test_res, test_res, test_res, 3))
+
+    pts_ov = np.zeros([test_res, test_res, test_res])
+    dirty = np.ones_like(pts_ov, dtype=np.bool)
+    grid_mask = np.zeros_like(pts_ov, dtype=np.bool)
+
+    reso = test_res // init_res
+    while reso > 0:
+        grid_mask[0:test_res:reso, 0:test_res:reso, 0:test_res:reso] = True
+        test_mask = np.logical_and(grid_mask, dirty)
+
+        pts_ = pts[test_mask]
+        pts_ov[test_mask] = np.reshape(query_fn(pts_), pts_ov[test_mask].shape)
+
+        if reso <= 1:
+            break
+        for x in range(0, test_res - reso, reso):
+            for y in range(0, test_res - reso, reso):
+                for z in range(0, test_res - reso, reso):
+                    # if center marked, return
+                    if not dirty[x + reso // 2, y + reso // 2, z + reso // 2]:
+                        continue
+                    v0 = pts_ov[x, y, z]
+                    v1 = pts_ov[x, y, z + reso]
+                    v2 = pts_ov[x, y + reso, z]
+                    v3 = pts_ov[x, y + reso, z + reso]
+                    v4 = pts_ov[x + reso, y, z]
+                    v5 = pts_ov[x + reso, y, z + reso]
+                    v6 = pts_ov[x + reso, y + reso, z]
+                    v7 = pts_ov[x + reso, y + reso, z + reso]
+                    v = np.array([v0, v1, v2, v3, v4, v5, v6, v7])
+                    v_min = np.min(v)
+                    v_max = np.max(v)
+                    # this cell is all the same
+                    if (v_max - v_min) < ignore_thres:
+                        pts_ov[x:x + reso, y:y + reso, z:z + reso] = (v_max + v_min) / 2
+                        dirty[x:x + reso, y:y + reso, z:z + reso] = False
+        reso //= 2
+    return pts_ov
+
+
+def batch_rod2quat(rot_vecs):
+    batch_size = rot_vecs.shape[0]
+
+    angle = torch.norm(rot_vecs + 1e-16, dim=1, keepdim=True)
+    rot_dir = rot_vecs / angle
+
+    cos = torch.cos(angle / 2)
+    sin = torch.sin(angle / 2)
+
+    # Bx1 arrays
+    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
+
+    qx = rx * sin
+    qy = ry * sin
+    qz = rz * sin
+    qw = cos-1.0
+
+    return torch.cat([qx,qy,qz,qw], dim=1)
+
+
+def batch_quat2matrix(rvec):
+    '''
+    args:
+        rvec: (B, N, 4)
+    '''
+    B, N, _ = rvec.size()
+
+    theta = torch.sqrt(1e-5 + torch.sum(rvec ** 2, dim=2))
+    rvec = rvec / theta[:, :, None]
+    return torch.stack((
+        1. - 2. * rvec[:, :, 1] ** 2 - 2. * rvec[:, :, 2] ** 2,
+        2. * (rvec[:, :, 0] * rvec[:, :, 1] - rvec[:, :, 2] * rvec[:, :, 3]),
+        2. * (rvec[:, :, 0] * rvec[:, :, 2] + rvec[:, :, 1] * rvec[:, :, 3]),
+
+        2. * (rvec[:, :, 0] * rvec[:, :, 1] + rvec[:, :, 2] * rvec[:, :, 3]),
+        1. - 2. * rvec[:, :, 0] ** 2 - 2. * rvec[:, :, 2] ** 2,
+        2. * (rvec[:, :, 1] * rvec[:, :, 2] - rvec[:, :, 0] * rvec[:, :, 3]),
+
+        2. * (rvec[:, :, 0] * rvec[:, :, 2] - rvec[:, :, 1] * rvec[:, :, 3]),
+        2. * (rvec[:, :, 0] * rvec[:, :, 3] + rvec[:, :, 1] * rvec[:, :, 2]),
+        1. - 2. * rvec[:, :, 0] ** 2 - 2. * rvec[:, :, 1] ** 2
+        ), dim=2).view(B, N, 3, 3)
+
+
+def get_posemap(map_type, n_joints, parents, n_traverse=1, normalize=True):
+    pose_map = torch.zeros(n_joints,n_joints-1)
+    if map_type == 'parent':
+        for i in range(n_joints-1):
+            pose_map[i+1,i] = 1.0
+    elif map_type == 'children':
+        for i in range(n_joints-1):
+            parent = parents[i+1]
+            for j in range(n_traverse):
+                pose_map[parent, i] += 1.0
+                if parent == 0:
+                    break
+                parent = parents[parent]
+        if normalize:
+            pose_map /= pose_map.sum(0,keepdim=True)+1e-16
+    elif map_type == 'both':
+        for i in range(n_joints-1):
+            pose_map[i+1,i] += 1.0
+            parent = parents[i+1]
+            for j in range(n_traverse):
+                pose_map[parent, i] += 1.0
+                if parent == 0:
+                    break
+                parent = parents[parent]
+        if normalize:
+            pose_map /= pose_map.sum(0,keepdim=True)+1e-16
+    else:
+        raise NotImplementedError('unsupported pose map type [%s]' % map_type)
+    pose_map = torch.cat([torch.zeros(n_joints, 1), pose_map], dim=1)
+    return pose_map
+
+
+def vertices_to_triangles(vertices, faces):
+    """
+    :param vertices: [batch size, number of vertices, 3]
+    :param faces: [batch size, number of faces, 3)
+    :return: [batch size, number of faces, 3, 3]
+    """
+    assert (vertices.ndimension() == 3)
+    assert (faces.ndimension() == 3)
+    assert (vertices.shape[0] == faces.shape[0])
+    assert (vertices.shape[2] == 3)
+    assert (faces.shape[2] == 3)
+
+    bs, nv = vertices.shape[:2]
+    bs, nf = faces.shape[:2]
+    device = vertices.device
+    faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None]
+    vertices = vertices.reshape((bs * nv, 3))
+    # pytorch only supports long and byte tensors for indexing
+    return vertices[faces.long()]
+
+
+def calc_face_normals(vertices, faces):
+    assert len(vertices.shape) == 3
+    assert len(faces.shape) == 2
+    if isinstance(faces, np.ndarray):
+        faces = torch.from_numpy(faces.astype(np.int64)).to(vertices.device)
+
+    batch_size, pt_num = vertices.shape[:2]
+    face_num = faces.shape[0]
+
+    triangles = vertices_to_triangles(vertices, faces.unsqueeze(0).expand(batch_size, -1, -1))
+    triangles = triangles.reshape((batch_size * face_num, 3, 3))
+    v10 = triangles[:, 0] - triangles[:, 1]
+    v12 = triangles[:, 2] - triangles[:, 1]
+    # pytorch normalize divides by max(norm, eps) instead of (norm+eps) in chainer
+    normals = F.normalize(torch.cross(v10, v12), eps=1e-5)
+    normals = normals.reshape((batch_size, face_num, 3))
+
+    return normals
+
+
+def calc_vert_normals(vertices, faces):
+    """
+    vertices: [B, N, 3]
+    faces: [F, 3]
+    """
+    normals = torch.zeros_like(vertices)
+    v0s = torch.index_select(vertices, dim=1, index=faces[:, 0])    # [B, F, 3]
+    v1s = torch.index_select(vertices, dim=1, index=faces[:, 1])
+    v2s = torch.index_select(vertices, dim=1, index=faces[:, 2])
+    normals = torch.index_add(normals, dim=1, index=faces[:, 1], source=torch.cross(v2s-v1s, v0s-v1s, dim=-1))
+    normals = torch.index_add(normals, dim=1, index=faces[:, 2], source=torch.cross(v0s-v2s, v1s-v2s, dim=-1))
+    normals = torch.index_add(normals, dim=1, index=faces[:, 0], source=torch.cross(v1s-v0s, v2s-v0s, dim=-1))
+    normals = F.normalize(normals, dim=-1)
+    return normals
+
+
+def calc_vert_normals_numpy(vertices, faces):
+    assert len(vertices.shape) == 2
+    assert len(faces.shape) == 2
+
+    nmls = np.zeros_like(vertices)
+    fv0 = vertices[faces[:, 0]]
+    fv1 = vertices[faces[:, 1]]
+    fv2 = vertices[faces[:, 2]]
+    face_nmls = np.cross(fv1-fv0, fv2-fv0, axis=-1)
+    face_nmls = face_nmls / (np.linalg.norm(face_nmls, axis=-1, keepdims=True) + 1e-20)
+    for f, fn in zip(faces, face_nmls):
+        nmls[f] += fn
+    nmls = nmls / (np.linalg.norm(nmls, axis=-1, keepdims=True) + 1e-20)
+    return nmls
+
+
+def glUV2torchUV(gl_uv):
+    torch_uv = torch.stack([
+        gl_uv[..., 0]*2.0-1.0,
+        gl_uv[..., 1]*-2.0+1.0
+    ], dim=-1)
+    return torch_uv
+
+
+def normalize_vert_bbox(verts, dim=-1, per_axis=False):
+    bbox_min = torch.min(verts, dim=dim, keepdim=True)[0]
+    bbox_max = torch.max(verts, dim=dim, keepdim=True)[0]
+    verts = verts - 0.5 * (bbox_max + bbox_min)
+    if per_axis:
+        verts = 2 * verts / (bbox_max - bbox_min)
+    else:
+        verts = 2 * verts / torch.max(bbox_max-bbox_min, dim=dim, keepdim=True)[0]
+    return verts
+
+
+def upsample_sdf_volume(sdf, upsample_factor):
+    assert sdf.shape[0] == sdf.shape[1] == sdf.shape[2]
+    coarse_resolution = sdf.shape[0]
+    fine_resolution = coarse_resolution * upsample_factor
+
+    sdf_interp_buffer = np.zeros([2, 2, 2, coarse_resolution, coarse_resolution, coarse_resolution],
+                                 dtype=np.float32)
+    dx_list = [0, 1, 0, 1, 0, 1, 0, 1]
+    dy_list = [0, 0, 1, 1, 0, 0, 1, 1]
+    dz_list = [0, 0, 0, 0, 1, 1, 1, 1]
+    for dx, dy, dz in zip(dx_list, dy_list, dz_list):
+        sdf_interp_buffer[dx, dy, dz, :, :, :] = np.roll(sdf, (-dx, -dy, -dz), axis=(0, 1, 2))
+
+    sdf_fine = np.zeros([fine_resolution, fine_resolution, fine_resolution], dtype=np.float32)
+    for dx in range(upsample_factor):
+        for dy in range(upsample_factor):
+            for dz in range(upsample_factor):
+                wx = (1.0 - dx / upsample_factor)
+                wy = (1.0 - dy / upsample_factor)
+                wz = (1.0 - dz / upsample_factor)
+                sdf_fine[dx::upsample_factor, dy::upsample_factor, dz::upsample_factor] += \
+                    wx * wy * wz * sdf_interp_buffer[0, 0, 0]
+                sdf_fine[dx::upsample_factor, dy::upsample_factor, dz::upsample_factor] += \
+                    (1.0 - wx) * wy * wz * sdf_interp_buffer[1, 0, 0]
+                sdf_fine[dx::upsample_factor, dy::upsample_factor, dz::upsample_factor] += \
+                    wx * (1.0 - wy) * wz * sdf_interp_buffer[0, 1, 0]
+                sdf_fine[dx::upsample_factor, dy::upsample_factor, dz::upsample_factor] += \
+                    (1.0 - wx) * (1.0 - wy) * wz * sdf_interp_buffer[1, 1, 0]
+                sdf_fine[dx::upsample_factor, dy::upsample_factor, dz::upsample_factor] += \
+                    wx * wy * (1.0 - wz) * sdf_interp_buffer[0, 0, 1]
+                sdf_fine[dx::upsample_factor, dy::upsample_factor, dz::upsample_factor] += \
+                    (1.0 - wx) * wy * (1.0 - wz) * sdf_interp_buffer[1, 0, 1]
+                sdf_fine[dx::upsample_factor, dy::upsample_factor, dz::upsample_factor] += \
+                    wx * (1.0 - wy) * (1.0 - wz) * sdf_interp_buffer[0, 1, 1]
+                sdf_fine[dx::upsample_factor, dy::upsample_factor, dz::upsample_factor] += \
+                    (1.0 - wx) * (1.0 - wy) * (1.0 - wz) * sdf_interp_buffer[1, 1, 1]
+    return sdf_fine
+
+
+def search_nearest_correspondence(src, tgt):
+    tgt_idx = np.zeros(len(src), dtype=np.int32)
+    tgt_dist = np.zeros(len(src), dtype=np.float32)
+    for i in range(len(src)):
+        dist = np.linalg.norm(tgt - src[i:(i+1)], axis=1, keepdims=False)
+        tgt_idx[i] = np.argmin(dist)
+        tgt_dist[i] = dist[tgt_idx[i]]
+    return tgt_idx, tgt_dist
+
+
+def estimate_rigid_transformation(src, tgt):
+    src = src.transpose()
+    tgt = tgt.transpose()
+    mu1, mu2 = src.mean(axis=1, keepdims=True), tgt.mean(axis=1, keepdims=True)
+    X1, X2 = src - mu1, tgt - mu2
+
+    K = X1.dot(X2.T)
+    U, s, Vh = np.linalg.svd(K)
+    V = Vh.T
+    Z = np.eye(U.shape[0])
+    Z[-1, -1] *= np.sign(np.linalg.det(U.dot(V.T)))
+    R = V.dot(Z.dot(U.T))
+    t = mu2 - R.dot(mu1)
+
+    # orient, _ = cv.Rodrigues(R)
+    # orient = orient.reshape([-1])
+    # t = t.reshape([-1])
+    # return orient, t
+
+    transf = np.eye(4, dtype=np.float32)
+    transf[:3, :3] = R
+    transf[:3, 3] = t.reshape([-1])
+    return transf

render_utils/lib/utils/graphics_utils.py

@@ -0,0 +1,181 @@
+#
+# Copyright (C) 2023, Inria
+# GRAPHDECO research group, https://team.inria.fr/graphdeco
+# All rights reserved.
+#
+# This software is free for non-commercial, research and evaluation use 
+# under the terms of the LICENSE.md file.
+#
+# For inquiries contact  george.drettakis@inria.fr
+#
+
+import torch
+import math
+import numpy as np
+from typing import NamedTuple
+
+class BasicPointCloud(NamedTuple):
+    points : np.array
+    colors : np.array
+    normals : np.array
+
+def geom_transform_points(points, transf_matrix):
+    P, _ = points.shape
+    ones = torch.ones(P, 1, dtype=points.dtype, device=points.device)
+    points_hom = torch.cat([points, ones], dim=1)
+    points_out = torch.matmul(points_hom, transf_matrix.unsqueeze(0))
+
+    denom = points_out[..., 3:] + 0.0000001
+    return (points_out[..., :3] / denom).squeeze(dim=0)
+
+def getWorld2View(R, t):
+    Rt = np.zeros((4, 4))
+    Rt[:3, :3] = R.transpose()
+    Rt[:3, 3] = t
+    Rt[3, 3] = 1.0
+    return np.float32(Rt)
+
+def getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0):
+    Rt = np.zeros((4, 4))
+    Rt[:3, :3] = R.transpose()
+    Rt[:3, 3] = t
+    Rt[3, 3] = 1.0
+
+    C2W = np.linalg.inv(Rt)
+    cam_center = C2W[:3, 3]
+    cam_center = (cam_center + translate) * scale
+    C2W[:3, 3] = cam_center
+    Rt = np.linalg.inv(C2W)
+    return np.float32(Rt)
+
+def getProjectionMatrix(znear, zfar, fovX, fovY, K = None, img_h = None, img_w = None):
+    if K is None:
+        tanHalfFovY = math.tan((fovY / 2))
+        tanHalfFovX = math.tan((fovX / 2))
+        top = tanHalfFovY * znear
+        bottom = -top
+        right = tanHalfFovX * znear
+        left = -right
+    else:
+        near_fx = znear / K[0, 0]
+        near_fy = znear / K[1, 1]
+
+        left = - (img_w - K[0, 2]) * near_fx
+        right = K[0, 2] * near_fx
+        bottom = (K[1, 2] - img_h) * near_fy
+        top = K[1, 2] * near_fy
+
+    P = torch.zeros(4, 4)
+
+    z_sign = 1.0
+
+    P[0, 0] = 2.0 * znear / (right - left)
+    P[1, 1] = 2.0 * znear / (top - bottom)
+    P[0, 2] = (right + left) / (right - left)
+    P[1, 2] = (top + bottom) / (top - bottom)
+    P[3, 2] = z_sign
+    P[2, 2] = z_sign * zfar / (zfar - znear)
+    P[2, 3] = -(zfar * znear) / (zfar - znear)
+    return P
+
+def fov2focal(fov, pixels):
+    return pixels / (2 * math.tan(fov / 2))
+
+def focal2fov(focal, pixels):
+    return 2*math.atan(pixels/(2*focal))
+
+
+def _so3_exp_map(
+    log_rot: torch.Tensor, eps: float = 0.0001
+):
+    """
+    A helper function that computes the so3 exponential map and,
+    apart from the rotation matrix, also returns intermediate variables
+    that can be re-used in other functions.
+    """
+
+    def hat(v: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the Hat operator [1] of a batch of 3D vectors.
+
+        Args:
+            v: Batch of vectors of shape `(minibatch , 3)`.
+
+        Returns:
+            Batch of skew-symmetric matrices of shape
+            `(minibatch, 3 , 3)` where each matrix is of the form:
+                `[    0  -v_z   v_y ]
+                 [  v_z     0  -v_x ]
+                 [ -v_y   v_x     0 ]`
+
+        Raises:
+            ValueError if `v` is of incorrect shape.
+
+        [1] https://en.wikipedia.org/wiki/Hat_operator
+        """
+
+        N, dim = v.shape
+        if dim != 3:
+            raise ValueError("Input vectors have to be 3-dimensional.")
+
+        h = torch.zeros((N, 3, 3), dtype=v.dtype, device=v.device)
+
+        x, y, z = v.unbind(1)
+
+        h[:, 0, 1] = -z
+        h[:, 0, 2] = y
+        h[:, 1, 0] = z
+        h[:, 1, 2] = -x
+        h[:, 2, 0] = -y
+        h[:, 2, 1] = x
+
+        return h
+
+    _, dim = log_rot.shape
+    if dim != 3:
+        raise ValueError("Input tensor shape has to be Nx3.")
+
+    nrms = (log_rot * log_rot).sum(1)
+    # phis ... rotation angles
+    rot_angles = torch.clamp(nrms, eps).sqrt()
+    rot_angles_inv = 1.0 / rot_angles
+    fac1 = rot_angles_inv * rot_angles.sin()
+    fac2 = rot_angles_inv * rot_angles_inv * (1.0 - rot_angles.cos())
+    skews = hat(log_rot)
+    skews_square = torch.bmm(skews, skews)
+
+    R = (
+        # pyre-fixme[16]: `float` has no attribute `__getitem__`.
+        fac1[:, None, None] * skews
+        + fac2[:, None, None] * skews_square
+        + torch.eye(3, dtype=log_rot.dtype, device=log_rot.device)[None]
+    )
+
+    return R, rot_angles, skews, skews_square
+
+
+def so3_exp_map(log_rot: torch.Tensor, eps: float = 0.0001) -> torch.Tensor:
+    """
+    Convert a batch of logarithmic representations of rotation matrices `log_rot`
+    to a batch of 3x3 rotation matrices using Rodrigues formula [1].
+
+    In the logarithmic representation, each rotation matrix is represented as
+    a 3-dimensional vector (`log_rot`) who's l2-norm and direction correspond
+    to the magnitude of the rotation angle and the axis of rotation respectively.
+
+    The conversion has a singularity around `log(R) = 0`
+    which is handled by clamping controlled with the `eps` argument.
+
+    Args:
+        log_rot: Batch of vectors of shape `(minibatch, 3)`.
+        eps: A float constant handling the conversion singularity.
+
+    Returns:
+        Batch of rotation matrices of shape `(minibatch, 3, 3)`.
+
+    Raises:
+        ValueError if `log_rot` is of incorrect shape.
+
+    [1] https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
+    """
+    return _so3_exp_map(log_rot, eps=eps)[0]

render_utils/lib/utils/rotation_conversions.py

@@ -0,0 +1,586 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional, Union
+
+import torch
+import torch.nn.functional as F
+
+Device = Union[str, torch.device]
+
+"""
+The transformation matrices returned from the functions in this file assume
+the points on which the transformation will be applied are column vectors.
+i.e. the R matrix is structured as
+
+    R = [
+            [Rxx, Rxy, Rxz],
+            [Ryx, Ryy, Ryz],
+            [Rzx, Rzy, Rzz],
+        ]  # (3, 3)
+
+This matrix can be applied to column vectors by post multiplication
+by the points e.g.
+
+    points = [[0], [1], [2]]  # (3 x 1) xyz coordinates of a point
+    transformed_points = R * points
+
+To apply the same matrix to points which are row vectors, the R matrix
+can be transposed and pre multiplied by the points:
+
+e.g.
+    points = [[0, 1, 2]]  # (1 x 3) xyz coordinates of a point
+    transformed_points = points * R.transpose(1, 0)
+"""
+
+
+def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as quaternions to rotation matrices.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def _copysign(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    Return a tensor where each element has the absolute value taken from the,
+    corresponding element of a, with sign taken from the corresponding
+    element of b. This is like the standard copysign floating-point operation,
+    but is not careful about negative 0 and NaN.
+
+    Args:
+        a: source tensor.
+        b: tensor whose signs will be used, of the same shape as a.
+
+    Returns:
+        Tensor of the same shape as a with the signs of b.
+    """
+    signs_differ = (a < 0) != (b < 0)
+    return torch.where(signs_differ, -a, a)
+
+
+def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
+    """
+    Returns torch.sqrt(torch.max(0, x))
+    but with a zero subgradient where x is 0.
+    """
+    ret = torch.zeros_like(x)
+    positive_mask = x > 0
+    ret[positive_mask] = torch.sqrt(x[positive_mask])
+    return ret
+
+
+def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to quaternions.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+
+    batch_dim = matrix.shape[:-2]
+    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
+        matrix.reshape(batch_dim + (9,)), dim=-1
+    )
+
+    q_abs = _sqrt_positive_part(
+        torch.stack(
+            [
+                1.0 + m00 + m11 + m22,
+                1.0 + m00 - m11 - m22,
+                1.0 - m00 + m11 - m22,
+                1.0 - m00 - m11 + m22,
+            ],
+            dim=-1,
+        )
+    )
+
+    # we produce the desired quaternion multiplied by each of r, i, j, k
+    quat_by_rijk = torch.stack(
+        [
+            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
+            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
+            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
+            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
+        ],
+        dim=-2,
+    )
+
+    # We floor here at 0.1 but the exact level is not important; if q_abs is small,
+    # the candidate won't be picked.
+    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
+    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
+
+    # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
+    # forall i; we pick the best-conditioned one (with the largest denominator)
+
+    return quat_candidates[
+        F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :  # pyre-ignore[16]
+    ].reshape(batch_dim + (4,))
+
+
+def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
+    """
+    Return the rotation matrices for one of the rotations about an axis
+    of which Euler angles describe, for each value of the angle given.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z".
+        angle: any shape tensor of Euler angles in radians
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+
+    cos = torch.cos(angle)
+    sin = torch.sin(angle)
+    one = torch.ones_like(angle)
+    zero = torch.zeros_like(angle)
+
+    if axis == "X":
+        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
+    elif axis == "Y":
+        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
+    elif axis == "Z":
+        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
+    else:
+        raise ValueError("letter must be either X, Y or Z.")
+
+    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
+
+
+def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor:
+    """
+    Convert rotations given as Euler angles in radians to rotation matrices.
+
+    Args:
+        euler_angles: Euler angles in radians as tensor of shape (..., 3).
+        convention: Convention string of three uppercase letters from
+            {"X", "Y", and "Z"}.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
+        raise ValueError("Invalid input euler angles.")
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    matrices = [
+        _axis_angle_rotation(c, e)
+        for c, e in zip(convention, torch.unbind(euler_angles, -1))
+    ]
+    # return functools.reduce(torch.matmul, matrices)
+    return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2])
+
+
+def _angle_from_tan(
+    axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool
+) -> torch.Tensor:
+    """
+    Extract the first or third Euler angle from the two members of
+    the matrix which are positive constant times its sine and cosine.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z" for the angle we are finding.
+        other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
+            convention.
+        data: Rotation matrices as tensor of shape (..., 3, 3).
+        horizontal: Whether we are looking for the angle for the third axis,
+            which means the relevant entries are in the same row of the
+            rotation matrix. If not, they are in the same column.
+        tait_bryan: Whether the first and third axes in the convention differ.
+
+    Returns:
+        Euler Angles in radians for each matrix in data as a tensor
+        of shape (...).
+    """
+
+    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
+    if horizontal:
+        i2, i1 = i1, i2
+    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
+    if horizontal == even:
+        return torch.atan2(data[..., i1], data[..., i2])
+    if tait_bryan:
+        return torch.atan2(-data[..., i2], data[..., i1])
+    return torch.atan2(data[..., i2], -data[..., i1])
+
+
+def _index_from_letter(letter: str) -> int:
+    if letter == "X":
+        return 0
+    if letter == "Y":
+        return 1
+    if letter == "Z":
+        return 2
+    raise ValueError("letter must be either X, Y or Z.")
+
+
+def matrix_to_euler_angles(matrix: torch.Tensor, convention: str) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to Euler angles in radians.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+        convention: Convention string of three uppercase letters.
+
+    Returns:
+        Euler angles in radians as tensor of shape (..., 3).
+    """
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+    i0 = _index_from_letter(convention[0])
+    i2 = _index_from_letter(convention[2])
+    tait_bryan = i0 != i2
+    if tait_bryan:
+        central_angle = torch.asin(
+            matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
+        )
+    else:
+        central_angle = torch.acos(matrix[..., i0, i0])
+
+    o = (
+        _angle_from_tan(
+            convention[0], convention[1], matrix[..., i2], False, tait_bryan
+        ),
+        central_angle,
+        _angle_from_tan(
+            convention[2], convention[1], matrix[..., i0, :], True, tait_bryan
+        ),
+    )
+    return torch.stack(o, -1)
+
+
+def random_quaternions(
+    n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
+) -> torch.Tensor:
+    """
+    Generate random quaternions representing rotations,
+    i.e. versors with nonnegative real part.
+
+    Args:
+        n: Number of quaternions in a batch to return.
+        dtype: Type to return.
+        device: Desired device of returned tensor. Default:
+            uses the current device for the default tensor type.
+
+    Returns:
+        Quaternions as tensor of shape (N, 4).
+    """
+    if isinstance(device, str):
+        device = torch.device(device)
+    o = torch.randn((n, 4), dtype=dtype, device=device)
+    s = (o * o).sum(1)
+    o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
+    return o
+
+
+def random_rotations(
+    n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
+) -> torch.Tensor:
+    """
+    Generate random rotations as 3x3 rotation matrices.
+
+    Args:
+        n: Number of rotation matrices in a batch to return.
+        dtype: Type to return.
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type.
+
+    Returns:
+        Rotation matrices as tensor of shape (n, 3, 3).
+    """
+    quaternions = random_quaternions(n, dtype=dtype, device=device)
+    return quaternion_to_matrix(quaternions)
+
+
+def random_rotation(
+    dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
+) -> torch.Tensor:
+    """
+    Generate a single random 3x3 rotation matrix.
+
+    Args:
+        dtype: Type to return
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type
+
+    Returns:
+        Rotation matrix as tensor of shape (3, 3).
+    """
+    return random_rotations(1, dtype, device)[0]
+
+
+def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
+
+
+def quaternion_raw_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    Multiply two quaternions.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions shape (..., 4).
+    """
+    aw, ax, ay, az = torch.unbind(a, -1)
+    bw, bx, by, bz = torch.unbind(b, -1)
+    ow = aw * bw - ax * bx - ay * by - az * bz
+    ox = aw * bx + ax * bw + ay * bz - az * by
+    oy = aw * by - ax * bz + ay * bw + az * bx
+    oz = aw * bz + ax * by - ay * bx + az * bw
+    return torch.stack((ow, ox, oy, oz), -1)
+
+
+def quaternion_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    Multiply two quaternions representing rotations, returning the quaternion
+    representing their composition, i.e. the versor with nonnegative real part.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions of shape (..., 4).
+    """
+    ab = quaternion_raw_multiply(a, b)
+    return standardize_quaternion(ab)
+
+
+def quaternion_invert(quaternion: torch.Tensor) -> torch.Tensor:
+    """
+    Given a quaternion representing rotation, get the quaternion representing
+    its inverse.
+
+    Args:
+        quaternion: Quaternions as tensor of shape (..., 4), with real part
+            first, which must be versors (unit quaternions).
+
+    Returns:
+        The inverse, a tensor of quaternions of shape (..., 4).
+    """
+
+    scaling = torch.tensor([1, -1, -1, -1], device=quaternion.device)
+    return quaternion * scaling
+
+
+def quaternion_apply(quaternion: torch.Tensor, point: torch.Tensor) -> torch.Tensor:
+    """
+    Apply the rotation given by a quaternion to a 3D point.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        quaternion: Tensor of quaternions, real part first, of shape (..., 4).
+        point: Tensor of 3D points of shape (..., 3).
+
+    Returns:
+        Tensor of rotated points of shape (..., 3).
+    """
+    if point.size(-1) != 3:
+        raise ValueError(f"Points are not in 3D, {point.shape}.")
+    real_parts = point.new_zeros(point.shape[:-1] + (1,))
+    point_as_quaternion = torch.cat((real_parts, point), -1)
+    out = quaternion_raw_multiply(
+        quaternion_raw_multiply(quaternion, point_as_quaternion),
+        quaternion_invert(quaternion),
+    )
+    return out[..., 1:]
+
+
+def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as axis/angle to rotation matrices.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
+
+
+def matrix_to_axis_angle(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to axis/angle.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        Rotations given as a vector in axis angle form, as a tensor
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
+
+
+def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as axis/angle to quaternions.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
+    half_angles = angles * 0.5
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
+    )
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48
+    )
+    quaternions = torch.cat(
+        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
+    )
+    return quaternions
+
+
+def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as quaternions to axis/angle.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotations given as a vector in axis angle form, as a tensor
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
+    half_angles = torch.atan2(norms, quaternions[..., :1])
+    angles = 2 * half_angles
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
+    )
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48
+    )
+    return quaternions[..., 1:] / sin_half_angles_over_angles
+
+
+def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
+    """
+    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
+    using Gram--Schmidt orthogonalization per Section B of [1].
+    Args:
+        d6: 6D rotation representation, of size (*, 6)
+
+    Returns:
+        batch of rotation matrices of size (*, 3, 3)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+
+    a1, a2 = d6[..., :3], d6[..., 3:]
+    b1 = F.normalize(a1, dim=-1)
+    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+    b2 = F.normalize(b2, dim=-1)
+    b3 = torch.cross(b1, b2, dim=-1)
+    return torch.stack((b1, b2, b3), dim=-2)
+
+
+def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
+    by dropping the last row. Note that 6D representation is not unique.
+    Args:
+        matrix: batch of rotation matrices of size (*, 3, 3)
+
+    Returns:
+        6D rotation representation, of size (*, 6)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+    batch_dim = matrix.size()[:-2]
+    return matrix[..., :2, :].clone().reshape(batch_dim + (6,))

render_utils/lib/utils/sh_utils.py

@@ -0,0 +1,118 @@
+#  Copyright 2021 The PlenOctree Authors.
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#  this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#  this list of conditions and the following disclaimer in the documentation
+#  and/or other materials provided with the distribution.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+
+import torch
+
+C0 = 0.28209479177387814
+C1 = 0.4886025119029199
+C2 = [
+    1.0925484305920792,
+    -1.0925484305920792,
+    0.31539156525252005,
+    -1.0925484305920792,
+    0.5462742152960396
+]
+C3 = [
+    -0.5900435899266435,
+    2.890611442640554,
+    -0.4570457994644658,
+    0.3731763325901154,
+    -0.4570457994644658,
+    1.445305721320277,
+    -0.5900435899266435
+]
+C4 = [
+    2.5033429417967046,
+    -1.7701307697799304,
+    0.9461746957575601,
+    -0.6690465435572892,
+    0.10578554691520431,
+    -0.6690465435572892,
+    0.47308734787878004,
+    -1.7701307697799304,
+    0.6258357354491761,
+]   
+
+
+def eval_sh(deg, sh, dirs):
+    """
+    Evaluate spherical harmonics at unit directions
+    using hardcoded SH polynomials.
+    Works with torch/np/jnp.
+    ... Can be 0 or more batch dimensions.
+    Args:
+        deg: int SH deg. Currently, 0-3 supported
+        sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2]
+        dirs: jnp.ndarray unit directions [..., 3]
+    Returns:
+        [..., C]
+    """
+    assert deg <= 4 and deg >= 0
+    coeff = (deg + 1) ** 2
+    assert sh.shape[-1] >= coeff
+
+    result = C0 * sh[..., 0]
+    if deg > 0:
+        x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]
+        result = (result -
+                C1 * y * sh[..., 1] +
+                C1 * z * sh[..., 2] -
+                C1 * x * sh[..., 3])
+
+        if deg > 1:
+            xx, yy, zz = x * x, y * y, z * z
+            xy, yz, xz = x * y, y * z, x * z
+            result = (result +
+                    C2[0] * xy * sh[..., 4] +
+                    C2[1] * yz * sh[..., 5] +
+                    C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] +
+                    C2[3] * xz * sh[..., 7] +
+                    C2[4] * (xx - yy) * sh[..., 8])
+
+            if deg > 2:
+                result = (result +
+                C3[0] * y * (3 * xx - yy) * sh[..., 9] +
+                C3[1] * xy * z * sh[..., 10] +
+                C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] +
+                C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] +
+                C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] +
+                C3[5] * z * (xx - yy) * sh[..., 14] +
+                C3[6] * x * (xx - 3 * yy) * sh[..., 15])
+
+                if deg > 3:
+                    result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] +
+                            C4[1] * yz * (3 * xx - yy) * sh[..., 17] +
+                            C4[2] * xy * (7 * zz - 1) * sh[..., 18] +
+                            C4[3] * yz * (7 * zz - 3) * sh[..., 19] +
+                            C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] +
+                            C4[5] * xz * (7 * zz - 3) * sh[..., 21] +
+                            C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] +
+                            C4[7] * xz * (xx - 3 * yy) * sh[..., 23] +
+                            C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24])
+    return result
+
+def RGB2SH(rgb):
+    return (rgb - 0.5) / C0
+
+def SH2RGB(sh):
+    return sh * C0 + 0.5

render_utils/stitch_body_and_head.py

@@ -0,0 +1,433 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import trimesh
+from sklearn.neighbors import KDTree
+import tqdm
+import cv2 as cv
+import os, glob
+
+from .lib.networks.faceverse_torch import FaceVerseModel
+from .lib.networks.smpl_torch import SmplTorch
+from .lib.utils.gaussian_np_utils import GaussianAttributes, load_gaussians_from_ply, save_gaussians_as_ply, \
+    apply_transformation_to_gaussians, combine_gaussians, select_gaussians, update_gaussian_attributes
+from .lib.utils.geometry import search_nearest_correspondence, estimate_rigid_transformation
+from .lib.utils.sh_utils import SH2RGB
+
+
+def process_smpl_head():
+    smpl = SmplTorch(model_file='./AnimatableGaussians/smpl_files/smplx/SMPLX_NEUTRAL.npz')
+    smpl_v_template = smpl.v_template.detach().cpu().numpy()
+    smpl_faces = smpl.faces.detach().cpu().numpy()
+    # head_skinning_weights = smpl.weights[:, 15] + smpl.weights[:, 22] + smpl.weights[:, 23] + smpl.weights[:, 24]
+    #
+    # blend_weight = np.clip(head_skinning_weights* 1.2 - 0.2, 0, 1)
+    # head_ids = np.where(blend_weight > 0)[0]
+    #
+    # np.savez('./data/smplx_head_vidx_and_blendweight.npz', blend_weight=blend_weight, head_ids=head_ids)
+
+    if not os.path.exists('./data/smpl_models/smplx_head_3.obj'):
+        trimesh.Trimesh(vertices=smpl_v_template, faces=smpl_faces).export('./data/smpl_models/smplx_head_3.obj')
+        print('Please cut out SMPL head!!!')
+        import pdb; pdb.set_trace()
+
+    smplx_head = trimesh.load('./data/smpl_models/smplx_head_3.obj')
+    smplx_to_head_dist = np.zeros([smpl_v_template.shape[0]])
+    for vi, v in enumerate(smpl_v_template):
+        nndist = np.min(np.linalg.norm(v.reshape([1, 3]) - smplx_head.vertices, axis=1))
+        smplx_to_head_dist[vi] = nndist
+
+    head_ids = np.where(smplx_to_head_dist < 0.001)[0]
+    blend_weight = np.exp(-smplx_to_head_dist*smplx_to_head_dist * 2000)
+
+    np.savez('./data/smpl_models/smplx_head_vidx_and_blendweight.npz', blend_weight=blend_weight, head_ids=head_ids)
+
+    return
+
+
+def load_body_params(path):
+    param = dict(np.load(path))
+    global_orient = param['global_orient']
+    transl = param['transl']
+    body_pose = param['body_pose']
+    betas = param['betas']
+    return global_orient, transl, body_pose, betas
+
+
+def load_face_params(path):
+    param = dict(np.load(path))
+    pose = param['pose']
+    scale = param['scale']
+    id_coeff = param['id_coeff']
+    exp_coeff = param['exp_coeff']
+    return pose, scale, id_coeff, exp_coeff
+
+
+def get_smpl_verts_and_head_transformation(smpl, global_orient, body_pose, transl, betas):
+    pose = torch.cat([
+        torch.from_numpy(global_orient.astype(np.float32)),
+        torch.from_numpy(body_pose.astype(np.float32)),
+        torch.zeros([(3+15+15)*3], dtype=torch.float32)], dim=-1)
+    beta = torch.from_numpy(betas.astype(np.float32))
+    verts, skinning_dict = smpl.forward(pose.reshape(1, -1), beta.reshape(1, -1))
+    verts = verts[0].detach().cpu().numpy()
+    head_joint_transfmat = skinning_dict['G'][0, 15].detach().cpu().numpy()
+    verts += transl.reshape([1, 3])
+    head_joint_transfmat[:3, 3] += transl.reshape([3])
+    return verts, head_joint_transfmat
+
+
+def crop_facial_area(faceverse_verts, points, dist_thres=0.025):
+    min_x, min_y, min_z = np.min(faceverse_verts, axis=0)
+    max_x, max_y, max_z = np.max(faceverse_verts, axis=0)
+    pad = dist_thres*2
+    in_bbox_mask = (points[:, 0] > min_x - pad) * (points[:, 0] < max_x + pad) * \
+                   (points[:, 1] > min_y - pad) * (points[:, 1] < max_y + pad) * \
+                   (points[:, 2] > min_z - pad) * (points[:, 2] < max_z + pad)
+    in_bbox_idx = np.where(in_bbox_mask)[0]
+    facial_points = points[in_bbox_mask]
+    nndist = np.ones([len(facial_points)]) * 1e10
+    for i in tqdm.trange(len(facial_points), desc='calculating facial area'):
+        nndist[i] = np.min(np.linalg.norm(faceverse_verts - facial_points[i:(i+1)], axis=1, keepdims=False))
+    close_to_face_mask = nndist < dist_thres
+    facial_points = facial_points[close_to_face_mask]
+    facial_idx = in_bbox_idx[close_to_face_mask]
+    return facial_points, facial_idx
+
+
+def crop_facial_area2(smpl_verts, smpl_head_vids, points):
+    min_x, min_y, min_z = np.min(smpl_verts[smpl_head_vids], axis=0)
+    max_x, max_y, max_z = np.max(smpl_verts[smpl_head_vids], axis=0)
+    pad = 0.05
+    in_bbox_mask = (points[:, 0] > min_x - pad) * (points[:, 0] < max_x + pad) * \
+                   (points[:, 1] > min_y - pad) * (points[:, 1] < max_y + pad) * \
+                   (points[:, 2] > min_z - pad) * (points[:, 2] < max_z + pad)
+    in_bbox_idx = np.where(in_bbox_mask)[0]
+    facial_points = points[in_bbox_mask]
+    smpl_head_mask = np.zeros([len(smpl_verts)], dtype=np.bool_)
+    smpl_head_mask[smpl_head_vids] = True
+    close_to_face_mask = np.zeros([len(facial_points)], dtype=np.bool_)
+    for i in tqdm.trange(len(facial_points)):
+        nnid = np.argmin(np.linalg.norm(smpl_verts - facial_points[i:(i+1)], axis=1, keepdims=False))
+        close_to_face_mask[i] = smpl_head_mask[nnid]
+    facial_points = facial_points[close_to_face_mask]
+    facial_idx = in_bbox_idx[close_to_face_mask]
+    return facial_points, facial_idx
+
+
+def transform_faceverse_to_live_body_space(faceverse_verts, faceverse_to_smplx, head_joint_transfmat):
+    faceverse_verts = np.matmul(faceverse_verts, faceverse_to_smplx[:3, :3].transpose()) + faceverse_to_smplx[:3, 3].reshape(1, 3)
+    faceverse_verts = np.matmul(faceverse_verts, head_joint_transfmat[:3, :3].transpose()) + head_joint_transfmat[:3, 3].reshape(1, 3)
+    return faceverse_verts
+
+
+def calc_livehead2livebody(head_pose, smplx_to_faceverse, head_joint_transfmat):
+    head_cano2live = np.eye(4, dtype=np.float32)
+    head_cano2live[:3, :3] = cv.Rodrigues(head_pose[:3])[0]
+    head_cano2live[:3, 3] = head_pose[3:]
+    head_live2cano = np.linalg.inv(head_cano2live)
+
+    faceverse_to_smplx = np.linalg.inv(smplx_to_faceverse)
+
+    total_transf = np.eye(4, dtype=np.float32)
+    for t in [head_live2cano, np.diag([1, -1, -1, 1]), faceverse_to_smplx, head_joint_transfmat]:
+        total_transf = np.matmul(t, total_transf)
+
+    return total_transf
+
+
+def get_face_blend_weight(head_facial_points, smpl_verts, sigma=0.015):
+    # dists = np.load('./data/faceverse/smplx_verts_to_faceverse_dist.npy').astype(np.float32)
+    # face_nerf_blend_weight = np.exp(-dists**2/(2*sigma**2))
+    # face_nerf_blend_weight = np.clip(face_nerf_blend_weight*1.2 - 0.1, 0, 1)
+
+    smpl_blend_weight = dict(np.load('./data/smpl_models/smplx_head_vidx_and_blendweight.npz'))['blend_weight']
+
+    corr_idx_, _ = search_nearest_correspondence(head_facial_points, smpl_verts)
+    corr_bw = smpl_blend_weight[corr_idx_]
+
+    for _ in tqdm.trange(10):
+        corr_bw_ = np.zeros_like(corr_bw)
+        tree = KDTree(head_facial_points, leaf_size=2)
+        for i in range(len(head_facial_points)):
+            _, idx = tree.query(head_facial_points[i:(i+1)], k=4)
+            corr_bw_[i] = np.mean(corr_bw[idx])
+        corr_bw = np.copy(corr_bw_)
+
+    # corr_bw = np.clip(corr_bw*1.2 - 0.15, 0, 1)
+
+    # with open('./debug/debug_head_facial_bw.obj', 'w') as fp:
+    #     for p, w in zip(head_facial_points, corr_bw):
+    #         fp.write('v %f %f %f %f %f %f\n' % (p[0], p[1], p[2], w, w, w))
+    # import pdb; pdb.set_trace()
+    return corr_bw
+
+
+def get_face_blend_weight2(head_facial_points, body_points, body_facial_idx):
+    body_facial_bbox_min = np.min(body_points[body_facial_idx], axis=0)
+    body_facial_bbox_max = np.max(body_points[body_facial_idx], axis=0)
+    body_facial_bbox_min = body_facial_bbox_min - 0.1
+    body_facial_bbox_max = body_facial_bbox_max + 0.1
+    inside_bbox_flag = \
+        np.int32(body_points[:, 0] > body_facial_bbox_min[0]) * \
+        np.int32(body_points[:, 0] < body_facial_bbox_max[0]) * \
+        np.int32(body_points[:, 1] > body_facial_bbox_min[1]) * \
+        np.int32(body_points[:, 1] < body_facial_bbox_max[1]) * \
+        np.int32(body_points[:, 2] > body_facial_bbox_min[2]) * \
+        np.int32(body_points[:, 2] < body_facial_bbox_max[2])
+    point_idx_inside_bbox = np.nonzero(inside_bbox_flag >0)[0]
+    body_blend_weight = np.zeros([len(body_points)], dtype=np.float32)
+    body_blend_weight[body_facial_idx] = 1
+
+    body_points_in_bbox = body_points[point_idx_inside_bbox]
+    body_blend_weight_in_bbox = body_blend_weight[point_idx_inside_bbox]
+    for _ in tqdm.trange(1, desc='Calculating body facial blend weight'):
+        corr_bw_ = np.zeros_like(body_blend_weight_in_bbox)
+        tree = KDTree(body_points_in_bbox, leaf_size=2)
+        for i in tqdm.trange(len(body_points_in_bbox)):
+            ind = tree.query_radius(body_points_in_bbox[i:(i+1)], r=0.035)
+            corr_bw_[i] = np.mean(body_blend_weight_in_bbox[ind[0]])
+        body_blend_weight_in_bbox = np.copy(corr_bw_)
+    body_blend_weight[point_idx_inside_bbox] = body_blend_weight_in_bbox
+
+    with open('./debug/debug_body_facial_bw.obj', 'w') as fp:
+        for p, w in zip(body_points, body_blend_weight):
+            fp.write('v %f %f %f %f %f %f\n' % (p[0], p[1], p[2], w, w, w))
+
+    tree = KDTree(body_points, leaf_size=2)
+    corr_bw = np.zeros([len(head_facial_points)], dtype=np.float32)
+    for i in range(len(head_facial_points)):
+        _, idx = tree.query(head_facial_points[i:(i+1)], k=4)
+        corr_bw[i] = np.mean(body_blend_weight[idx])
+    # corr_bw = np.clip(corr_bw*1.2 - 0.15, 0, 1)
+
+    corr_bw_bmin, corr_bw_bmax = np.percentile(corr_bw, 5), np.percentile(corr_bw, 95)
+    corr_bw = np.clip((corr_bw-corr_bw_bmin)/(corr_bw_bmax-corr_bw_bmin), 0, 1)
+    with open('./debug/debug_head_facial_bw.obj', 'w') as fp:
+        for p, w in zip(head_facial_points, corr_bw):
+            fp.write('v %f %f %f %f %f %f\n' % (p[0], p[1], p[2], w, w, w))
+    return corr_bw
+
+
+def estimate_color_transfer(head_facial_points, body_facial_points, head_facial_color, body_facial_color, head_facial_opacity):
+    head_facial_color = head_facial_color * 0.28209479177387814 + 0.5
+    body_facial_color = body_facial_color * 0.28209479177387814 + 0.5
+
+    corr_idx, _ = search_nearest_correspondence(head_facial_points, body_facial_points)
+    corr_color = body_facial_color[corr_idx]
+
+    opacity = 1/(1+np.exp(-head_facial_opacity))
+    weight = np.float32(opacity > 0.35)
+    head_facial_color = head_facial_color.reshape(len(head_facial_color), 3) * weight.reshape([-1, 1])
+    corr_color = corr_color.reshape(len(corr_color), 3) * weight.reshape([-1, 1])
+
+    head_facial_color = np.concatenate([head_facial_color, np.zeros_like(head_facial_color[:, :1])], axis=1)
+    corr_color = np.concatenate([corr_color, np.zeros_like(corr_color[:, :1])], axis=1)
+
+    transfer = nn.Parameter(torch.eye(4, dtype=torch.float32))
+    head_facial_color_th = torch.from_numpy(head_facial_color).float()
+    corr_color_th = torch.from_numpy(corr_color).float()
+    weight_th = torch.from_numpy(weight).float()
+    optim = torch.optim.Adam([transfer], lr=1e-2)
+
+    for i in range(500):
+        optim.zero_grad()
+        loss = torch.mean(torch.abs(corr_color_th - torch.matmul(head_facial_color_th, transfer.permute(1, 0)))*weight_th)
+        loss = loss + torch.sum(torch.square(transfer - torch.eye(4, dtype=torch.float32))) * 5e-2
+        if i % 25 == 0:
+            print(loss.item())
+        loss.backward()
+        optim.step()
+    transfer = transfer.detach().cpu().numpy()
+    print(transfer)
+
+    # with open('./debug/debug_body_facial_color_updated.obj', 'w') as fp:
+    #     for p, c in zip(body_facial_points, body_facial_color):
+    #         # c = c * 0.28209479177387814 + 0.5
+    #         c = np.clip(c, 0, 1)
+    #         fp.write('v %f %f %f %f %f %f\n' % (p[0], p[1], p[2], c[0], c[1], c[2]))
+    # with open('./debug/debug_head_facial_color_updated.obj', 'w') as fp:
+    #     head_facial_color = np.matmul(head_facial_color, transfer)
+    #     for p, c, w in zip(head_facial_points, head_facial_color, weight):
+    #         if w < 0.1:
+    #             continue
+    #         # c = c * 0.28209479177387814 + 0.5
+    #         c = np.clip(c, 0, 1)
+    #         fp.write('v %f %f %f %f %f %f\n' % (p[0], p[1], p[2], c[0], c[1], c[2]))
+    # import pdb; pdb.set_trace()
+
+    return transfer
+
+
+def blend_color(head_facial_color, body_facial_color, blend_weight):
+    blend_weight = blend_weight.reshape([len(blend_weight)] + [1]*(len(head_facial_color.shape)-1))
+    result = head_facial_color * blend_weight + body_facial_color * (1-blend_weight)
+    return result
+
+def save_body_face_stitching_data(
+        result_path, smplx_to_faceverse, residual_transf, body_nonface_mask, head_nonface_mask,
+        head_facial_idx, body_facial_idx, corr_idx, face_color_bw, color_transfer):
+    # os.makedirs('./data/%s' % result_suffix, exist_ok=True)
+    # np.savez('./data/%s/body_face_blending_param.npz' % result_suffix,
+    #          smplx_to_faceverse=smplx_to_faceverse.astype(np.float32),
+    #          residual_transf=residual_transf.astype(np.float32),
+    #          body_nonface_mask=body_nonface_mask.astype(np.int32),
+    #          head_facial_idx=head_facial_idx.astype(np.int32),
+    #          body_facial_idx=body_facial_idx.astype(np.int32),
+    #          head_body_facial_corr_idx=corr_idx.astype(np.int32),
+    #          face_color_bw=face_color_bw.astype(np.float32),
+    #          color_transfer=color_transfer.astype(np.float32))
+    head_color_bw = np.zeros([len(head_nonface_mask)])
+    head_color_bw[head_facial_idx] = face_color_bw
+    head_corr_idx = np.zeros([len(head_nonface_mask)])
+    head_corr_idx[head_facial_idx] = body_facial_idx[corr_idx]
+    np.savez(result_path,
+             smplx_to_faceverse=smplx_to_faceverse.astype(np.float32),
+             residual_transf=residual_transf.astype(np.float32),
+             body_nonface_mask=body_nonface_mask.astype(np.int32),
+             head_nonface_mask=head_nonface_mask.astype(np.int32),
+             head_facial_idx=head_facial_idx.astype(np.int32),
+             body_facial_idx=body_facial_idx.astype(np.int32),
+             head_body_corr_idx=head_corr_idx.astype(np.int32),
+             head_color_bw=head_color_bw.astype(np.float32),
+             color_transfer=color_transfer.astype(np.float32))
+    return
+
+
+def manual_refine_facial_cropping(head_facial_points, head_facial_idx, head_facial_colors, body_facial_points, body_facial_idx, body_facial_colors):
+    def _save_points_as_obj(fpath, points, points_color):
+        points_color = np.clip(points_color, 0, 1)
+        with open(fpath, 'w') as fp:
+            for p, c in zip(points, points_color):
+                fp.write('v %f %f %f %f %f %f\n' % (p[0], p[1], p[2], c[0], c[1], c[2]))
+        return
+
+    _save_points_as_obj('./debug/head_facial_points.obj', head_facial_points, head_facial_colors)
+    _save_points_as_obj('./debug/body_facial_points.obj', body_facial_points, body_facial_colors)
+    # trimesh.Trimesh(vertices=head_facial_points, vertex_colors=head_facial_colors).export('./debug/head_facial_points.obj')
+    # trimesh.Trimesh(vertices=body_facial_points, vertex_colors=body_facial_colors).export('./debug/body_facial_points.obj')
+    if True:
+        print('Saving facial points cropped by algorithms. Please remove unnecessary points manually!')
+        import pdb; pdb.set_trace()
+
+    head_facial_points_ = trimesh.load('./debug/head_facial_points.obj').vertices
+    body_facial_points_ = trimesh.load('./debug/body_facial_points.obj').vertices
+    _, head_nndist = search_nearest_correspondence(head_facial_points, head_facial_points_)
+    _, body_nndist = search_nearest_correspondence(body_facial_points, body_facial_points_)
+    head_flag = head_nndist < 1e-4
+    body_flag = body_nndist < 1e-4
+    return head_facial_points[head_flag], head_facial_idx[head_flag], body_facial_points[body_flag], body_facial_idx[body_flag]
+
+
+def stitch_body_and_head(ref_body_gaussian_path, ref_head_gaussian_path,
+                         ref_body_param_path, ref_head_param_path,
+                         smplx2faceverse_path, result_folder):
+    device = torch.device("cuda")
+
+    body_gaussians = load_gaussians_from_ply(ref_body_gaussian_path)
+    head_gaussians = load_gaussians_from_ply(ref_head_gaussian_path)
+    global_orient, transl, body_pose, betas = load_body_params(ref_body_param_path)
+    head_pose, head_scale, id_coeff, exp_coeff = load_face_params(ref_head_param_path)
+    smplx_to_faceverse = np.load(smplx2faceverse_path)
+    faceverse_to_smplx = np.linalg.inv(smplx_to_faceverse)
+
+    smpl = SmplTorch(model_file='./AnimatableGaussians/smpl_files/smplx/SMPLX_NEUTRAL.npz')
+    smpl_verts, head_joint_transfmat = get_smpl_verts_and_head_transformation(
+        smpl, global_orient, body_pose, transl, betas)
+    smpl_head_vids = dict(np.load('./data/smpl_models/smplx_head_vidx_and_blendweight.npz'))['head_ids']
+    smpl_head_verts = smpl_verts[smpl_head_vids]
+
+    model_dict = np.load('./data/faceverse_models/faceverse_simple_v2.npy', allow_pickle=True).item()
+    faceverse_model = FaceVerseModel(model_dict, batch_size=1)
+    faceverse_model.init_coeff_tensors(
+        id_coeff=torch.from_numpy(id_coeff).reshape([1, -1]).to(device),
+        scale_coeff=torch.from_numpy(head_scale).reshape([1, 1]).to(device),
+    )
+    faceverse_verts = faceverse_model.forward()['v'][0].detach().cpu().numpy()
+    faceverse_verts = transform_faceverse_to_live_body_space(faceverse_verts, faceverse_to_smplx, head_joint_transfmat)
+
+    livehead2livebody = calc_livehead2livebody(
+        head_pose, smplx_to_faceverse, head_joint_transfmat)
+    head_gaussians_xyz = np.matmul(head_gaussians.xyz, livehead2livebody[:3, :3].transpose()) \
+                         + livehead2livebody[:3, 3].reshape(1, 3)
+
+    # head_facial_points, head_facial_idx = crop_facial_area(smpl_head_verts, head_gaussians_xyz)
+    # body_facial_points, body_facial_idx = crop_facial_area(smpl_head_verts, body_gaussians.xyz)
+    head_facial_points, head_facial_idx = crop_facial_area2(smpl_verts, smpl_head_vids, head_gaussians_xyz)
+    body_facial_points, body_facial_idx = crop_facial_area2(smpl_verts, smpl_head_vids, body_gaussians.xyz)
+    
+    residual_transf = np.eye(4)
+    head_facial_points, head_facial_idx, body_facial_points, body_facial_idx = manual_refine_facial_cropping(
+            head_facial_points, head_facial_idx, SH2RGB(head_gaussians.features_dc[head_facial_idx]), 
+            body_facial_points, body_facial_idx, SH2RGB(body_gaussians.features_dc[body_facial_idx]))
+    while True:
+        for _ in tqdm.trange(4, desc='Fitting residual transformation'):
+            corr_idx, _ = search_nearest_correspondence(head_facial_points, body_facial_points)
+            corr = body_facial_points[corr_idx]
+            transf = estimate_rigid_transformation(head_facial_points, corr)
+            residual_transf = np.matmul(transf, residual_transf)
+            head_facial_points = np.matmul(head_facial_points, transf[:3, :3].transpose()) + transf[:3, 3].reshape(1, 3)
+        if_crop_well = input('If the facial cropping is good enough? (y/n): ')
+        if if_crop_well == 'y':
+            break
+        else:
+            head_facial_points, head_facial_idx, body_facial_points, body_facial_idx = manual_refine_facial_cropping(
+                head_facial_points, head_facial_idx, SH2RGB(head_gaussians.features_dc[head_facial_idx]), 
+                body_facial_points, body_facial_idx, SH2RGB(body_gaussians.features_dc[body_facial_idx]))
+    
+    # head_facial_points, head_facial_idx, body_facial_points, body_facial_idx = manual_refine_facial_cropping(
+    #     head_facial_points, head_facial_idx, SH2RGB(head_gaussians.features_dc[head_facial_idx]), 
+    #     body_facial_points, body_facial_idx, SH2RGB(body_gaussians.features_dc[body_facial_idx]))
+    
+    # 更改一下逻辑，改成直到对齐为止。
+    
+    
+    print(np.matmul(residual_transf, livehead2livebody))
+    residual_transf = np.matmul(np.linalg.inv(livehead2livebody), np.matmul(residual_transf, livehead2livebody))
+    corr_idx, _ = search_nearest_correspondence(head_facial_points, body_facial_points)
+
+    # head_gaussians_xyz = np.matmul(head_gaussians_xyz, residual_transf[:3, :3].transpose()) + residual_transf[:3, 3].reshape(1, 3)
+    # faceverse_verts = np.matmul(faceverse_verts, residual_transf[:3, :3].transpose()) + residual_transf[:3, 3].reshape(1, 3)
+
+    # total_transf = np.matmul(residual_transf, livehead2livebody)
+    total_transf = np.matmul(livehead2livebody, residual_transf)
+    print(total_transf)
+
+    color_transfer = estimate_color_transfer(
+        head_facial_points, body_facial_points,
+        head_gaussians.features_dc[head_facial_idx], body_gaussians.features_dc[body_facial_idx],
+        head_gaussians.opacities[head_facial_idx]
+    )
+    # face_color_bw = get_face_blend_weight(head_facial_points, smpl_verts, sigma=0.015)
+    face_color_bw = get_face_blend_weight2(head_facial_points, body_gaussians.xyz, body_facial_idx)
+
+    body_nonface_mask = np.ones([len(body_gaussians.xyz)], dtype=np.bool_)
+    body_nonface_mask[body_facial_idx] = 0
+    head_nonface_mask = np.ones([len(head_gaussians.xyz)], dtype=np.bool_)
+    head_nonface_mask[head_facial_idx] = 0
+
+    save_body_face_stitching_data(
+        os.path.join(result_folder, 'body_head_blending_param.npz'),
+        smplx_to_faceverse, residual_transf, body_nonface_mask, head_nonface_mask,
+        head_facial_idx, body_facial_idx, corr_idx, face_color_bw, color_transfer)
+
+    body_gaussians = apply_transformation_to_gaussians(body_gaussians, np.eye(4))
+    head_gaussians = apply_transformation_to_gaussians(head_gaussians, total_transf, np.eye(3))
+    body_gaussians_wo_face = select_gaussians(body_gaussians, body_nonface_mask)
+    head_gaussians_face_only = select_gaussians(head_gaussians, head_facial_idx)
+    head_gaussians_face_only_new_color = blend_color(
+        head_gaussians_face_only.features_dc, body_gaussians.features_dc[body_facial_idx][corr_idx], face_color_bw)
+    head_gaussians_face_only_new_xyz = blend_color(
+        head_gaussians_face_only.xyz, body_gaussians.xyz[body_facial_idx][corr_idx], face_color_bw)
+    head_gaussians_face_only_new_opacities = blend_color(
+        head_gaussians_face_only.opacities, body_gaussians.opacities[body_facial_idx][corr_idx], face_color_bw)
+    head_gaussians_face_only_new_scales = blend_color(
+        head_gaussians_face_only.scales, body_gaussians.scales[body_facial_idx][corr_idx], face_color_bw)
+
+    head_gaussians_face_only = update_gaussian_attributes(
+        head_gaussians_face_only, new_rgb=head_gaussians_face_only_new_color,
+        new_xyz=head_gaussians_face_only_new_xyz, new_opacity=head_gaussians_face_only_new_opacities,
+        new_scale=head_gaussians_face_only_new_scales)
+
+    full_gaussians = combine_gaussians([body_gaussians_wo_face, head_gaussians_face_only])
+    save_gaussians_as_ply(os.path.join(result_folder, 'full_gaussians.ply'), full_gaussians)

render_utils/stitch_funcs.py

@@ -0,0 +1,145 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import argparse
+import tqdm
+import json
+import cv2 as cv
+import os, glob
+import math
+
+
+from render_utils.lib.utils.graphics_utils import focal2fov, getProjectionMatrix
+from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
+
+
+def render3(
+    gaussian_vals: dict,
+    bg_color: torch.Tensor,
+    extr: torch.Tensor,
+    intr: torch.Tensor,
+    img_w: int,
+    img_h: int,
+    scaling_modifier = 1.0,
+    override_color = None,
+    compute_cov3D_python = False
+):
+    means3D = gaussian_vals['positions']
+    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
+    screenspace_points = torch.zeros_like(means3D, dtype = means3D.dtype, requires_grad = True, device = "cuda") + 0
+    try:
+        screenspace_points.retain_grad()
+    except:
+        pass
+    means2D = screenspace_points
+    opacity = gaussian_vals['opacity']
+
+    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
+    # scaling / rotation by the rasterizer.
+    scales = None
+    rotations = None
+    cov3D_precomp = None
+    scales = gaussian_vals['scales']
+    rotations = gaussian_vals['rotations']
+
+    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
+    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
+    shs = None
+    # colors_precomp = None
+    # if override_color is None:
+    #     shs = gaussian_vals['shs']
+    # else:
+    #     colors_precomp = override_color
+    if 'colors' in gaussian_vals:
+        colors_precomp = gaussian_vals['colors']
+    else:
+        colors_precomp = None
+
+    # Set up rasterization configuration
+    FoVx = focal2fov(intr[0, 0].item(), img_w)
+    FoVy = focal2fov(intr[1, 1].item(), img_h)
+    tanfovx = math.tan(FoVx * 0.5)
+    tanfovy = math.tan(FoVy * 0.5)
+    world_view_transform = extr.transpose(1, 0).cuda()
+    projection_matrix = getProjectionMatrix(znear = 0.1, zfar = 100, fovX = FoVx, fovY = FoVy, K = intr, img_w = img_w, img_h = img_h).transpose(0, 1).cuda()
+    full_proj_transform = (world_view_transform.unsqueeze(0).bmm(projection_matrix.unsqueeze(0))).squeeze(0)
+    camera_center = torch.linalg.inv(extr)[:3, 3]
+
+    raster_settings = GaussianRasterizationSettings(
+        image_height = img_h,
+        image_width = img_w,
+        tanfovx = tanfovx,
+        tanfovy = tanfovy,
+        bg = bg_color,
+        scale_modifier = scaling_modifier,
+        viewmatrix = world_view_transform,
+        projmatrix = full_proj_transform,
+        sh_degree = gaussian_vals['max_sh_degree'],
+        campos = camera_center,
+        prefiltered = False,
+        debug = False
+    )
+
+    rasterizer = GaussianRasterizer(raster_settings = raster_settings)
+
+    # Rasterize visible Gaussians to image, obtain their radii (on screen).
+    rendered_image, radii = rasterizer(
+        means3D = means3D,
+        means2D = means2D,
+        shs = shs,
+        colors_precomp = colors_precomp,
+        opacities = opacity,
+        scales = scales,
+        rotations = rotations,
+        cov3D_precomp = cov3D_precomp)
+
+    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
+    # They will be excluded from value updates used in the splitting criteria.
+    return {
+        "render": rendered_image,
+        "viewspace_points": screenspace_points,
+        "visibility_filter": radii > 0,
+        "radii": radii
+    }
+
+
+def blend_color(head_facial_color, body_facial_color, blend_weight):
+    blend_weight = blend_weight.reshape([len(blend_weight)] + [1]*(len(head_facial_color.shape)-1))
+    result = head_facial_color * blend_weight + body_facial_color * (1-blend_weight)
+    return result
+
+
+@torch.no_grad()
+def paste_back_with_linear_interp(pasteback_scale, pasteback_center, src, tgt_size):
+    pasteback_topleft = [pasteback_center[0] - src.shape[1]/2/pasteback_scale, 
+                         pasteback_center[1] - src.shape[0]/2/pasteback_scale]
+
+    h, w = src.shape[0], src.shape[1]
+    grayscale = False
+    if len(src.shape) == 2:
+        src = src.reshape([h, w, 1])
+        grayscale = True
+    src = torch.from_numpy(src)
+    src = src.permute(2, 0, 1).unsqueeze(0)
+    grid = torch.meshgrid(torch.arange(0, tgt_size[0]), torch.arange(0, tgt_size[1]), indexing='xy')
+    grid = torch.stack(grid, dim = -1).float().to(src.device).unsqueeze(0)
+    grid[..., 0] = (grid[..., 0] - pasteback_topleft[0]) * pasteback_scale
+    grid[..., 1] = (grid[..., 1] - pasteback_topleft[1]) * pasteback_scale
+
+    grid[..., 0] = grid[..., 0] / (src.shape[-1] / 2.0) - 1.0
+    grid[..., 1] = grid[..., 1] / (src.shape[-2] / 2.0) - 1.0
+    out = F.grid_sample(src, grid, align_corners = True)
+    out = out[0].detach().permute(1, 2, 0).cpu().numpy()
+    if grayscale:
+        out = out[:, :, 0]
+    return out
+
+
+def soften_blending_mask(blending_mask, valid_mask):
+    blending_mask = np.clip(blending_mask*2.0, 0.0, 1.0)
+    blending_mask = cv.erode(blending_mask, np.ones((5, 5))) * valid_mask
+    blending_mask_bk = np.copy(blending_mask)
+    blending_mask = cv.blur(blending_mask*valid_mask, (25, 25))
+    valid_mask = cv.blur(valid_mask, (25, 25))
+    blending_mask = blending_mask / (valid_mask + 1e-6) * blending_mask_bk
+    return blending_mask