avatar_generator.py

from calendar import c
import os
# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
# os.environ['TORCH_USE_CUDA_DSA'] = '1'
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
import yaml
import shutil
import collections
import torch
import torch.utils.data
import torch.nn.functional as F
import numpy as np
import cv2 as cv
import glob
import datetime
import trimesh
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
import importlib
# import config
from omegaconf import OmegaConf
import json
import math
import cv2

# AnimatableGaussians part
from AnimatableGaussians.network.lpips import LPIPS
from AnimatableGaussians.dataset.dataset_pose import PoseDataset
import AnimatableGaussians.utils.net_util as net_util
# import AnimatableGaussians.utils.visualize_util as visualize_util
from AnimatableGaussians.utils.camera_dir import get_camera_dir
from AnimatableGaussians.utils.renderer import Renderer
from AnimatableGaussians.utils.net_util import to_cuda
from AnimatableGaussians.utils.obj_io import save_mesh_as_ply
from AnimatableGaussians.gaussians.obj_io import save_gaussians_as_ply
import AnimatableGaussians.config as ag_config

# Gaussian-Head-Avatar part
from GHA.config.config import config_reenactment
from GHA.lib.dataset.Dataset import ReenactmentDataset
from GHA.lib.dataset.DataLoaderX import DataLoaderX
from GHA.lib.module.GaussianHeadModule import GaussianHeadModule
from GHA.lib.module.SuperResolutionModule import SuperResolutionModule
from GHA.lib.module.CameraModule import CameraModule
from GHA.lib.recorder.Recorder import ReenactmentRecorder
from GHA.lib.apps.Reenactment import Reenactment
from GHA.lib.utils.graphics_utils import getWorld2View2, getProjectionMatrix

# cat utils
from calc_offline_rendering_param import calc_offline_rendering_param
from calc_offline_rendering_param import load_camera_data
from render_utils.lib.networks.smpl_torch import SmplTorch
from render_utils.lib.utils.gaussian_np_utils import load_gaussians_from_ply
from render_utils.stitch_body_and_head import load_body_params, load_face_params, get_smpl_verts_and_head_transformation, calc_livehead2livebody
from render_utils.stitch_funcs import soften_blending_mask,paste_back_with_linear_interp


import ipdb

class Avatar:
    def __init__(self, config):
        self.config = config
        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
        
        # animateble gaussians part init
        self.body = config.animatablegaussians
        self.body.mode = 'test'
        ag_config.set_opt(self.body)
        avatar_module = self.body['model'].get('module', 'AnimatableGaussians.network.avatar')
        print('Import AvatarNet from %s' % avatar_module)
        AvatarNet = importlib.import_module(avatar_module).AvatarNet
        self.avatar_net = AvatarNet(self.body.model).to(self.device)
        self.random_bg_color = self.body['train'].get('random_bg_color', True)
        self.bg_color = (1., 1., 1.)
        self.bg_color_cuda = torch.from_numpy(np.asarray(self.bg_color)).to(torch.float32).to(self.device)
        self.loss_weight = self.body['train']['loss_weight']
        self.finetune_color = self.body['train']['finetune_color']
        print('# Parameter number of AvatarNet is %d' % (sum([p.numel() for p in self.avatar_net.parameters()])))
        
        # gaussian head avatar part init
        self.head = config.gha
        # cat utils part init
        self.cat = config.cat
        
    def build_dataset(self, body_pose_path=None, face_exp_path=None):
         # build body_dataset
         
        if body_pose_path is not None:
            self.body['test']['pose_data']['data_path'] = body_pose_path
            body_pose = np.load(body_pose_path, allow_pickle = True)
            # print('body_pose keys:', body_pose.keys())
            # print('body_pose shape:', body_pose['poses'].shape)
            self.body['test']['pose_data']['frame_range'] = [0,body_pose['poses'].shape[0]]
            
        dataset_module = self.body.get('dataset', 'MvRgbDatasetAvatarReX')
        MvRgbDataset = importlib.import_module('AnimatableGaussians.dataset.dataset_mv_rgb').__getattribute__(dataset_module)
        self.body_training_dataset = MvRgbDataset(**self.body['train']['data'], training = False)
        if self.body['test'].get('n_pca', -1) >= 1:
            self.body_training_dataset.compute_pca(n_components = self.body['test']['n_pca'])
        if 'pose_data' in self.body.test:
            testing_dataset = PoseDataset(**self.body['test']['pose_data'], smpl_shape = self.body_training_dataset.smpl_data['betas'][0])
            dataset_name = testing_dataset.dataset_name
            seq_name = testing_dataset.seq_name
        else:
            # throw an error
            raise ValueError('No pose data in test config')
        self.body_dataset = testing_dataset
        iter_idx = self.load_ckpt(self.body['test']['prev_ckpt'], False)[1]
        
        
        self.head_config = config_reenactment()
        self.head_config.load(self.head.config_path)
        if face_exp_path is not None:
            self.head_config.cfg.dataset.exp_path = face_exp_path
        self.head_config.freeze()
        self.head_config = self.head_config.get_cfg()        
        # build face dataset
        self.head_dataset = ReenactmentDataset(self.head_config.dataset)
        self.head_dataloader = DataLoaderX(self.head_dataset, batch_size=1, shuffle=False, pin_memory=True) 

        # device = torch.device('cuda:%d' % cfg.gpu_id)

        gaussianhead_state_dict = torch.load(self.head_config.load_gaussianhead_checkpoint, map_location=lambda storage, loc: storage)
        self.gaussianhead = GaussianHeadModule(self.head_config.gaussianheadmodule, 
                                              xyz=gaussianhead_state_dict['xyz'], 
                                              feature=gaussianhead_state_dict['feature'],
                                              landmarks_3d_neutral=gaussianhead_state_dict['landmarks_3d_neutral']).to(self.device)
        self.gaussianhead.load_state_dict(gaussianhead_state_dict)

        self.supres = SuperResolutionModule(self.head_config.supresmodule).to(self.device)
        self.supres.load_state_dict(torch.load(self.head_config.load_supres_checkpoint, map_location=lambda storage, loc: storage))

        self.head_camera = CameraModule()
        self.head_recorder = ReenactmentRecorder(self.head_config.recorder)
        
    def render_all(self):
        # len = short one
        lenth = min(len(self.body_dataset), len(self.head_dataloader))
        # build a tqdm bar
        for idx in tqdm(range(lenth)):
            self.reder_frame(idx)
        
        # for idx in range(lenth):
        #     self.reder_frame(idx)
        
    def reder_frame(self, idx):
        # 渲染身体和各种mask
        body_output = self.build_body(idx)
        # 计算头的渲染参数
        head_param = self.build_param(idx,body_output)
        # 渲染头
        head_output = self.build_head(idx, head_param)
        # 把头和身体拼接起来
        body_rendering= body_output['rgb_map_wo_hand'].astype(np.float32) / 255.0
        # save body_rendering
        # cv.imwrite('./output' + '/body_rgb_%08d.jpg' % idx, (body_output['rgb_map']).astype(np.uint8))
        # cv.imwrite('./output' + '/body_rgb_wo_hand%08d.jpg' % idx, (body_output['rgb_map_wo_hand']).astype(np.uint8))
        body_mask = body_output['mask_map'].astype(np.float32) / 255.0
        body_torso_mask = body_output['torso_map'].astype(np.float32) / 255.0
        head_rendering = head_output['render_images'].astype(np.float32) / 255.0
        head_blending_mask = head_output['render_bw'].astype(np.float32) / 255.0
        body_head_blending_params = np.load(self.cat.body_head_blending_param_path)
        head_offline_rendering_param = head_param
        stitch_output = self.stich_head_body(body_rendering, body_mask, body_torso_mask, head_rendering, head_blending_mask, body_head_blending_params, head_offline_rendering_param)
        cv.imwrite('./output' + '/%08d.jpg' % idx, stitch_output)
        
        # 渲染手和手的mask
        
        # 把手拼上去
           
        return stitch_output
        
        pass
    
    def load_ckpt(self, path, load_optm = True):
        print('Loading networks from ', path + '/net.pt')
        net_dict = torch.load(path + '/net.pt')
        if 'avatar_net' in net_dict:
            self.avatar_net.load_state_dict(net_dict['avatar_net'])
        else:
            print('[WARNING] Cannot find "avatar_net" from the network checkpoint!')
        epoch_idx = net_dict['epoch_idx']
        iter_idx = net_dict['iter_idx']

        # if load_optm and os.path.exists(path + '/optm.pt'):
        #     print('Loading optimizers from ', path + '/optm.pt')
        #     optm_dict = torch.load(path + '/optm.pt')
        #     if 'avatar_net' in optm_dict:
        #         self.optm.load_state_dict(optm_dict['avatar_net'])
        #     else:
        #         print('[WARNING] Cannot find "avatar_net" from the optimizer checkpoint!')

        return epoch_idx, iter_idx
    
    @torch.no_grad()
    def build_body(self,idx):
        self.avatar_net.eval()
        geo_renderer = None
        item_0 = self.body_dataset.getitem(0, training = False)
        object_center = item_0['live_bounds'].mean(0)
        global_orient = item_0['global_orient'].cpu().numpy() if isinstance(item_0['global_orient'], torch.Tensor) else item_0['global_orient']
        use_pca = self.body['test'].get('n_pca', -1) >= 1
        # set x and z to 0
        global_orient[0] = 0
        global_orient[2] = 0
        
        global_orient = cv.Rodrigues(global_orient)[0]
        time_start = torch.cuda.Event(enable_timing = True)
        time_start_all = torch.cuda.Event(enable_timing = True)
        time_end = torch.cuda.Event(enable_timing = True)
        
        if self.body['test'].get('fix_hand', False):
            self.avatar_net.generate_mean_hands()
            
        img_scale = self.body['test'].get('img_scale', 1.0)
        view_setting = self.body['test'].get('view_setting', 'free')
        extr, intr, img_h, img_w = get_camera_dir(idx, object_center, global_orient, img_scale, view_setting)
        w2c = extr
        c2w = np.linalg.inv(w2c)
        pos = c2w[:3, 3]
        rot = c2w[:3, :3]
        serializable_array_2d = [x.tolist() for x in rot]
        camera_entry = {
            'width': int(img_w),
            'height': int(img_h),
            'position': pos.tolist(),
            'rotation': serializable_array_2d,
            'fy': float(intr[1, 1]),
            'fx': float(intr[0, 0]),
        }
        
        getitem_func = self.body_dataset.getitem_fast if hasattr(self.body_dataset, 'getitem_fast') else self.body_dataset.getitem
        item = getitem_func(
            idx,
            training = False,
            extr = extr,
            intr = intr,
            img_w = img_w,
            img_h = img_h
        )
        items = to_cuda(item, add_batch = False)
        
        if 'smpl_pos_map' not in items:
            self.avatar_net.get_pose_map(items)

        # pca
        if use_pca:
            mask = self.body_training_dataset.pos_map_mask
            live_pos_map = items['smpl_pos_map'].permute(1, 2, 0).cpu().numpy()
            front_live_pos_map, back_live_pos_map = np.split(live_pos_map, [3], 2)
            pose_conds = front_live_pos_map[mask]
            new_pose_conds = self.body_training_dataset.transform_pca(pose_conds, sigma_pca = float(self.body['test'].get('sigma_pca', 2.)))
            front_live_pos_map[mask] = new_pose_conds
            live_pos_map = np.concatenate([front_live_pos_map, back_live_pos_map], 2)
            items.update({
                'smpl_pos_map_pca': torch.from_numpy(live_pos_map).to(self.device).permute(2, 0, 1)
            })
            
        # print items
        # print(items.keys())
        # print(items.values())
        # exit()
        
        # get render result
        output = self.avatar_net.render(items, bg_color = self.bg_color, use_pca = use_pca)
        output_wo_hand = self.avatar_net.render_wo_hand(items, bg_color = self.bg_color, use_pca = use_pca)
        mask_output = self.avatar_net.render_mask(items, bg_color = self.bg_color, use_pca = use_pca)
        
        # do some postprocess
        rgb_map_wo_hand = output_wo_hand['rgb_map']
        
        full_body_mask = mask_output['full_body_rgb_map']
        full_body_mask.clip_(0., 1.)
        full_body_mask = (full_body_mask * 255).to(torch.uint8)
        
        hand_only_mask = mask_output['hand_only_rgb_map']
        hand_only_mask.clip_(0., 1.)
        hand_only_mask = (hand_only_mask * 255).to(torch.uint8)
        
        # build the covered hand mask and the hand visualbility flag
        body_red_mask = (mask_output['full_body_rgb_map'] - torch.tensor([1., 0., 0.], device = mask_output['full_body_rgb_map'].device))
        body_red_mask = (body_red_mask*body_red_mask).sum(dim=2) < 0.01 # need save
        
        hand_red_mask = (mask_output['hand_only_rgb_map'] - torch.tensor([1., 0., 0.], device = mask_output['hand_only_rgb_map'].device))
        hand_red_mask = (hand_red_mask*hand_red_mask).sum(dim=2) < 0.0
        if_mask_r_hand = abs(body_red_mask.sum() - hand_red_mask.sum()) / hand_red_mask.sum() > 0.95
        if_mask_r_hand = if_mask_r_hand.cpu().numpy()
        
        body_blue_mask = (mask_output['full_body_rgb_map'] - torch.tensor([0., 0., 1.], device = mask_output['full_body_rgb_map'].device))
        body_blue_mask = (body_blue_mask*body_blue_mask).sum(dim=2) < 0.01 # need save
        
        hand_blue_mask = (mask_output['hand_only_rgb_map'] - torch.tensor([0., 0., 1.], device = mask_output['hand_only_rgb_map'].device))
        hand_blue_mask = (hand_blue_mask*hand_blue_mask).sum(dim=2) < 0.01
        
        if_mask_l_hand = abs(body_blue_mask.sum() - hand_blue_mask.sum()) / hand_blue_mask.sum() > 0.95
        if_mask_l_hand = if_mask_l_hand.cpu().numpy()
        
        # 保存左右手被遮挡部分的mask
        red_mask = hand_red_mask ^ (hand_red_mask & body_red_mask)
        blue_mask = hand_blue_mask ^ (hand_blue_mask & body_blue_mask)
        all_mask = red_mask | blue_mask
        
        all_mask = (all_mask * 255).to(torch.uint8) 
        r_hand_mask = (body_red_mask * 255).to(torch.uint8)
        l_hand_mask = (body_blue_mask * 255).to(torch.uint8)
        hand_visual = [if_mask_r_hand, if_mask_l_hand]
        
        # build sleeve mask
        mask = (r_hand_mask>128) | (l_hand_mask>128)| (all_mask>128)
        mask = mask.cpu().numpy().astype(np.uint8)
        # 定义一个结构元素，可以调整其大小以改变膨胀的程度
        kernel = np.ones((5, 5), np.uint8)
        # 应用膨胀操作
        mask = cv.dilate(mask, kernel, iterations=3)
        mask = torch.tensor(mask).to(self.device)
        
        left_hand_mask = mask_output['left_hand_rgb_map']
        left_hand_mask.clip_(0., 1.)
        # non white part is mask
        left_hand_mask = (torch.tensor([1., 1., 1.], device = left_hand_mask.device) - left_hand_mask)
        left_hand_mask = (left_hand_mask*left_hand_mask).sum(dim=2) > 0.01
        # dele two hand mask
        left_hand_mask = left_hand_mask & ~mask
        
        right_hand_mask = mask_output['right_hand_rgb_map']
        right_hand_mask.clip_(0., 1.)
        right_hand_mask = (torch.tensor([1., 1., 1.], device = right_hand_mask.device) - right_hand_mask)
        right_hand_mask = (right_hand_mask*right_hand_mask).sum(dim=2) > 0.01
        right_hand_mask = right_hand_mask & ~mask
        
        left_sleeve_mask = (left_hand_mask * 255).to(torch.uint8)
        right_sleeve_mask = (right_hand_mask * 255).to(torch.uint8)
        
        # 利用 r_hand_mask 和 l_hand_mask，将wo_hand图像中的mask部分覆盖rgb_map
        rgb_map = output['rgb_map']
        rgb_map.clip_(0., 1.)
        rgb_map = (rgb_map * 255).to(torch.uint8).cpu().numpy()
        
        rgb_map_wo_hand = output_wo_hand['rgb_map']
        rgb_map_wo_hand.clip_(0., 1.)
        rgb_map_wo_hand = (rgb_map_wo_hand * 255).to(torch.uint8).cpu().numpy()
        
        r_mask = (r_hand_mask>128).cpu().numpy()
        l_mask = (l_hand_mask>128).cpu().numpy()
        mask = r_mask | l_mask
        mask = mask.astype(np.uint8)
        # 定义一个结构元素，可以调整其大小以改变膨胀的程度
        kernel = np.ones((5, 5), np.uint8)
        # 应用膨胀操作
        mask = cv.dilate(mask, kernel, iterations=3)
        mask = mask.astype(np.bool_)
        mask = np.expand_dims(mask, axis=2)
        # get the final rgb_map without hand
        mix = rgb_map_wo_hand.copy() * mask + rgb_map * ~mask
        
        torso_map = output['torso_map'][:, :, 0]
        torso_map.clip_(0., 1.)
        torso_map = (torso_map * 255).to(torch.uint8).cpu().numpy()
        
        
        mask_map = output['mask_map'][:, :, 0]
        mask_map.clip_(0., 1.)
        mask_map = (mask_map * 255).to(torch.uint8).cpu().numpy()
        
        output={
            # smpl
            'betas':self.body_training_dataset.smpl_data['betas'].reshape([-1]).detach().cpu().numpy(),
            'global_orient':item['global_orient'].reshape([-1]).detach().cpu().numpy(),
            'transl':item['transl'].reshape([-1]).detach().cpu().numpy(),
            'body_pose':item['body_pose'].reshape([-1]).detach().cpu().numpy(),
            
            # camera
            'extr':extr,
            'intr':intr,
            'img_h':img_h,
            'img_w':img_w,
            'camera_entry':camera_entry,
            
            # rgb and masks
            'rgb_map':rgb_map,
            'rgb_map_wo_hand':mix,
            'torso_map':torso_map,
            'mask_map':mask_map,
            'all_mask':all_mask,
            'left_sleeve_mask':left_sleeve_mask,
            'right_sleeve_mask':right_sleeve_mask,
            'hand_visual':hand_visual        
        }
        
        return output
        
    def build_param(self,idx,body_output):
        head_gaussians = load_gaussians_from_ply(self.cat.ref_head_gaussian_path)
        head_pose, head_scale, id_coeff, exp_coeff = load_face_params(self.cat.ref_head_param_path)
        body_head_blending_params = np.load(self.cat.body_head_blending_param_path)
        smplx_to_faceverse = body_head_blending_params['smplx_to_faceverse']
        residual_transf = body_head_blending_params['residual_transf']
        head_color_bw = body_head_blending_params['head_color_bw']
        
        smpl = SmplTorch(model_file='./AnimatableGaussians/smpl_files/smplx/SMPLX_NEUTRAL.npz')
        global_orient, transl, body_pose, betas = body_output['global_orient'], body_output['transl'], body_output['body_pose'], body_output['betas']
        smpl_verts, head_joint_transfmat = get_smpl_verts_and_head_transformation(
            smpl, global_orient, body_pose, transl, betas)
        livehead2livebody = calc_livehead2livebody(head_pose, smplx_to_faceverse, head_joint_transfmat)
        total_transf = np.matmul(livehead2livebody, residual_transf)
        
        cam, image_size = load_camera_data(body_output['camera_entry'])
        cam_extr = np.matmul(cam[0], total_transf)
        cam_intr = np.copy(cam[1])
        
        pts = np.copy(head_gaussians.xyz)
        pts_proj = np.matmul(pts, cam_extr[:3, :3].transpose()) + cam_extr[:3, 3]
        pts_proj = np.matmul(pts_proj, cam_intr.transpose())
        pts_proj = pts_proj / pts_proj[:, 2:]
        
        pts_min, pts_max = np.min(pts_proj, axis=0), np.max(pts_proj, axis=0)
        pts_center = (pts_min + pts_max) // 2
        pts_size = np.max(pts_max - pts_min)
        tgt_pts_size = 350
        tgt_image_size = 512
        zoom_scale = tgt_pts_size / pts_size
        cam_intr_zoom = np.copy(cam_intr)
        cam_intr_zoom[:2] *= zoom_scale
        cam_intr_zoom[0, 2] = cam_intr_zoom[0, 2] - (pts_center[0]*zoom_scale - tgt_image_size/2)
        cam_intr_zoom[1, 2] = cam_intr_zoom[1, 2] - (pts_center[1]*zoom_scale - tgt_image_size/2)
        
        output = {
            'cam_extr':cam_extr,
            'cam_intr':cam_intr,
            'image_size':image_size,
            'cam_intr_zoom':cam_intr_zoom,
            'zoom_image_size':[tgt_image_size, tgt_image_size],
            'zoom_center':pts_center,
            'zoom_scale':zoom_scale,
            'head_pose':head_pose,
            'head_scale':head_scale,
            'head_color_bw':head_color_bw,
        }
        
        return output
    
    def build_head(self, idx, head_offline_rendering_param):
        # head_offline_rendering_param = np.load(offline_rendering_param_fpath)
        cam_extr = head_offline_rendering_param['cam_extr']
        cam_intr = head_offline_rendering_param['cam_intr']
        cam_intr_zoom = head_offline_rendering_param['cam_intr_zoom']
        zoom_image_size = head_offline_rendering_param['zoom_image_size']
        head_pose = head_offline_rendering_param['head_pose']
        head_scale = head_offline_rendering_param['head_scale']
        head_color_bw = head_offline_rendering_param['head_color_bw']
        zoom_scale = head_offline_rendering_param['zoom_scale']
        head_pose = torch.from_numpy(head_pose.astype(np.float32)).to(self.device)
        head_color_bw = torch.from_numpy(head_color_bw.astype(np.float32)).to(self.device)
        render_size = 512
        
        # data = self.head_dataloader[idx]
        data = self.head_dataset[idx]
        # add batch dim
        data = {k: v.unsqueeze(0) for k, v in data.items() if isinstance(v, torch.Tensor)}
        # print(data.keys())
        
        new_gs_camera_param_dict = self.prepare_camera_data_for_gs_rendering(cam_extr, cam_intr_zoom, render_size, render_size)
        for k in new_gs_camera_param_dict.keys():
            if isinstance(new_gs_camera_param_dict[k], torch.Tensor):
                new_gs_camera_param_dict[k] = new_gs_camera_param_dict[k].unsqueeze(0).to(self.device)
        new_gs_camera_param_dict['pose'] = head_pose.unsqueeze(0).to(self.device)

        to_cuda = ['images', 'intrinsics', 'extrinsics', 'world_view_transform', 'projection_matrix', 'full_proj_transform', 'camera_center', 
                   'pose', 'scale', 'exp_coeff', 'pose_code']
        for data_item in to_cuda:
            data[data_item] = data[data_item].to(device=self.device)

        data.update(new_gs_camera_param_dict)
        
        with torch.no_grad():
            data = self.gaussianhead.generate(data)
            data = self.head_camera.render_gaussian(data, 512)
            render_images = data['render_images']
            supres_images = self.supres(render_images)
            data['supres_images'] = supres_images
            data['bg_color'] = torch.zeros([1, 32], device=self.device, dtype=torch.float32)
            data['color_bk'] = data.pop('color')
            data['color'] = torch.ones_like(data['color_bk']) * head_color_bw.reshape([1, -1, 1]) * 2.0
            data['color'][:, :, 1] = 1
            data['color'] = torch.clamp(data['color'], 0., 1.)
            data = self.head_camera.render_gaussian(data, render_size)
            render_bw = data['render_images'][:, :3, :, :]
            data['color'] = data.pop('color_bk')
            data['render_bw'] = render_bw
            
        supres_image = data['supres_images'][0].permute(1, 2, 0).detach().cpu().numpy()
        supres_image = (supres_image * 255).astype(np.uint8)[:,:,::-1]
        
        render_bw = data['render_bw'][0].permute(1, 2, 0).detach().cpu().numpy()
        render_bw = np.clip(render_bw * 255, 0, 255).astype(np.uint8)[:,:,::-1]
        render_bw = cv2.resize(render_bw, (supres_image.shape[0], supres_image.shape[1]))

            
        output = {
            'render_images':supres_image,
            'render_bw':render_bw,
        }
        
        return output
    
    def prepare_camera_data_for_gs_rendering(self, extrinsic, intrinsic, original_resolution, new_resolution):
        extrinsic = np.copy(extrinsic)
        intrinsic = np.copy(intrinsic)
        new_intrinsic = np.copy(intrinsic)
        new_intrinsic[:2] *= new_resolution / original_resolution

        intrinsic[0, 0] = intrinsic[0, 0] * 2 / original_resolution
        intrinsic[0, 2] = intrinsic[1, 2] * 2 / original_resolution - 1
        intrinsic[1, 1] = intrinsic[1, 1] * 2 / original_resolution
        intrinsic[1, 2] = intrinsic[1, 2] * 2 / original_resolution - 1
        fovx = 2 * math.atan(1 / intrinsic[0, 0])
        fovy = 2 * math.atan(1 / intrinsic[1, 1])

        world_view_transform = torch.tensor(getWorld2View2(extrinsic[:3, :3].transpose(), extrinsic[:3, 3])).transpose(0, 1)
        projection_matrix = getProjectionMatrix(
            znear=0.01, zfar=100, fovX=None, fovY=None, 
            K=new_intrinsic, img_h=new_resolution, img_w=new_resolution).transpose(0,1)
        full_proj_transform = (world_view_transform.unsqueeze(0).bmm(projection_matrix.unsqueeze(0))).squeeze(0)
        camera_center = world_view_transform.inverse()[3, :3]

        c2w = np.linalg.inv(extrinsic)
        viewdir = np.matmul(c2w[:3, :3], np.array([0, 0, -1], np.float32).reshape([3, 1])).reshape([-1])
        viewdir = torch.from_numpy(viewdir.astype(np.float32))

        return {
            'extrinsics': torch.from_numpy(extrinsic.astype(np.float32)),
            'intrinsics': torch.from_numpy(intrinsic.astype(np.float32)),
            'viewdir': viewdir, 
            'fovx': torch.Tensor([fovx]),
            'fovy': torch.Tensor([fovy]),
            'world_view_transform': world_view_transform,
            'projection_matrix': projection_matrix,
            'full_proj_transform': full_proj_transform,
            'camera_center': camera_center
            }
        
    def stich_head_body(self,body_rendering,body_mask,body_torso_mask,head_rendering,head_blending_mask,body_head_blending_params,head_offline_rendering_param):
        color_transfer = body_head_blending_params['color_transfer']
        zoom_image_size = head_offline_rendering_param['zoom_image_size']
        zoom_center = head_offline_rendering_param['zoom_center']
        zoom_scale = head_offline_rendering_param['zoom_scale']
        
        
        if len(body_mask.shape) == 3:
            body_mask = body_mask[:, :, 0]
        if len(body_torso_mask.shape) == 3:
            body_torso_mask = body_torso_mask[:, :, 0]
        
        head_rendering = cv2.resize(head_rendering, (int(zoom_image_size[0]), int(zoom_image_size[1])))
        head_blending_mask = cv2.resize(head_blending_mask, (int(zoom_image_size[0]), int(zoom_image_size[1])))
        head_mask = head_blending_mask[:, :, 1]
        head_blending_mask = head_blending_mask[:, :, 0]
        head_blending_mask = soften_blending_mask(head_blending_mask, head_mask)
        
        pasteback_center = zoom_center
        pasteback_scale = zoom_scale
        
        head_rendering_back = paste_back_with_linear_interp(pasteback_scale, pasteback_center, head_rendering, [body_rendering.shape[1], body_rendering.shape[0]])
        head_blending_mask_back = paste_back_with_linear_interp(pasteback_scale, pasteback_center, head_blending_mask, [body_rendering.shape[1], body_rendering.shape[0]])
        head_mask_back = paste_back_with_linear_interp(pasteback_scale, pasteback_center, head_mask, [body_rendering.shape[1], body_rendering.shape[0]])
        # head_blending_mask_back *= body_mask
        # head_mask_back *= body_mask
        head_blending_mask_back = head_blending_mask_back * (1 - body_torso_mask)

        head_rendering_back_shape = head_rendering_back.shape
        head_rendering_back = np.matmul(head_rendering_back.reshape(-1, 3), color_transfer[:3, :3].transpose()) + color_transfer[:3, 3][None]
        head_rendering_back = head_rendering_back.reshape(head_rendering_back_shape)
        head_rendering_back = head_rendering_back * head_mask_back[:, :, None] + (1 - head_mask_back[:, :, None])

        body_rendering = body_rendering * (1 - head_blending_mask_back[:, :, None]) + head_rendering_back * head_blending_mask_back[:, :, None]
        
        return np.uint8(np.clip(body_rendering, 0, 1)*255)
    
    # def build_hand(betas,poses,camera):
        
    #     # build hand here
        
    #     output = {
    #         'hand_render':render,
    #         'hand_mask':mask,
            
    #     }        
        
    #     return output
    

if __name__ == '__main__':
    conf = OmegaConf.load('configs/example.yaml')
    avatar = Avatar(conf)
    avatar.build_dataset()
    # avatar.test_body()
    avatar.render_all()