render_utils/stitch_body_and_head.py

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)