avatar.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

# 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.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

# cat utils
from calc_offline_rendering_param import calc_offline_rendering_param

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
        self.head_config = config_reenactment()
        self.head_config.load(self.head.config_path)
        self.head_config = self.head_config.get_cfg()
        
        # cat utils part init
        self.cat = config.cat
    
    @torch.no_grad()
    def test_body(self):
        # run the animatable gaussian test
        self.avatar_net.eval()
        dataset_module = self.body.get('dataset', 'MvRgbDatasetAvatarReX')
        MvRgbDataset = importlib.import_module('AnimatableGaussians.dataset.dataset_mv_rgb').__getattribute__(dataset_module)
        training_dataset = MvRgbDataset(**self.body['train']['data'], training = False)
        if self.body['test'].get('n_pca', -1) >= 1:
            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 = 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.dataset = testing_dataset
        # iter_idx = self.load_ckpt(self.body['test']['prev_ckpt'], False)[1]

        output_dir = self.body['test'].get('output_dir', None)
        if output_dir is None:
            raise ValueError('No output_dir in test config')
        use_pca = self.body['test'].get('n_pca', -1) >= 1
        if use_pca:
            output_dir += '/pca_%d_sigma_%.2f' % (self.body['test'].get('n_pca', -1), float(self.body['test'].get('sigma_pca', 1.)))
        else:
            output_dir += '/vanilla'
        print('# Output dir: \033[1;31m%s\033[0m' % output_dir)

        os.makedirs(output_dir + '/live_skeleton', exist_ok = True)
        os.makedirs(output_dir + '/rgb_map', exist_ok = True)
        os.makedirs(output_dir + '/rgb_map_wo_hand', exist_ok = True)
        os.makedirs(output_dir + '/torso_map', exist_ok = True)
        os.makedirs(output_dir + '/mask_map', exist_ok = True)
        os.makedirs(output_dir + '/posed_gaussians', exist_ok = True)
        os.makedirs(output_dir + '/posed_params', exist_ok = True)
        os.makedirs(output_dir + '/full_body_mask', exist_ok = True)
        os.makedirs(output_dir + '/hand_only_mask', exist_ok = True)

        geo_renderer = None
        item_0 = self.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']
        
        # 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)

        data_num = len(self.dataset)
        if self.body['test'].get('fix_hand', False):
            self.avatar_net.generate_mean_hands()
        log_time = False
        extr_list = []
        intr_list = []
        img_h_list = []
        img_w_list = []
        

        for idx in tqdm(range(data_num), desc = 'Rendering avatars...'):
            if log_time:
                time_start.record()
                time_start_all.record()

            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!')
            
            
            self.dump_renderer_info(output_dir, extr_list, intr_list, img_h_list, img_w_list)
            # also save the extr and intr and img_h and img_w to json
            camera_info = []
            for i in range(len(extr_list)):
                camera = {}
                camera['extr'] = extr_list[i].tolist()
                camera['intr'] = intr_list[i].tolist()
                camera['img_h'] = img_h_list[i]
                camera['img_w'] = img_w_list[i]
                camera_info.append(camera)
            with open(os.path.join(output_dir, 'camera_info.json'), 'w') as fp:
                json.dump(camera_info, fp)
            

            getitem_func = self.dataset.getitem_fast if hasattr(self.dataset, 'getitem_fast') else self.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 view_setting.startswith('moving') or view_setting == 'free_moving':
                current_center = items['live_bounds'].cpu().numpy().mean(0)
                delta = current_center - object_center

                object_center[0] += delta[0]
                # object_center[1] += delta[1]
                # object_center[2] += delta[2]

            if log_time:
                time_end.record()
                torch.cuda.synchronize()
                print('Loading data costs %.4f secs' % (time_start.elapsed_time(time_end) / 1000.))
                time_start.record()

            if self.body['test'].get('render_skeleton', False):
                from AnimatableGaussians.utils.visualize_skeletons import construct_skeletons
                skel_vertices, skel_faces = construct_skeletons(item['joints'].cpu().numpy(), item['kin_parent'].cpu().numpy())
                skel_mesh = trimesh.Trimesh(skel_vertices, skel_faces, process = False)

                if geo_renderer is None:
                    geo_renderer = Renderer(item['img_w'], item['img_h'], shader_name = 'phong_geometry', bg_color = (1, 1, 1))
                extr, intr = item['extr'], item['intr']
                geo_renderer.set_camera(extr, intr)
                geo_renderer.set_model(skel_vertices[skel_faces.reshape(-1)], skel_mesh.vertex_normals.astype(np.float32)[skel_faces.reshape(-1)])
                skel_img = geo_renderer.render()[:, :, :3]
                skel_img = (skel_img * 255).astype(np.uint8)
                cv.imwrite(output_dir + '/live_skeleton/%08d.jpg' % item['data_idx'], skel_img)

            if log_time:
                time_end.record()
                torch.cuda.synchronize()
                print('Rendering skeletons costs %.4f secs' % (time_start.elapsed_time(time_end) / 1000.))
                time_start.record()

            if 'smpl_pos_map' not in items:
                self.avatar_net.get_pose_map(items)

            # pca
            if use_pca:
                mask = 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 = 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)
                })

            if log_time:
                time_end.record()
                torch.cuda.synchronize()
                print('Rendering pose conditions costs %.4f secs' % (time_start.elapsed_time(time_end) / 1000.))
                time_start.record()

            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)
            
            if log_time:
                time_end.record()
                torch.cuda.synchronize()
                print('Rendering avatar costs %.4f secs' % (time_start.elapsed_time(time_end) / 1000.))
                time_start.record()
            
            if 'rgb_map' in output_wo_hand:
                rgb_map_wo_hand = output_wo_hand['rgb_map']

            if 'full_body_rgb_map' in mask_output:
                os.makedirs(output_dir + '/full_body_mask', exist_ok = True)
                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)
                cv.imwrite(output_dir + '/full_body_mask/%08d.png' % item['data_idx'], full_body_mask.cpu().numpy())
            
            if 'hand_only_rgb_map' in mask_output:
                os.makedirs(output_dir + '/hand_only_mask', exist_ok = True)
                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)
                cv.imwrite(output_dir + '/hand_only_mask/%08d.png' % item['data_idx'], hand_only_mask.cpu().numpy())

            if 'full_body_rgb_map' in mask_output and 'hand_only_rgb_map' in mask_output:
                # mask only covers hand
                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.01

                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
                
                # now save 3 mask to 3 folders
                os.makedirs(output_dir + '/hand_mask', exist_ok = True)
                os.makedirs(output_dir + '/r_hand_mask', exist_ok = True)
                os.makedirs(output_dir + '/l_hand_mask', exist_ok = True)
                os.makedirs(output_dir + '/hand_visual', exist_ok = True)
                
                all_mask = (all_mask * 255).to(torch.uint8)   
                cv.imwrite(output_dir + '/hand_mask/%08d.png' % item['data_idx'], all_mask.cpu().numpy())
                r_hand_mask = (body_red_mask * 255).to(torch.uint8)
                cv.imwrite(output_dir + '/r_hand_mask/%08d.png' % item['data_idx'], r_hand_mask.cpu().numpy())
                l_hand_mask = (body_blue_mask * 255).to(torch.uint8)
                cv.imwrite(output_dir + '/l_hand_mask/%08d.png' % item['data_idx'], l_hand_mask.cpu().numpy())
                hand_visual = [if_mask_r_hand, if_mask_l_hand]
                # save to npy
                with open(output_dir + '/hand_visual/%08d.npy' % item['data_idx'], 'wb') as f:
                    np.save(f, hand_visual)
                
                    
            # now build sleeve_mask
            if 'left_hand_rgb_map' in mask_output and 'right_hand_rgb_map' in mask_output:
                os.makedirs(output_dir + '/left_sleeve_mask', exist_ok = True)
                os.makedirs(output_dir + '/right_sleeve_mask', exist_ok = True)
                
                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
                
                # save
                left_hand_mask = (left_hand_mask * 255).to(torch.uint8)
                cv.imwrite(output_dir + '/left_sleeve_mask/%08d.png' % item['data_idx'], left_hand_mask.cpu().numpy())
                right_hand_mask = (right_hand_mask * 255).to(torch.uint8)
                cv.imwrite(output_dir + '/right_sleeve_mask/%08d.png' % item['data_idx'], right_hand_mask.cpu().numpy())
                    
            rgb_map = output['rgb_map']
            rgb_map.clip_(0., 1.)
            rgb_map = (rgb_map * 255).to(torch.uint8).cpu().numpy()
            cv.imwrite(output_dir + '/rgb_map/%08d.jpg' % item['data_idx'], rgb_map)
            
            # 利用 r_hand_mask 和 l_hand_mask，将wo_hand图像中的mask部分覆盖rgb_map
            if 'rgb_map' in output_wo_hand and 'full_body_rgb_map' in mask_output and 'hand_only_rgb_map' in mask_output:
                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)
                # print('mask shape: ', mask.shape)
                import ipdb
                # ipdb.set_trace()
                mix = rgb_map_wo_hand.copy() * mask + rgb_map * ~mask
                cv.imwrite(output_dir + '/rgb_map_wo_hand/%08d.png' % item['data_idx'], mix)
                
            if 'torso_map' in output:
                os.makedirs(output_dir + '/torso_map', exist_ok = True)
                torso_map = output['torso_map'][:, :, 0]
                torso_map.clip_(0., 1.)
                torso_map = (torso_map * 255).to(torch.uint8)
                cv.imwrite(output_dir + '/torso_map/%08d.png' % item['data_idx'], torso_map.cpu().numpy())

            if 'mask_map' in output:
                os.makedirs(output_dir + '/mask_map', exist_ok = True)
                mask_map = output['mask_map'][:, :, 0]
                mask_map.clip_(0., 1.)
                mask_map = (mask_map * 255).to(torch.uint8)
                cv.imwrite(output_dir + '/mask_map/%08d.png' % item['data_idx'], mask_map.cpu().numpy())

            if self.body['test'].get('save_tex_map', False):
                os.makedirs(output_dir + '/cano_tex_map', exist_ok = True)
                cano_tex_map = output['cano_tex_map']
                cano_tex_map.clip_(0., 1.)
                cano_tex_map = (cano_tex_map * 255).to(torch.uint8)
                cv.imwrite(output_dir + '/cano_tex_map/%08d.png' % item['data_idx'], cano_tex_map.cpu().numpy())

            if self.body['test'].get('save_ply', False):
                if item['data_idx'] == 0:
                    save_gaussians_as_ply(output_dir + '/posed_gaussians/%08d.ply' % item['data_idx'], output['posed_gaussians'])
                    for k in output['posed_gaussians'].keys():
                        if isinstance(output['posed_gaussians'][k], torch.Tensor):
                            output['posed_gaussians'][k] = output['posed_gaussians'][k].detach().cpu().numpy()
                    np.savez(output_dir + '/posed_gaussians/%08d.npz' % item['data_idx'], **output['posed_gaussians'])
                np.savez(output_dir + ('/posed_params/%08d.npz' % item['data_idx']), 
                         betas=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())

            if log_time:
                time_end.record()
                torch.cuda.synchronize()
                print('Saving images costs %.4f secs' % (time_start.elapsed_time(time_end) / 1000.))
                print('Animating one frame costs %.4f secs' % (time_start_all.elapsed_time(time_end) / 1000.))

            torch.cuda.empty_cache()
    
    def dump_renderer_info(self, dump_dir, extrs, intrs, img_heights, img_widths):
        with open(os.path.join(dump_dir, 'cfg_args'), 'w') as fp:
            outstr = "Namespace(sh_degree=%d, source_path='%s', model_path='%s', images='images', resolution=-1, " \
                     "white_background=False, data_device='cuda', eval=False)" % (
                      3, self.body['train']['data']['data_dir'], dump_dir)
            fp.write(outstr)
        with open(os.path.join(dump_dir, 'cameras.json'), 'w') as fp:
            cam_jsons = []
            for ci in range(len(extrs)):
                extr, intr = extrs[ci], intrs[ci]
                img_h, img_w = img_heights[ci], img_widths[ci]

                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 = {
                    'id': ci,
                    'img_name': '%08d' % ci,
                    '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]),
                }
                cam_jsons.append(camera_entry)
            json.dump(cam_jsons, fp)
        return
    
    def test_head(self):
        dataset = ReenactmentDataset(self.head_config.dataset)
        dataloader = DataLoaderX(dataset, batch_size=1, shuffle=False, pin_memory=True) 

        device = torch.device('cuda:%d' % self.head_config.gpu_id)

        gaussianhead_state_dict = torch.load(self.head_config.load_gaussianhead_checkpoint, map_location=lambda storage, loc: storage)
        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(device)
        gaussianhead.load_state_dict(gaussianhead_state_dict)

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

        camera = CameraModule()
        recorder = ReenactmentRecorder(self.head_config.recorder)

        app = Reenactment(dataloader, gaussianhead, supres, camera, recorder, self.head_config.gpu_id, dataset.freeview)
        if self.head.offline_rendering_param_fpath is None:
            app.run(stop_fid=800)
        else:
            app.run_for_offline_stitching(self.head.offline_rendering_param_fpath)
            
    def cal_cat_param(self):
        calc_offline_rendering_param(
        self.cat.body_gaussian_root_dir, 
        self.cat.ref_head_gaussian_path, 
        self.cat.ref_head_param_path, 
        self.cat.render_cam_fpath, 
        self.cat.body_head_blending_param_path
    )
        
        


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