Kleinhe/SemanticBoost
0
1from torch import nn2import torch3import numpy as np4from SMPLX.visualize_joint2smpl.simplify_loc2rot import joints2smpl5from motion.hybrik_loc2rot import HybrIKJointsToRotmat6from SMPLX import smplx7from SMPLX.read_from_npy import npy2info, info2dict8from SMPLX.rotation_conversions import *9from motion.dataset.recover_smr import *10from motion.dataset.recover_joints import recover_from_ric as recover_joints11from motion.dataset.paramUtil import t2m_kinematic_chain12from motion.plot3d import plot_3d_motion13import os14import subprocess15import platform16from PIL import Image17from motion.sample import Predictor as mdm_predictor18from TADA.anime import Animation19 20class Visualize(nn.Module):21 def __init__(self, **kargs):22 super(Visualize, self).__init__()23 self.mode = kargs.get("mode", "cadm")24 if self.mode in ["mdm", "cadm", "cadm-augment"]:25 self.predictor = mdm_predictor(**kargs)26 self.rep = self.predictor.rep27 self.smpl_path = kargs.get("smpl_path")28 self.device = kargs.get("device", "cpu")29 self.rotate = kargs.get("rotate", 0)30 self.pose_generator = HybrIKJointsToRotmat()31 self.path = kargs["path"]32 33 self.tada_base = kargs.get("tada_base", None)34 self.tada_role = kargs.get("tada_role", None)35 36 if self.tada_base is not None and self.tada_role is not None:37 self.anime = Animation(self.tada_role, self.tada_base, self.device)38 self.face = None39 else:40 self.face = np.load(os.path.join(self.path["dataset_dir"], "smplh.faces"))41 self.anime = None42 43 def fit2smpl(self, motion, mode="fast"):44 print(">>>>>>>>>>>>>>> fit joints to smpl >>>>>>>>>>>>>>>>>>>>")45 if mode == "slow":46 frames = motion.shape[0]47 j2s = joints2smpl(num_frames=frames, device=self.device, model_path=self.smpl_path, json_dict=self.path)48 motion_tensor, translation = j2s.joint2smpl(motion)49 else:50 translation = motion[:, 0:1, :] - motion[0, 0:1, :]51 motion = self.pose_generator(motion)52 motion = torch.from_numpy(motion)53 hand = torch.eye(3).unsqueeze(0).unsqueeze(0).repeat(motion.shape[0], 2, 1, 1)54 motion = torch.cat([motion, hand], dim=1)55 motion_tensor = matrix_to_axis_angle(motion)56 motion_tensor = motion_tensor.numpy()57 58 return motion_tensor, translation59 60 def predict(self, sentence, path, render_mode="pyrender", joint_path=None, smpl_path=None):61 if self.mode == "pose":62 motion_tensor = np.load(path)63 if render_mode == "joints":64 _, joints = self.get_mesh(motion_tensor)65 motion_tensor = joints66 67 elif self.mode == "joints":68 joints = np.load(path)69 if render_mode == "joints":70 motion_tensor = joints71 else:72 motion_tensor, translation = self.fit2smpl(joints, render_mode.split("_")[-1])73 motion_tensor = np.concatenate([motion_tensor, translation], axis=1)74 motion_tensor = motion_tensor.reshape(motion_tensor.shape[0], -1) 75 elif self.mode in ["mdm", "cadm", "cadm-augment"]:76 motion_tensor = self.predictor.predict(sentence, 1, path)77 if self.rep == "t2m":78 motion_tensor = motion_tensor[0].detach().cpu().numpy() #### [nframes, 263]79 80 if joint_path is not None:81 np.save(joint_path, motion_tensor)82 83 if render_mode == "joints":84 motion_tensor = motion_tensor85 else:86 motion_tensor, translation = self.fit2smpl(motion_tensor, render_mode.split("_")[-1])87 motion_tensor = np.concatenate([motion_tensor, translation], axis=1) 88 motion_tensor = motion_tensor.reshape(motion_tensor.shape[0], -1)89 90 if smpl_path is not None:91 np.save(smpl_path, motion_tensor)92 93 elif self.rep == "smr":94 motion_tensor = motion_tensor[0][0].detach().cpu().numpy()95 joints = recover_from_ric(motion_tensor, 22)96 97 if joint_path is not None:98 np.save(joint_path, joints)99 100 if render_mode == "joints":101 motion_tensor = joints102 else:103 pose = recover_pose_from_smr(motion_tensor, 22)104 pose = pose.reshape(pose.shape[0], -1, 3)105 motion_tensor, translation = self.fit2smpl(joints, render_mode.split("_")[-1])106 motion_tensor = np.concatenate([motion_tensor, translation], axis=1)107 motion_tensor = motion_tensor.reshape(motion_tensor.shape[0], -1, 3)108 replace = [12, 15, 20, 21]109 motion_tensor[:, replace, :] = pose[:, replace, :]110 motion_tensor = motion_tensor.reshape(motion_tensor.shape[0], -1)111 112 if smpl_path is not None:113 np.save(smpl_path, motion_tensor)114 115 return motion_tensor.astype(np.float32)116 117 def joints_process(self, joints, text, width=1024, height=1024):118 os.makedirs("temp", exist_ok=True)119 plot_3d_motion(t2m_kinematic_chain, joints, text, figsize=(width/100, height/100))120 files = os.listdir("temp")121 files = sorted(files)122 pics = []123 for i in range(len(files)):124 pic = Image.open(os.path.join("temp", files[i]))125 pic = np.asarray(pic)126 pics.append(pic.copy())127 128 cmd = "rm -r temp"129 subprocess.call(cmd, shell=platform.system() != 'Windows')130 pics = np.stack(pics, axis=0)131 return pics132 133 def pyrender_process(self, vertices, height=1024, weight=1024):134 import trimesh135 from trimesh import Trimesh136 import pyrender137 from pyrender.constants import RenderFlags138 import os139 os.environ['PYOPENGL_PLATFORM'] = "egl"140 from shapely import geometry141 from tqdm import tqdm142 143 faces = self.face144 145 vertices = vertices.astype(np.float32)146 MINS = np.min(np.min(vertices, axis=0), axis=0)147 MAXS = np.max(np.max(vertices, axis=0), axis=0)148 149 #################### position initial at zero point150 vertices[:, :, 0] -= (MAXS + MINS)[0] / 2151 vertices[:, :, 2] -= (MAXS + MINS)[2] / 2152 153 MINS = np.min(np.min(vertices, axis=0), axis=0)154 MAXS = np.max(np.max(vertices, axis=0), axis=0)155 156 pics = []157 158 ############### ground initial ###########159 minx = MINS[0] - 0.5160 maxx = MAXS[0] + 0.5161 minz = MINS[2] - 0.5 162 maxz = MAXS[2] + 0.5163 polygon = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz], [maxx, minz]])164 polygon_mesh = trimesh.creation.extrude_polygon(polygon, 1e-5)165 polygon_mesh.visual.face_colors = [0, 0, 0, 0.21]166 polygon_render = pyrender.Mesh.from_trimesh(polygon_mesh, smooth=False)167 168 r = pyrender.OffscreenRenderer(weight, height)169 170 for i in tqdm(range(vertices.shape[0])):171 end_color = np.array([30, 128, 255]) / 255.0172 173 bg_color = [1, 1, 1, 0.8]174 scene = pyrender.Scene(bg_color=bg_color, ambient_light=(0.4, 0.4, 0.4))175 176 if self.anime is None:177 mesh = Trimesh(vertices=vertices[i, :, :].tolist(), faces=faces) 178 base_color = end_color.tolist()179 material = pyrender.MetallicRoughnessMaterial(180 metallicFactor=0.7, roughnessFactor=0.7,181 alphaMode='OPAQUE',182 baseColorFactor=base_color183 )184 mesh = pyrender.Mesh.from_trimesh(mesh, material=material)185 else:186 mesh = Trimesh(vertices=vertices[i, :, :].tolist(), faces=faces, visual=self.anime.trimesh_visual, process=False)187 mesh = pyrender.Mesh.from_trimesh(mesh, smooth=True, material=None) 188 189 scene.add(mesh)190 191 ########################### ground ##################192 c = np.pi / 2193 scene.add(polygon_render, pose=np.array([[ 1, 0, 0, 0],194 [ 0, np.cos(c), -np.sin(c), MINS[1]],195 [ 0, np.sin(c), np.cos(c), 0],196 [ 0, 0, 0, 1]]))197 198 ################ light ############199 light = pyrender.DirectionalLight(color=[1,1,1], intensity=300)200 light_pose = np.eye(4)201 light_pose[:3, 3] = [0, -1, 1]202 scene.add(light, pose=light_pose.copy())203 light_pose[:3, 3] = [0, 1, 1]204 scene.add(light, pose=light_pose.copy())205 light_pose[:3, 3] = [1, 1, 2]206 scene.add(light, pose=light_pose.copy())207 208 ################ camera ##############209 camera = pyrender.PerspectiveCamera(yfov=(np.pi / 3.0))210 c = -np.pi / 6211 scene.add(camera, pose=[[ 1, 0, 0, (minx+maxx)],212 [ 0, np.cos(c), -np.sin(c), 2.5],213 [ 0, np.sin(c), np.cos(c), max(4, minz+(1.5-MINS[1])*2, (maxx-minx))],214 [ 0, 0, 0, 1]215 ])216 217 pic, _ = r.render(scene, flags=RenderFlags.RGBA)218 pics.append(pic)219 220 pics = np.stack(pics, axis=0)221 return pics222 223 @torch.no_grad()224 def get_mesh(self, motions):225 if self.anime is not None:226 vertices, faces = self.anime.forward_mdm(motions)227 joints = vertices228 self.face = faces229 else:230 motions, trans, gender, betas = npy2info(motions, 10)231 232 betas = None233 gender = "neutral"234 235 if motions.shape[1] == 72:236 mode = "smpl"237 elif motions.shape[1] == 156:238 mode = "smplh"239 elif motions.shape[1] == 165:240 motions = np.concatenate([motions[:, :66], motions[:, 75::]], axis=1)241 mode = "smplh"242 243 if self.rotate != 0:244 motions = motions.reshape(motions.shape[0], -1, 3)245 motions = torch.from_numpy(motions).float()246 first_frame_root_pose_matrix = axis_angle_to_matrix(motions[0][0])247 all_root_poses_matrix = axis_angle_to_matrix(motions[:, 0, :])248 aligned_root_poses_matrix = torch.matmul(torch.transpose(first_frame_root_pose_matrix, 0, 1),249 all_root_poses_matrix)250 motions[:, 0, :] = matrix_to_axis_angle(aligned_root_poses_matrix)251 motions = motions.reshape(motions.shape[0], -1)252 motions = motions.numpy()253 254 print("Visualize Mode -> ", mode)255 model = smplx.create(self.smpl_path, model_type=mode,256 gender=gender, use_face_contour=True,257 num_betas=10,258 num_expression_coeffs=10,259 ext="npz", use_pca=False, batch_size=motions.shape[0])260 model = model.eval().to(self.device)261 262 inputs = info2dict(motions, trans, betas, mode, self.device)263 264 output = model(**inputs)265 266 vertices = output.vertices.cpu().numpy()267 joints = output.joints.cpu().numpy()268 269 return vertices, joints