all

.gitattributes

@@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoints/face_0929/gaussianhead_latest filter=lfs diff=lfs merge=lfs -text
+checkpoints/face_0929/supres_latest filter=lfs diff=lfs merge=lfs -text
+checkpoints/face_0929/delta_poses_latest filter=lfs diff=lfs merge=lfs -text
+checkpoints/pos_map_ys/body_mix/smpl_pos_map/cano_smpl_nml_map.exr filter=lfs diff=lfs merge=lfs -text
+checkpoints/pos_map_ys/body_mix/smpl_pos_map/cano_smpl_pos_map.exr filter=lfs diff=lfs merge=lfs -text
+checkpoints/ref_gaussian/head/000000.ply filter=lfs diff=lfs merge=lfs -text
+checkpoints/ filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,7 +1,147 @@
 import gradio as gr
+import moviepy.editor as mpy
+import numpy as np
+import os
+from omegaconf import OmegaConf
+from tqdm import tqdm
+import shutil
+import time
+from avatar_generator import Avatar
 
-def greet(name):
-    return "Hello " + name + "!!"
+# # 指定保存文件的目录
+# SAVE_DIR = "./uploaded_files"
+# os.makedirs(SAVE_DIR, exist_ok=True)  # 创建目录（如果不存在）
+# 全局变量，用于控制任务是否应当终止
+should_stop = False
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+
+# 定义逐帧处理的函数
+def process_files(file1, file2):
+    global should_stop
+    should_stop = False  # 重置停止标志
+    
+    yield None, None, None, "Starting Process!"
+    
+    file_path1 = file1.name
+    file_path2 = file2.name 
+    pose_data = np.load(file_path1)
+    exp_data = np.load(file_path2)
+    
+    # save
+    pose_path = './test_data/AMASS/online_test_pose_data.npz'
+    exp_path = './test_data/face_exp/online_test_exp_data.npy'
+
+    np.savez(pose_path, **pose_data)
+    np.save(exp_path, exp_data)
+
+ 
+    # with open(file1.name, 'rb') as fsrc:
+    #     with open(file_path1, 'wb') as fdst:
+    #         shutil.copyfileobj(fsrc, fdst)
+    
+    # with open(file2.name, 'rb') as fsrc:
+    #     with open(file_path2, 'wb') as fdst:
+    #         shutil.copyfileobj(fsrc, fdst)
+    
+    conf = OmegaConf.load('configs/example.yaml')
+    avatar = Avatar(conf)
+    avatar.build_dataset(pose_path, exp_path)
+    
+    lenth = min(len(avatar.body_dataset), len(avatar.head_dataloader),20)
+    output_frames = []
+    
+    start_time = time.time()
+    for idx in tqdm(range(lenth)):
+        if should_stop:
+            yield None, None, None, None
+            break  # 任务应当终止时跳出循环
+        frame = avatar.reder_frame(idx)
+        # rgb2bgr
+        frame = frame[..., ::-1]
+        output_frames.append(frame)
+        
+        elapsed_time = time.time() - start_time
+        estimated_total_time = (elapsed_time / (idx + 1)) * lenth
+        remaining_time = estimated_total_time - elapsed_time
+        
+        yield frame, None, (idx + 1) / lenth * 100, f"{elapsed_time:.2f} sec/{estimated_total_time:.2f} sec"
+    
+    if not should_stop:
+        output_path = "./output/output_video.mp4"
+        final_video = mpy.ImageSequenceClip(output_frames, fps=25)
+        final_video.write_videofile(output_path, codec='libx264')
+
+        yield output_frames[-1], output_path, 100.0, "Processing completed!"
+
+# 清除操作
+def clear_files():
+    global should_stop
+    should_stop = True  # 设置停止标志
+
+    # 返回空值以清空界面元素
+    return None, None, None, None, None, None
+
+
+# 创建 Gradio 接口
+with gr.Blocks(css="""
+    .equal-height {
+        height: 425px;  /* 设置为你希望的高度 */
+        display: flex;
+        flex-direction: column;
+        justify-content: center;
+        align-items: center;
+    }
+    .equal-height input {
+        height: 100%;  /* 输入框占满整个容器高度 */
+    }
+    .output-container {
+        height: 400px;  /* 输出框的高度 */
+    }
+    .custom-text {
+        height: 80px;  /* 输出框的高度 */
+    }
+""") as demo:
+    with gr.Row():
+        # 左侧列，用于放置文件输入
+        with gr.Column(scale=1):
+            with gr.Row(elem_classes="equal-height"):
+                file_input1 = gr.File(label="Upload File (Body Pose)")
+                file_input2 = gr.File(label="Upload File (Face EXP)")
+
+        with gr.Column(scale=2):
+            with gr.Row():
+                # 中间列，用于放置帧输出
+                with gr.Column(scale=1):
+                    frame_output = gr.Image(label="Current Frame Output", elem_classes="output-container")  # 输出当前帧图像
+                # 右侧列，用于放置视频输出
+                with gr.Column(scale=1):
+                    video_output = gr.Video(label="Processed Video Output", elem_classes="output-container")  # 输出视频
+                    # progress_bar = gr.Label(label="Progress")
+            with gr.Row():
+                with gr.Column(scale=2):
+                    progress_bar = gr.Slider(visible=True, minimum=0, maximum=100, step=1, label="Progress %",elem_classes="custom-text")  # 使用Slider模拟进度条
+                with gr.Column(scale=1):
+                    output_time = gr.Textbox(label='Processing Time/Estimate Time', elem_classes="custom-text")
+                    # time_label = gr.Label(value="", label="Estimated Time Remaining", elem_classes="custom-label")
+            # with gr.Row():
+            #     progress_bar = gr.Progress()  # 添加进度条
+    with gr.Row():
+        process_button = gr.Button("Start Processing Files")
+        clear_button = gr.Button("Clear or Stop Processing")
+    
+    # 定义按钮的功能
+    process_button.click(
+        fn=process_files, 
+        inputs=[file_input1, file_input2], 
+        outputs=[frame_output, video_output, progress_bar, output_time],
+        show_progress=False
+        )
+
+    clear_button.click(
+        fn= clear_files, 
+        inputs=[], 
+        outputs=[file_input1, file_input2, frame_output, video_output, progress_bar, output_time]
+        )
+
+# 启动应用
+demo.launch()

avatar.py

@@ -0,0 +1,642 @@
+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()

calc_offline_rendering_param.py

@@ -0,0 +1,188 @@
+import numpy as np
+import tqdm
+import os, glob
+import json
+import argparse
+
+from render_utils.lib.networks.smpl_torch import SmplTorch
+from render_utils.lib.utils.gaussian_np_utils import load_gaussians_from_ply
+from render_utils.stitch_body_and_head import load_body_params, load_face_params, get_smpl_verts_and_head_transformation, calc_livehead2livebody
+
+
+def load_rendering_camera(camera_fpath):
+    with open(camera_fpath, 'r') as fp:
+        camera_data = json.load(fp)
+    camera_data = camera_data[0]
+    image_size = [camera_data['width'], camera_data['height']]
+    cam_f = [camera_data['fx'], camera_data['fy']]
+    cam_pos = np.array(camera_data['position'])
+    cam_rot = np.array(camera_data['rotation']).reshape(3, 3)
+    c2w = np.eye(4)
+    c2w[:3, :3] = cam_rot
+    c2w[:3, 3] = cam_pos
+    cam_extr = np.linalg.inv(c2w)
+    cam_intr = np.eye(3)
+    cam_intr[0, 0] = cam_f[0]
+    cam_intr[1, 1] = cam_f[1]
+    cam_intr[0, 2] = image_size[0] / 2
+    cam_intr[1, 2] = image_size[1] / 2
+    return cam_extr, cam_intr, image_size
+
+def load_camera_list(camera_fpath):
+    with open(camera_fpath, 'r') as fp:
+        camera_data = json.load(fp)
+    image_size = [camera_data[0]['width'], camera_data[0]['height']]
+    cam_list = []
+    for cam in camera_data:
+        cam_f = [cam['fx'], cam['fy']]
+        cam_pos = np.array(cam['position'])
+        cam_rot = np.array(cam['rotation']).reshape(3, 3)
+        c2w = np.eye(4)
+        c2w[:3, :3] = cam_rot
+        c2w[:3, 3] = cam_pos
+        cam_extr = np.linalg.inv(c2w)
+        cam_intr = np.eye(3)
+        cam_intr[0, 0] = cam_f[0]
+        cam_intr[1, 1] = cam_f[1]
+        cam_intr[0, 2] = image_size[0] / 2
+        cam_intr[1, 2] = image_size[1] / 2
+        cam_list.append((cam_extr, cam_intr))
+    return cam_list, image_size
+
+def load_camera_data(cam):
+    image_size = [cam['width'], cam['height']]
+    cam_f = [cam['fx'], cam['fy']]
+    cam_pos = np.array(cam['position'])
+    cam_rot = np.array(cam['rotation']).reshape(3, 3)
+    c2w = np.eye(4)
+    c2w[:3, :3] = cam_rot
+    c2w[:3, 3] = cam_pos
+    cam_extr = np.linalg.inv(c2w)
+    cam_intr = np.eye(3)
+    cam_intr[0, 0] = cam_f[0]
+    cam_intr[1, 1] = cam_f[1]
+    cam_intr[0, 2] = image_size[0] / 2
+    cam_intr[1, 2] = image_size[1] / 2
+    
+    return (cam_extr, cam_intr), image_size
+
+def calc_offline_rendering_param(
+        body_gaussian_root_dir, ref_head_gaussian_path, ref_head_param_path, render_cam_fpath, 
+        body_head_blending_param_path):
+    body_param_flist = sorted(glob.glob(os.path.join(body_gaussian_root_dir, 'posed_params/*.npz')))
+
+    head_gaussians = load_gaussians_from_ply(ref_head_gaussian_path)
+    head_pose, head_scale, id_coeff, exp_coeff = load_face_params(ref_head_param_path)
+    # cam_extr_body, cam_intr_body, image_size = load_rendering_camera(render_cam_fpath)
+    cam_list, image_size = load_camera_list(render_cam_fpath)
+
+    body_head_blending_params = np.load(body_head_blending_param_path)
+    smplx_to_faceverse = body_head_blending_params['smplx_to_faceverse']
+    residual_transf = body_head_blending_params['residual_transf']
+    body_nonface_mask = body_head_blending_params['body_nonface_mask']
+    head_nonface_mask = body_head_blending_params['head_nonface_mask']
+    head_facial_idx = body_head_blending_params['head_facial_idx']
+    body_facial_idx = body_head_blending_params['body_facial_idx']
+    head_body_corr_idx = body_head_blending_params['head_body_corr_idx']
+    head_color_bw = body_head_blending_params['head_color_bw']
+    color_transfer = body_head_blending_params['color_transfer']
+
+    smpl = SmplTorch(model_file='./AnimatableGaussians/smpl_files/smplx/SMPLX_NEUTRAL.npz')
+
+    head_cam_extr = []
+    head_cam_intr = []
+    head_cam_intr_zoom = []
+    head_zoom_center = []
+    head_zoom_scale = []
+    
+    for i, body_param_fpath in enumerate(tqdm.tqdm(body_param_flist)):
+        global_orient, transl, body_pose, betas = load_body_params(body_param_fpath)
+        # body_gaussians = load_gaussians_from_ply(body_gaussian_fpath)
+
+        smpl_verts, head_joint_transfmat = get_smpl_verts_and_head_transformation(
+            smpl, global_orient, body_pose, transl, betas)
+        livehead2livebody = calc_livehead2livebody(head_pose, smplx_to_faceverse, head_joint_transfmat)
+        total_transf = np.matmul(livehead2livebody, residual_transf)
+
+        cam_extr = np.matmul(cam_list[i][0], total_transf)
+        cam_intr = np.copy(cam_list[i][1])
+
+        head_cam_extr.append(cam_extr)
+        head_cam_intr.append(cam_intr)
+
+        pts = np.copy(head_gaussians.xyz)
+        pts_proj = np.matmul(pts, cam_extr[:3, :3].transpose()) + cam_extr[:3, 3]
+        pts_proj = np.matmul(pts_proj, cam_intr.transpose())
+        pts_proj = pts_proj / pts_proj[:, 2:]
+        # pts_proj = np.int32(np.round(pts_proj[:, :2]))
+
+        # img = np.zeros([image_size[1], image_size[0], 3], dtype=np.uint8)
+        # for p in pts_proj[::50]:
+        #     p = np.clip(p, 0, image_size[0] - 1)
+        #     cv.circle(img, (int(p[0]), int(p[1])), 2, (0, 255, 0), -1)
+        # cv.imshow('img', img)
+
+        pts_min, pts_max = np.min(pts_proj, axis=0), np.max(pts_proj, axis=0)
+        pts_center = (pts_min + pts_max) // 2
+        pts_size = np.max(pts_max - pts_min)
+        tgt_pts_size = 350
+        tgt_image_size = 512
+        zoom_scale = tgt_pts_size / pts_size
+        cam_intr_zoom = np.copy(cam_intr)
+        cam_intr_zoom[:2] *= zoom_scale
+        cam_intr_zoom[0, 2] = cam_intr_zoom[0, 2] - (pts_center[0]*zoom_scale - tgt_image_size/2)
+        cam_intr_zoom[1, 2] = cam_intr_zoom[1, 2] - (pts_center[1]*zoom_scale - tgt_image_size/2)
+        head_cam_intr_zoom.append(cam_intr_zoom)
+        head_zoom_center.append(pts_center)
+        head_zoom_scale.append(zoom_scale)
+
+        # pts_proj = np.matmul(pts, cam_extr[:3, :3].transpose()) + cam_extr[:3, 3]
+        # pts_proj = np.matmul(pts_proj, cam_intr_zoom.transpose())
+        # pts_proj = pts_proj / pts_proj[:, 2:]
+        # pts_proj = np.int32(np.round(pts_proj[:, :2]))
+        # img = np.zeros([512, 512, 3], dtype=np.uint8)
+        # for p in pts_proj[::50]:
+        #     p = np.clip(p, 0, image_size[0] - 1)
+        #     cv.circle(img, (int(p[0]), int(p[1])), 2, (0, 255, 0), -1)
+        # cv.imshow('img_zoom', img)
+        # cv.waitKey()
+
+    np.savez(os.path.join(os.path.dirname(body_head_blending_param_path), 'head_zoomin_render_param.npz'), 
+             cam_extr=head_cam_extr, cam_intr=head_cam_intr, image_size=image_size,
+             cam_intr_zoom=head_cam_intr_zoom, zoom_image_size=[tgt_image_size, tgt_image_size], 
+             zoom_center=head_zoom_center, 
+             zoom_scale=head_zoom_scale, 
+             head_pose=head_pose, head_scale=head_scale, head_color_bw=head_color_bw)
+    
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    """
+    body_gaussian_root_dir, ref_head_gaussian_path, ref_head_param_path, render_cam_fpath, 
+        body_head_blending_param_path
+    """
+    parser.add_argument('--body_gaussian_root_dir', type=str)
+    parser.add_argument('--ref_head_gaussian_path', type=str)
+    parser.add_argument('--ref_head_param_path', type=str)
+    parser.add_argument('--render_cam_fpath', type=str)
+    parser.add_argument('--body_head_blending_param_path', type=str)
+    args = parser.parse_args()
+    calc_offline_rendering_param(
+        args.body_gaussian_root_dir, 
+        args.ref_head_gaussian_path, 
+        args.ref_head_param_path, 
+        args.render_cam_fpath, 
+        args.body_head_blending_param_path
+    )
+
+"""
+python calc_offline_rendering_param.py ^
+--body_gaussian_root_dir ./AnimatableGaussians/test_results/huawei0425/checkpoints/AMASS__test_poses_ours_front_view/batch_750000/pca_20_sigma_2.00/ ^
+--ref_head_gaussian_path ./Gaussian-Head-Avatar/results/reenactment/huawei0425_self/posed_gaussians/000000.ply ^
+--ref_head_param_path ./Gaussian-Head-Avatar/results/reenactment/huawei0425_self/params/000000_param.npz ^
+--render_cam_fpath ./AnimatableGaussians/test_results/huawei0425/checkpoints/AMASS__test_poses_ours_front_view/batch_750000/pca_20_sigma_2.00/cameras.json ^
+--body_head_blending_param_path ./data/body_face_stitching_sr/body_head_blending_param.npz
+
+"""

configs/example.yaml

@@ -0,0 +1,89 @@
+trial_name: "body_head_avatar"
+device: cuda
+animatablegaussians:
+  train:
+    dataset: MvRgbDatasetAvatarReX
+    data:
+      subject_name: 1007_slow10
+      data_dir: ./checkpoints/pos_map_ys/body_mix
+      frame_range: &id001
+      - 0
+      - 200
+      - 1
+      used_cam_ids:
+      - 0
+      - 1
+      - 2
+      - 3
+      - 4
+      - 5
+      - 6
+      - 8
+      - 9
+      - 10
+      - 11
+      - 12
+      - 14
+      - 15
+      load_smpl_pos_map: true
+    pretrained_dir: null
+    net_ckpt_dir: ./results/huawei0425/avatar2
+    prev_ckpt: null
+    ckpt_interval:
+      epoch: 10
+      batch: 50000
+    eval_interval: 1000
+    eval_training_ids:
+    - 190
+    - 7
+    eval_testing_ids:
+    - 354
+    - 7
+    eval_img_factor: 1.0
+    lr_init: 0.0005
+    loss_weight:
+      l1: 1.0
+      lpips: 0.1
+      offset: 0.005
+    finetune_color: false
+    batch_size: 1
+    num_workers: 8
+    random_bg_color: true
+  test:
+    output_dir: ./test_results/temp_test
+    dataset: MvRgbDatasetAvatarReX
+    data:
+      data_dir: ./checkpoints/pos_map_ys/body_mix
+      frame_range: [0, 800]
+      subject_name: huawei0425
+    pose_data:
+      data_path: ./test_data/AMASS/1007_train_data_slow10.npz
+      frame_range: [0, 2000]
+    view_setting: degree90
+    render_view_idx: 13
+    global_orient: true
+    img_scale: 2.0
+    save_mesh: false
+    render_skeleton: false
+    save_tex_map: false
+    save_ply: true
+    fix_hand: true
+    fix_hand_id: 23
+    n_pca: 20
+    sigma_pca: 2.0
+    prev_ckpt: ./checkpoints/checkpoints/body_ys
+  model:
+    with_viewdirs: true
+    random_style: false
+
+gha:
+  config_path: configs/head.yaml
+  offline_rendering_param_fpath: ./checkpoints/render_param/head_zoomin_render_param.npz
+
+cat:
+  body_gaussian_root_dir: ./checkpoints/pos_map_ys/body_mix
+  ref_head_gaussian_path: ./checkpoints/ref_gaussian/head/000000.ply
+  ref_head_param_path: ./checkpoints/ref_gaussian/head/000000_param.npz
+  render_cam_fpath: /home/pengc02/pengcheng/projects/gaussian_avatar/avatar_final/AnimatableGaussians/test_results/1007_slow10/checkpoints/AMASS__1007_train_data_slow10_degree90_view/batch_789377/pca_20_sigma_2.00/cameras.json
+  body_head_blending_param_path: ./checkpoints/render_param/body_head_blending_param.npz
+

configs/head.yaml

@@ -0,0 +1,39 @@
+gpu_id: 0
+load_supres_checkpoint: './checkpoints/face_0929/supres_latest'
+load_gaussianhead_checkpoint: './checkpoints/face_0929/gaussianhead_latest'
+
+dataset:
+  dataroot: './test_data/face1001'
+  image_files: 'images/*/wrong_image.jpg'
+  param_files: 'params/*/params.npz'
+  camera_path: './test_data/face1001/cameras/0000/camera_22070938.npz'
+  pose_code_path: './test_data/face1001/params/0000/params.npz'
+  exp_path: '/home/pengc02/pengcheng/projects/gaussian_avatar/avatar_final/data/1005_thu_slow/thuSlow10.npy'
+  freeview: False
+  resolution: 2048
+  original_resolution: 2048
+
+supresmodule:
+  input_dim: 32
+  output_dim: 3
+  network_capacity: 32
+
+gaussianheadmodule:
+  num_add_mouth_points: 3000
+  exp_color_mlp: [180, 256, 256, 32]
+  pose_color_mlp: [182, 128, 32]
+  exp_deform_mlp: [79, 256, 256, 256, 256, 256, 3]
+  pose_deform_mlp: [81, 256, 256, 3]
+  exp_attributes_mlp: [180, 256, 256, 256, 8]
+  pose_attributes_mlp: [182, 128, 128, 8]
+  exp_coeffs_dim: 52
+  pos_freq: 4
+  dist_threshold_near: 0.05
+  dist_threshold_far: 0.12
+  deform_scale: 0.3
+  attributes_scale: 0.2
+
+recorder:
+  name: 'thu_exp_slow'
+  result_path: 'results/reenactment'
+

gradio_debug.py

@@ -0,0 +1,21 @@
+import gradio as gr
+import os
+
+def load_and_display_video(video_path):
+    if os.path.exists(video_path):
+        return video_path
+    else:
+        return "Invalid video path."
+
+with gr.Blocks() as demo:
+    video_input = gr.Textbox(label="Enter Video Path")
+    video_output = gr.Video(label="Video Output")
+
+    load_button = gr.Button("Load Video")
+    
+    load_button.click(fn=load_and_display_video,
+                      inputs=video_input,
+                      outputs=video_output)
+
+# 启动应用
+demo.launch()

other_requirement.sh

@@ -0,0 +1,19 @@
+# pip install kaolin==0.16.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.4.0_cu121.html
+
+
+cd AnimatableGaussians
+# install diff-gaussian-rasterization-depth-alpha
+cd gaussians/diff_gaussian_rasterization_depth_alpha
+python setup.py install
+cd ../..
+
+# install styleunet
+cd network/styleunet
+python setup.py install
+cd ../..
+
+# HTTPS
+git clone https://github.com/graphdeco-inria/gaussian-splatting --recursive
+# Modify "submodules/diff-gaussian-rasterization/cuda_rasterizer/config.h" from "NUM_CHANNELS 3" to "NUM_CHANNELS 32"
+pip install submodules/diff-gaussian-rasterization
+pip install submodules/simple-knn