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()