# Copyright (c) 2023-2024, Zexin He # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from collections import defaultdict import os import glob from typing import Union import random import numpy as np import torch # from megfile import smart_path_join, smart_open import json from PIL import Image import cv2 from lam.datasets.base import BaseDataset from lam.datasets.cam_utils import build_camera_standard, build_camera_principle, camera_normalization_objaverse from lam.utils.proxy import no_proxy from typing import Optional, Union __all__ = ['VideoHeadDataset'] class VideoHeadDataset(BaseDataset): def __init__(self, root_dirs: str, meta_path: Optional[Union[str, list]], sample_side_views: int, render_image_res_low: int, render_image_res_high: int, render_region_size: int, source_image_res: int, repeat_num=1, crop_range_ratio_hw=[1.0, 1.0], aspect_standard=1.0, # h/w enlarge_ratio=[0.8, 1.2], debug=False, is_val=False, **kwargs): super().__init__(root_dirs, meta_path) self.sample_side_views = sample_side_views self.render_image_res_low = render_image_res_low self.render_image_res_high = render_image_res_high if not (isinstance(render_region_size, list) or isinstance(render_region_size, tuple)): render_region_size = render_region_size, render_region_size # [H, W] self.render_region_size = render_region_size self.source_image_res = source_image_res self.uids = self.uids * repeat_num self.crop_range_ratio_hw = crop_range_ratio_hw self.debug = debug self.aspect_standard = aspect_standard assert self.render_image_res_low == self.render_image_res_high self.render_image_res = self.render_image_res_low self.enlarge_ratio = enlarge_ratio print(f"VideoHeadDataset, data_len:{len(self.uids)}, repeat_num:{repeat_num}, debug:{debug}, is_val:{is_val}") self.multiply = kwargs.get("multiply", 14) # set data deterministic self.is_val = is_val @staticmethod def _load_pose(frame_info, transpose_R=False): c2w = torch.eye(4) c2w = np.array(frame_info["transform_matrix"]) c2w[:3, 1:3] *= -1 c2w = torch.FloatTensor(c2w) """ if transpose_R: w2c = torch.inverse(c2w) w2c[:3, :3] = w2c[:3, :3].transpose(1, 0).contiguous() c2w = torch.inverse(w2c) """ intrinsic = torch.eye(4) intrinsic[0, 0] = frame_info["fl_x"] intrinsic[1, 1] = frame_info["fl_y"] intrinsic[0, 2] = frame_info["cx"] intrinsic[1, 2] = frame_info["cy"] intrinsic = intrinsic.float() return c2w, intrinsic def img_center_padding(self, img_np, pad_ratio): ori_w, ori_h = img_np.shape[:2] w = round((1 + pad_ratio) * ori_w) h = round((1 + pad_ratio) * ori_h) if len(img_np.shape) > 2: img_pad_np = np.zeros((w, h, img_np.shape[2]), dtype=np.uint8) else: img_pad_np = np.zeros((w, h), dtype=np.uint8) offset_h, offset_w = (w - img_np.shape[0]) // 2, (h - img_np.shape[1]) // 2 img_pad_np[offset_h: offset_h + img_np.shape[0]:, offset_w: offset_w + img_np.shape[1]] = img_np return img_pad_np def resize_image_keepaspect_np(self, img, max_tgt_size): """ similar to ImageOps.contain(img_pil, (img_size, img_size)) # keep the same aspect ratio """ h, w = img.shape[:2] ratio = max_tgt_size / max(h, w) new_h, new_w = round(h * ratio), round(w * ratio) return cv2.resize(img, dsize=(new_w, new_h), interpolation=cv2.INTER_AREA) def center_crop_according_to_mask(self, img, mask, aspect_standard, enlarge_ratio): """ img: [H, W, 3] mask: [H, W] """ ys, xs = np.where(mask > 0) if len(xs) == 0 or len(ys) == 0: raise Exception("empty mask") x_min = np.min(xs) x_max = np.max(xs) y_min = np.min(ys) y_max = np.max(ys) center_x, center_y = img.shape[1]//2, img.shape[0]//2 half_w = max(abs(center_x - x_min), abs(center_x - x_max)) half_h = max(abs(center_y - y_min), abs(center_y - y_max)) aspect = half_h / half_w if aspect >= aspect_standard: half_w = round(half_h / aspect_standard) else: half_h = round(half_w * aspect_standard) if abs(enlarge_ratio[0] - 1) > 0.01 or abs(enlarge_ratio[1] - 1) > 0.01: enlarge_ratio_min, enlarge_ratio_max = enlarge_ratio enlarge_ratio_max_real = min(center_y / half_h, center_x / half_w) enlarge_ratio_max = min(enlarge_ratio_max_real, enlarge_ratio_max) enlarge_ratio_min = min(enlarge_ratio_max_real, enlarge_ratio_min) enlarge_ratio_cur = np.random.rand() * (enlarge_ratio_max - enlarge_ratio_min) + enlarge_ratio_min half_h, half_w = round(enlarge_ratio_cur * half_h), round(enlarge_ratio_cur * half_w) assert half_h <= center_y assert half_w <= center_x assert abs(half_h / half_w - aspect_standard) < 0.03 offset_x = center_x - half_w offset_y = center_y - half_h new_img = img[offset_y: offset_y + 2*half_h, offset_x: offset_x + 2*half_w] new_mask = mask[offset_y: offset_y + 2*half_h, offset_x: offset_x + 2*half_w] return new_img, new_mask, offset_x, offset_y def load_rgb_image_with_aug_bg(self, rgb_path, mask_path, bg_color, pad_ratio, max_tgt_size, aspect_standard, enlarge_ratio, render_tgt_size, multiply, intr): rgb = np.array(Image.open(rgb_path)) interpolation = cv2.INTER_AREA if rgb.shape[0] != 1024 and rgb.shape[0] == rgb.shape[1]: rgb = cv2.resize(rgb, (1024, 1024), interpolation=interpolation) if pad_ratio > 0: rgb = self.img_center_padding(rgb, pad_ratio) rgb = rgb / 255.0 if mask_path is not None: if os.path.exists(mask_path): mask = np.array(Image.open(mask_path)) > 180 if len(mask.shape) == 3: mask = mask[..., 0] assert pad_ratio == 0 # if pad_ratio > 0: # mask = self.img_center_padding(mask, pad_ratio) # mask = mask / 255.0 else: # print("no mask file") mask = (rgb >= 0.99).sum(axis=2) == 3 mask = np.logical_not(mask) # erode mask = (mask * 255).astype(np.uint8) kernel_size, iterations = 3, 7 kernel = np.ones((kernel_size, kernel_size), np.uint8) mask = cv2.erode(mask, kernel, iterations=iterations) / 255.0 else: # rgb: [H, W, 4] assert rgb.shape[2] == 4 mask = rgb[:, :, 3] # [H, W] if len(mask.shape) > 2: mask = mask[:, :, 0] mask = (mask > 0.5).astype(np.float32) rgb = rgb[:, :, :3] * mask[:, :, None] + bg_color * (1 - mask[:, :, None]) # crop image to enlarge face area. try: rgb, mask, offset_x, offset_y = self.center_crop_according_to_mask(rgb, mask, aspect_standard, enlarge_ratio) except Exception as ex: print(rgb_path, mask_path, ex) intr[0, 2] -= offset_x intr[1, 2] -= offset_y # resize to render_tgt_size for training tgt_hw_size, ratio_y, ratio_x = self.calc_new_tgt_size_by_aspect(cur_hw=rgb.shape[:2], aspect_standard=aspect_standard, tgt_size=render_tgt_size, multiply=multiply) rgb = cv2.resize(rgb, dsize=(tgt_hw_size[1], tgt_hw_size[0]), interpolation=interpolation) mask = cv2.resize(mask, dsize=(tgt_hw_size[1], tgt_hw_size[0]), interpolation=interpolation) intr = self.scale_intrs(intr, ratio_x=ratio_x, ratio_y=ratio_y) assert abs(intr[0, 2] * 2 - rgb.shape[1]) < 2.5, f"{intr[0, 2] * 2}, {rgb.shape[1]}" assert abs(intr[1, 2] * 2 - rgb.shape[0]) < 2.5, f"{intr[1, 2] * 2}, {rgb.shape[0]}" intr[0, 2] = rgb.shape[1] // 2 intr[1, 2] = rgb.shape[0] // 2 rgb = torch.from_numpy(rgb).float().permute(2, 0, 1).unsqueeze(0) mask = torch.from_numpy(mask[:, :, None]).float().permute(2, 0, 1).unsqueeze(0) return rgb, mask, intr def scale_intrs(self, intrs, ratio_x, ratio_y): if len(intrs.shape) >= 3: intrs[:, 0] = intrs[:, 0] * ratio_x intrs[:, 1] = intrs[:, 1] * ratio_y else: intrs[0] = intrs[0] * ratio_x intrs[1] = intrs[1] * ratio_y return intrs def uniform_sample_in_chunk(self, sample_num, sample_data): chunks = np.array_split(sample_data, sample_num) select_list = [] for chunk in chunks: select_list.append(np.random.choice(chunk)) return select_list def uniform_sample_in_chunk_det(self, sample_num, sample_data): chunks = np.array_split(sample_data, sample_num) select_list = [] for chunk in chunks: select_list.append(chunk[len(chunk)//2]) return select_list def calc_new_tgt_size(self, cur_hw, tgt_size, multiply): ratio = tgt_size / min(cur_hw) tgt_size = int(ratio * cur_hw[0]), int(ratio * cur_hw[1]) tgt_size = int(tgt_size[0] / multiply) * multiply, int(tgt_size[1] / multiply) * multiply ratio_y, ratio_x = tgt_size[0] / cur_hw[0], tgt_size[1] / cur_hw[1] return tgt_size, ratio_y, ratio_x def calc_new_tgt_size_by_aspect(self, cur_hw, aspect_standard, tgt_size, multiply): assert abs(cur_hw[0] / cur_hw[1] - aspect_standard) < 0.03 tgt_size = tgt_size * aspect_standard, tgt_size tgt_size = int(tgt_size[0] / multiply) * multiply, int(tgt_size[1] / multiply) * multiply ratio_y, ratio_x = tgt_size[0] / cur_hw[0], tgt_size[1] / cur_hw[1] return tgt_size, ratio_y, ratio_x def load_flame_params(self, flame_file_path, teeth_bs=None): flame_param = dict(np.load(flame_file_path), allow_pickle=True) flame_param_tensor = {} flame_param_tensor['expr'] = torch.FloatTensor(flame_param['expr'])[0] flame_param_tensor['rotation'] = torch.FloatTensor(flame_param['rotation'])[0] flame_param_tensor['neck_pose'] = torch.FloatTensor(flame_param['neck_pose'])[0] flame_param_tensor['jaw_pose'] = torch.FloatTensor(flame_param['jaw_pose'])[0] flame_param_tensor['eyes_pose'] = torch.FloatTensor(flame_param['eyes_pose'])[0] flame_param_tensor['translation'] = torch.FloatTensor(flame_param['translation'])[0] if teeth_bs is not None: flame_param_tensor['teeth_bs'] = torch.FloatTensor(teeth_bs) # flame_param_tensor['expr'] = torch.cat([flame_param_tensor['expr'], flame_param_tensor['teeth_bs']], dim=0) return flame_param_tensor @no_proxy def inner_get_item(self, idx): """ Loaded contents: rgbs: [M, 3, H, W] poses: [M, 3, 4], [R|t] intrinsics: [3, 2], [[fx, fy], [cx, cy], [weight, height]] """ crop_ratio_h, crop_ratio_w = self.crop_range_ratio_hw uid = self.uids[idx] if len(uid.split('/')) == 1: uid = os.path.join(self.root_dirs, uid) mode_str = "train" if not self.is_val else "test" transforms_json = os.path.join(uid, f"transforms_{mode_str}.json") with open(transforms_json) as fp: data = json.load(fp) cor_flame_path = transforms_json.replace('transforms_{}.json'.format(mode_str),'canonical_flame_param.npz') flame_param = np.load(cor_flame_path) shape_param = torch.FloatTensor(flame_param['shape']) # data['static_offset'] = flame_param['static_offset'] all_frames = data["frames"] sample_total_views = self.sample_side_views + 1 if len(all_frames) >= self.sample_side_views: if not self.is_val: if np.random.rand() < 0.7 and len(all_frames) > sample_total_views: frame_id_list = self.uniform_sample_in_chunk(sample_total_views, np.arange(len(all_frames))) else: replace = len(all_frames) < sample_total_views frame_id_list = np.random.choice(len(all_frames), size=sample_total_views, replace=replace) else: if len(all_frames) > sample_total_views: frame_id_list = self.uniform_sample_in_chunk_det(sample_total_views, np.arange(len(all_frames))) else: frame_id_list = np.random.choice(len(all_frames), size=sample_total_views, replace=True) else: if not self.is_val: replace = len(all_frames) < sample_total_views frame_id_list = np.random.choice(len(all_frames), size=sample_total_views, replace=replace) else: if len(all_frames) > 1: frame_id_list = np.linspace(0, len(all_frames) - 1, num=sample_total_views, endpoint=True) frame_id_list = [round(e) for e in frame_id_list] else: frame_id_list = [0 for i in range(sample_total_views)] cam_id_list = frame_id_list assert self.sample_side_views + 1 == len(frame_id_list) # source images c2ws, intrs, rgbs, bg_colors, masks = [], [], [], [], [] flame_params = [] teeth_bs_pth = os.path.join(uid, "tracked_teeth_bs.npz") use_teeth = False if os.path.exists(teeth_bs_pth) and use_teeth: teeth_bs_lst = np.load(teeth_bs_pth)['expr_teeth'] else: teeth_bs_lst = None for cam_id, frame_id in zip(cam_id_list, frame_id_list): frame_info = all_frames[frame_id] frame_path = os.path.join(uid, frame_info["file_path"]) if 'nersemble' in frame_path or "tiktok_v34" in frame_path: mask_path = os.path.join(uid, frame_info["fg_mask_path"]) else: mask_path = os.path.join(uid, frame_info["fg_mask_path"]).replace("/export/", "/mask/").replace("/fg_masks/", "/mask/").replace(".png", ".jpg") if not os.path.exists(mask_path): mask_path = os.path.join(uid, frame_info["fg_mask_path"]) teeth_bs = teeth_bs_lst[frame_id] if teeth_bs_lst is not None else None flame_path = os.path.join(uid, frame_info["flame_param_path"]) flame_param = self.load_flame_params(flame_path, teeth_bs) # if cam_id == 0: # shape_param = flame_param["betas"] c2w, ori_intrinsic = self._load_pose(frame_info, transpose_R="nersemble" in frame_path) bg_color = random.choice([0.0, 0.5, 1.0]) # 1.0 # if self.is_val: # bg_color = 1.0 rgb, mask, intrinsic = self.load_rgb_image_with_aug_bg(frame_path, mask_path=mask_path, bg_color=bg_color, pad_ratio=0, max_tgt_size=None, aspect_standard=self.aspect_standard, enlarge_ratio=self.enlarge_ratio if (not self.is_val) or ("nersemble" in frame_path) else [1.0, 1.0], render_tgt_size=self.render_image_res, multiply=16, intr=ori_intrinsic.clone()) c2ws.append(c2w) rgbs.append(rgb) bg_colors.append(bg_color) intrs.append(intrinsic) flame_params.append(flame_param) masks.append(mask) c2ws = torch.stack(c2ws, dim=0) # [N, 4, 4] intrs = torch.stack(intrs, dim=0) # [N, 4, 4] rgbs = torch.cat(rgbs, dim=0) # [N, 3, H, W] bg_colors = torch.tensor(bg_colors, dtype=torch.float32).unsqueeze(-1).repeat(1, 3) # [N, 3] masks = torch.cat(masks, dim=0) # [N, 1, H, W] flame_params_tmp = defaultdict(list) for flame in flame_params: for k, v in flame.items(): flame_params_tmp[k].append(v) for k, v in flame_params_tmp.items(): flame_params_tmp[k] = torch.stack(v) flame_params = flame_params_tmp # TODO check different betas for same person flame_params["betas"] = shape_param # reference images prob_refidx = np.ones(self.sample_side_views + 1) if not self.is_val: prob_refidx[0] = 0.5 # front_prob else: prob_refidx[0] = 1.0 # print(frame_id_list, kinect_color_list, prob_refidx[0]) prob_refidx[1:] = (1 - prob_refidx[0]) / len(prob_refidx[1:]) ref_idx = np.random.choice(self.sample_side_views + 1, p=prob_refidx) cam_id_source_list = cam_id_list[ref_idx: ref_idx + 1] frame_id_source_list = frame_id_list[ref_idx: ref_idx + 1] source_c2ws, source_intrs, source_rgbs, source_flame_params = [], [], [], [] for cam_id, frame_id in zip(cam_id_source_list, frame_id_source_list): frame_info = all_frames[frame_id] frame_path = os.path.join(uid, frame_info["file_path"]) if 'nersemble' in frame_path: mask_path = os.path.join(uid, frame_info["fg_mask_path"]) else: mask_path = os.path.join(uid, frame_info["fg_mask_path"]).replace("/export/", "/mask/").replace("/fg_masks/", "/mask/").replace(".png", ".jpg") flame_path = os.path.join(uid, frame_info["flame_param_path"]) teeth_bs = teeth_bs_lst[frame_id] if teeth_bs_lst is not None else None flame_param = self.load_flame_params(flame_path, teeth_bs) c2w, ori_intrinsic = self._load_pose(frame_info) # bg_color = 1.0 # bg_color = 0.0 bg_color = random.choice([0.0, 0.5, 1.0]) # 1. rgb, mask, intrinsic = self.load_rgb_image_with_aug_bg(frame_path, mask_path=mask_path, bg_color=bg_color, pad_ratio=0, max_tgt_size=None, aspect_standard=self.aspect_standard, enlarge_ratio=self.enlarge_ratio if (not self.is_val) or ("nersemble" in frame_path) else [1.0, 1.0], render_tgt_size=self.source_image_res, multiply=self.multiply, intr=ori_intrinsic.clone()) source_c2ws.append(c2w) source_intrs.append(intrinsic) source_rgbs.append(rgb) source_flame_params.append(flame_param) source_c2ws = torch.stack(source_c2ws, dim=0) source_intrs = torch.stack(source_intrs, dim=0) source_rgbs = torch.cat(source_rgbs, dim=0) flame_params_tmp = defaultdict(list) for flame in source_flame_params: for k, v in flame.items(): flame_params_tmp['source_'+k].append(v) for k, v in flame_params_tmp.items(): flame_params_tmp[k] = torch.stack(v) source_flame_params = flame_params_tmp # TODO check different betas for same person source_flame_params["source_betas"] = shape_param render_image = rgbs render_mask = masks tgt_size = render_image.shape[2:4] # [H, W] assert abs(intrs[0, 0, 2] * 2 - render_image.shape[3]) <= 1.1, f"{intrs[0, 0, 2] * 2}, {render_image.shape}" assert abs(intrs[0, 1, 2] * 2 - render_image.shape[2]) <= 1.1, f"{intrs[0, 1, 2] * 2}, {render_image.shape}" ret = { 'uid': uid, 'source_c2ws': source_c2ws, # [N1, 4, 4] 'source_intrs': source_intrs, # [N1, 4, 4] 'source_rgbs': source_rgbs.clamp(0, 1), # [N1, 3, H, W] 'render_image': render_image.clamp(0, 1), # [N, 3, H, W] 'render_mask': render_mask.clamp(0, 1), #[ N, 1, H, W] 'c2ws': c2ws, # [N, 4, 4] 'intrs': intrs, # [N, 4, 4] 'render_full_resolutions': torch.tensor([tgt_size], dtype=torch.float32).repeat(self.sample_side_views + 1, 1), # [N, 2] 'render_bg_colors': bg_colors, # [N, 3] 'pytorch3d_transpose_R': torch.Tensor(["nersemble" in frame_path]), # [1] } #['root_pose', 'body_pose', 'jaw_pose', 'leye_pose', 'reye_pose', 'lhand_pose', 'rhand_pose', 'expr', 'trans', 'betas'] # 'flame_params': flame_params, # dict: body_pose:[N, 21, 3], ret.update(flame_params) ret.update(source_flame_params) return ret def gen_valid_id_json(): root_dir = "./train_data/vfhq_vhap/export" save_path = "./train_data/vfhq_vhap/label/valid_id_list.json" os.makedirs(os.path.dirname(save_path), exist_ok=True) valid_id_list = [] for file in os.listdir(root_dir): if not file.startswith("."): valid_id_list.append(file) print(len(valid_id_list), valid_id_list[:2]) with open(save_path, "w") as fp: json.dump(valid_id_list, fp) def gen_valid_id_json(): root_dir = "./train_data/vfhq_vhap/export" mask_root_dir = "./train_data/vfhq_vhap/mask" save_path = "./train_data/vfhq_vhap/label/valid_id_list.json" os.makedirs(os.path.dirname(save_path), exist_ok=True) valid_id_list = [] for file in os.listdir(root_dir): if not file.startswith(".") and ".txt" not in file: valid_id_list.append(file) print("raw:", len(valid_id_list), valid_id_list[:2]) mask_valid_id_list = [] for file in os.listdir(mask_root_dir): if not file.startswith(".") and ".txt" not in file: mask_valid_id_list.append(file) print("mask:", len(mask_valid_id_list), mask_valid_id_list[:2]) valid_id_list = list(set(valid_id_list).intersection(set(mask_valid_id_list))) print("intesection:", len(mask_valid_id_list), mask_valid_id_list[:2]) with open(save_path, "w") as fp: json.dump(valid_id_list, fp) save_train_path = "./train_data/vfhq_vhap/label/valid_id_train_list.json" save_val_path = "./train_data/vfhq_vhap/label/valid_id_val_list.json" valid_id_list = sorted(valid_id_list) idxs = np.linspace(0, len(valid_id_list)-1, num=20, endpoint=True).astype(np.int64) valid_id_train_list = [] valid_id_val_list = [] for i in range(len(valid_id_list)): if i in idxs: valid_id_val_list.append(valid_id_list[i]) else: valid_id_train_list.append(valid_id_list[i]) print(len(valid_id_train_list), len(valid_id_val_list), valid_id_val_list) with open(save_train_path, "w") as fp: json.dump(valid_id_train_list, fp) with open(save_val_path, "w") as fp: json.dump(valid_id_val_list, fp) if __name__ == "__main__": import trimesh import cv2 root_dir = "./train_data/vfhq_vhap/export" meta_path = "./train_data/vfhq_vhap/label/valid_id_list.json" dataset = VideoHeadDataset(root_dirs=root_dir, meta_path=meta_path, sample_side_views=15, render_image_res_low=512, render_image_res_high=512, render_region_size=(512, 512), source_image_res=512, enlarge_ratio=[0.8, 1.2], debug=False, is_val=False) from lam.models.rendering.flame_model.flame import FlameHeadSubdivided # subdivided flame subdivide = 2 flame_sub_model = FlameHeadSubdivided( 300, 100, add_teeth=True, add_shoulder=False, flame_model_path='pretrained_models/human_model_files/flame_assets/flame/flame2023.pkl', flame_lmk_embedding_path="pretrained_models/human_model_files/flame_assets/flame/landmark_embedding_with_eyes.npy", flame_template_mesh_path="pretrained_models/human_model_files/flame_assets/flame/head_template_mesh.obj", flame_parts_path="pretrained_models/human_model_files/flame_assets/flame/FLAME_masks.pkl", subdivide_num=subdivide, teeth_bs_flag=False, ).cuda() source_key = "source_rgbs" render_key = "render_image" for idx, data in enumerate(dataset): import boxx boxx.tree(data) if idx > 0: exit(0) os.makedirs("debug_vis/dataloader", exist_ok=True) for i in range(data[source_key].shape[0]): cv2.imwrite(f"debug_vis/dataloader/{source_key}_{i}_b{idx}.jpg", ((data[source_key][i].permute(1, 2, 0).numpy()[:, :, (2, 1, 0)] * 255).astype(np.uint8))) for i in range(data[render_key].shape[0]): cv2.imwrite(f"debug_vis/dataloader/rgbs{i}_b{idx}.jpg", ((data[render_key][i].permute(1, 2, 0).numpy()[:, :, (2, 1, 0)] * 255).astype(np.uint8))) save_root = "./debug_vis/dataloader" os.makedirs(save_root, exist_ok=True) shape = data['betas'].to('cuda') flame_param = {} flame_param['expr'] = data['expr'].to('cuda') flame_param['rotation'] = data['rotation'].to('cuda') flame_param['neck'] = data['neck_pose'].to('cuda') flame_param['jaw'] = data['jaw_pose'].to('cuda') flame_param['eyes'] = data['eyes_pose'].to('cuda') flame_param['translation'] = data['translation'].to('cuda') v_cano = flame_sub_model.get_cano_verts( shape.unsqueeze(0) ) ret = flame_sub_model.animation_forward( v_cano.repeat(flame_param['expr'].shape[0], 1, 1), shape.unsqueeze(0).repeat(flame_param['expr'].shape[0], 1), flame_param['expr'], flame_param['rotation'], flame_param['neck'], flame_param['jaw'], flame_param['eyes'], flame_param['translation'], zero_centered_at_root_node=False, return_landmarks=False, return_verts_cano=True, # static_offset=batch_data['static_offset'].to('cuda'), static_offset=None, ) import boxx boxx.tree(data) boxx.tree(ret) for i in range(ret["animated"].shape[0]): mesh = trimesh.Trimesh() mesh.vertices = np.array(ret["animated"][i].cpu().squeeze()) mesh.faces = np.array(flame_sub_model.faces.cpu().squeeze()) mesh.export(f'{save_root}/animated_sub{subdivide}_{i}.obj') intr = data["intrs"][i] from lam.models.rendering.utils.vis_utils import render_mesh cam_param = {"focal": torch.tensor([intr[0, 0], intr[1, 1]]), "princpt": torch.tensor([intr[0, 2], intr[1, 2]])} render_shape = data[render_key].shape[2:] # int(cam_param['princpt'][1]* 2), int(cam_param['princpt'][0] * 2) face = flame_sub_model.faces.cpu().squeeze().numpy() vertices = ret["animated"][i].cpu().squeeze() c2ws = data["c2ws"][i] w2cs = torch.inverse(c2ws) if data['pytorch3d_transpose_R'][0] > 0: R = w2cs[:3, :3].transpose(1, 0) else: R = w2cs[:3, :3] T = w2cs[:3, 3] vertices = vertices @ R + T mesh_render, is_bkg = render_mesh(vertices, face, cam_param=cam_param, bkg=np.ones((render_shape[0],render_shape[1], 3), dtype=np.float32) * 255, return_bg_mask=True) rgb_mesh = mesh_render.astype(np.uint8) t_image = (data[render_key][i].permute(1, 2, 0)*255).numpy().astype(np.uint8) blend_ratio = 0.7 vis_img = np.concatenate([rgb_mesh, t_image, (blend_ratio * rgb_mesh + (1 - blend_ratio) * t_image).astype(np.uint8)], axis=1) cam_idx = int(data.get('cam_idxs', [i for j in range(16)])[i]) cv2.imwrite(os.path.join(save_root, f"render_{cam_idx}.jpg"), vis_img[:, :, (2, 1, 0)])