def demo_single_video_switch(target_obj_path = None, target_front_lmark_path = None, target_rt_path = None, original_front_lmark_path = None, original_key_id = None, pr_path = None, ani_save_path = None): itvl = 1000.0/25.0 # 25fps overlay = False # extract the frontal facial landmarks for key frame original_lmk3d_all = np.load(original_front_lmark_path) original_lmk3d_target = original_lmk3d_all[original_key_id] # load the 3D facial landmarks on the PRNet 3D reconstructed face lmk3d_origin = np.load(pr_path) # load RTs for all frame rots, trans = recover(np.load( target_rt_path)) # calculate the affine transformation between PRNet 3D face and the frotal face landmarks lmk3d_origin_homo = np.hstack((lmk3d_origin, np.ones([lmk3d_origin.shape[0],1]))) # 68x4 p_affine = np.linalg.lstsq(lmk3d_origin_homo, original_lmk3d_target, rcond=1)[0].T # Affine matrix. 3 x 4 pr = p_affine[:,:3] # 3x3 pt = p_affine[:,3:] # 3x1 # load the original 3D face mesh then transform it to align frontal face landmarks vertices_org, triangles, colors = load_obj(target_obj_path) # get unfrontalized vertices position vertices_origin_affine = (pr @ (vertices_org.T) + pt).T # aligned vertices # set up the renderer renderer = setup_renderer() ani_path =ani_save_path fig = plt.figure() ims = [] temp_path = './tempp_00005' if os.path.exists(temp_path): shutil.rmtree(temp_path) os.mkdir(temp_path) face_mesh = sr.Mesh(vertices_org, triangles, colors, texture_type="vertex") image_render = get_np_uint8_image(face_mesh, renderer) # RGBA, (224,224,3), np.uint8 # #####save rgba image as bgr in cv2 # rgb_frame = (image_render).astype(int)[:,:,:-1][...,::-1] # # flipBoth = cv2.flip(rgb_frame, 1) # flipBoth = cv2.flip(rgb_frame, 0) # cv2.imwrite( temp_path + "/" + id + "_original.png", flipBoth) for i in range(rots.shape[0]): # if i == 510: # break # get rendered frame vertices = (rots[i].T @ (vertices_origin_affine.T - trans[i])).T face_mesh = sr.Mesh(vertices, triangles, colors, texture_type="vertex") image_render = get_np_uint8_image(face_mesh, renderer) # RGBA, (224,224,3), np.uint8 #save rgba image as bgr in cv2 rgb_frame = (image_render).astype(int)[:,:,:-1][...,::-1] cv2.imwrite( temp_path + "/%05d.png"%i, rgb_frame) command = 'ffmpeg -framerate 25 -i ' + temp_path + '/%5d.png -c:v libx264 -y -vf format=yuv420p ' + ani_path os.system(command) print (command)
def forward(self, vertices, faces, textures=None): vs = vertices vs[:, :, 1] *= -1 fs = faces if(textures is None): mesh_ = sr.Mesh(vs, fs) else: ts = textures mesh_ = sr.Mesh(vs, fs, ts) imgs = self.renderer.render_mesh(mesh_) return imgs
def forward(self, inputs): x = inputs.view(-1, 15 * 15 * 256) x = self.fc1(x) # x = self.fc2(x) # x = self.fc3(x) base = torch.log(self.vertices.abs() / (1 - self.vertices.abs())).repeat( inputs.size(0), 1, 1) centroid = torch.tanh(self.fc_center(x).view(-1, 1, 3)).repeat( 1, self.template_mesh.num_vertices, 1) displace = self.fc_displace(x).view(-1, self.template_mesh.num_vertices, 3) vertices = torch.sigmoid(base + displace) * torch.sign(self.vertices) vertices = F.relu(vertices) * (1 - centroid) - F.relu(-vertices) * ( centroid + 1) vertices = vertices + centroid # vertices = torch.tanh(self.fc_displace(x)).view(-1, self.template_mesh.num_vertices, 3) # apply Laplacian and flatten geometry constraints laplacian_loss = self.laplacian_loss(vertices).mean() flatten_loss = self.flatten_loss(vertices).mean() return sr.Mesh(vertices, self.faces.repeat(inputs.size(0), 1, 1)), laplacian_loss, flatten_loss
def forward(self, displace, center, numViews, numBatch, texture): # center, displace would be batchx3, batchxnumvertx3 #vertices = self.vertices + displace + center base = torch.log(self.vertices.abs() / (1 - self.vertices.abs())) centroid = torch.tanh(center) vertices = torch.sigmoid(base + displace) * torch.sign(self.vertices) # need to figure out this transformation vertices = F.relu(vertices) * (1 - centroid) - F.relu(-vertices) * ( centroid + 1) vertices = vertices + centroid # apply Laplacian and flatten geometry constraints laplacian_loss = self.laplacian_loss(vertices).mean() flatten_loss = self.flatten_loss(vertices).mean() return sr.Mesh(vertices.repeat(1, numViews, 1).reshape(numViews * numBatch, -1, 3), self.faces.repeat(1, numViews, 1).reshape(numViews * numBatch, -1, 3), textures=texture.repeat(1, numViews, 1).reshape( numViews * numBatch, -1, 3), texture_type='vertex'), laplacian_loss, flatten_loss
def forward(self, mesh): if mesh.texture_type == 'surface': light = torch.zeros_like(mesh.faces, dtype=torch.float32).to(mesh.device) light = light.contiguous() light = self.ambient(light) for directional in self.directionals: light = directional(light, mesh.surface_normals) new_textures = mesh.textures * light[:, :, None, :] elif mesh.texture_type == 'vertex': light = torch.zeros_like(mesh.vertices, dtype=torch.float32).to(mesh.device) light = light.contiguous() light = self.ambient(light) for directional in self.directionals: light = directional(light, mesh.vertex_normals) new_textures = mesh.textures * light vertices = mesh.vertices faces = mesh.faces texture_res = mesh.texture_res texture_type = mesh.texture_type mesh = sr.Mesh(vertices, faces, new_textures, texture_res, texture_type) return mesh
def save_current_visuals(self, unsup=False): def norm_ip(img, min, max): img.clamp_(min=min, max=max) img.add_(-min).div_(max - min + 1e-5) if(unsup): visuals = self.get_current_visuals_unsup() else: visuals = self.get_current_visuals() imgs_dir = osp.join(self.opts.results_vis_dir, 'vis_iter_{}'.format(self.vis_iter)) if not os.path.exists(imgs_dir): os.makedirs(imgs_dir) for k in visuals: if("mesh" in k): meshes = visuals[k] verts = meshes.vertices faces = meshes.faces tex = meshes.textures for cnt in range(verts.size(0)): mesh_path = osp.join(imgs_dir, k + '_' + str(cnt) + '.obj') mesh_ = sr.Mesh(verts[cnt], faces[cnt], tex[cnt]) mesh_.save_obj(mesh_path, save_texture=self.opts.use_texture) elif("GIF" in k): for cnt in range(self.opts.sample_num): vis_image = torch.cat((visuals['gt_imgs'][cnt].unsqueeze(0).cpu(), visuals[k][cnt].cpu()), dim = 0) nrow = vis_image.size(0) vutils.save_image(vis_image, os.path.join(imgs_dir, "{}.png".format(cnt)), normalize=True, scale_each=True, nrow=nrow) else: img_path = osp.join(imgs_dir, k + '.png') vutils.save_image(visuals[k], img_path, normalize=False) self.vis_iter += 1
def forward(self, mesh): new_vertices = self.transform(mesh.vertices) faces = mesh.faces textures = mesh.textures texture_res = mesh.texture_res texture_type = mesh.texture_type return sr.Mesh(new_vertices, faces, textures, texture_res, texture_type)
def forward(self, vertices, faces, textures=None, return_buffers=False): mesh = sr.Mesh(vertices, faces, textures=textures) if self.is_transform: mesh = self.transform(mesh) face_vertices = mesh.face_vertices face_vertices[..., 1] = -face_vertices[..., 1] face_vertices = mesh.face_vertices * self.image_size / 2 + self.image_size / 2 # face_colors = srf.face_vertices(faces, textures) # images = # fnorm = F.normalize(torch.cross(v12, v10), p=2, dim=2, eps=1e-6) # return self.rasterizer(face_vertices, textures, return_buffers)
def execute(self, vertices, faces, textures=None, mode=None, texture_type='surface'): mesh = sr.Mesh(vertices, faces, textures=textures, texture_type=texture_type) return self.render_mesh(mesh, mode)
def main(): filename_input = os.path.join(data_dir, 'banana.obj') filename_output = os.path.join(output_directory, 'example1.gif') ########################### # camera settings ########################### camera_distance = 2 elevation = 30 ########################### # load object ########################### mesh = TriangleMesh.from_obj(filename_input) vertices = mesh.vertices faces = mesh.faces.int() face_textures = (faces).clone() vertices = vertices[None, :, :].cuda() faces = faces[None, :, :].cuda() face_textures[None, :, :].cuda() ########################### # normalize verts ########################### vertices_max = vertices.max() vertices_min = vertices.min() vertices_middle = (vertices_max + vertices_min) / 2. vertices = vertices - vertices_middle coef = 5 vertices = vertices * coef ########################### # Soft Rasterizer ########################### textures = torch.ones(1, faces.shape[1], 2, 3, dtype=torch.float32).cuda() mesh = sr.Mesh(vertices, faces, textures) renderer = sr.SoftRenderer(camera_mode='look_at') loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing SR') writer = imageio.get_writer(os.path.join(output_directory_sr, 'rotation.gif'), mode='I') for azimuth in loop: mesh.reset_() renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) images = renderer.render_mesh(mesh) image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) writer.close()
def forward(self, vertices, faces, textures=None, mode=None, texture_type='surface', R=None, t=None): mesh = sr.Mesh(vertices, faces, textures=textures, texture_type=texture_type) return self.render_mesh(mesh, mode)
def __init__(self, faces, vertices): super(MeshModel, self).__init__() # set template mesh # Assuming faces, vertices are batch, num, 3 # The mesh class they have is such that they work with multiple meshes.. # So the vertices and faces here would be for multiple(batchNum) meshes self.template_mesh = sr.Mesh(vertices, faces) self.vertices = self.template_mesh.vertices self.faces = self.template_mesh.faces #TODO : check if this works for multiple meshes.. # Mesh connectivitiy is same!!! self.laplacian_loss = losses.LaplacianLoss(self.vertices[0].cpu(), self.faces[0].cpu()) self.flatten_loss = losses.FlattenLoss(self.faces[0].cpu())
def get_current_visuals(self): vis_dict = {} # UV maps if self.opts.use_texture: uv_flows = self.uvimage_pred uv_flows = uv_flows.permute(0, 2, 3, 1) uv_images = torch.nn.functional.grid_sample(self.imgs, uv_flows) vis_dict['uv_images'] = uv_images # mask vis_dict['mask_pred'] = self.mask_pred_seen.unsqueeze(1) nb, nf, _, nc = self.tex.size() tex = self.tex.detach().view(nb, nf, opts.tex_size, opts.tex_size, nc).unsqueeze(4).repeat(1, 1, 1, 1, opts.tex_size, 1) vis_dict['mask_gt'] = self.masks.unsqueeze(1) # image vis_dict['image_pred'] = self.vis_renderer(self.pred_vs.detach(), self.faces, self.proj_cam.detach(), tex) vis_dict['image_gt'] = self.imgs * self.masks.unsqueeze(1).repeat(1, 3, 1, 1) # instance mesh if(self.opts.use_texture): mesh_ = sr.Mesh(self.pred_vs[0], self.faces[0], self.tex[0].view(self.faces.size(1),-1,3)) else: mesh_ = sr.Mesh(self.pred_vs[0], self.faces[0]) vis_dict['mesh'] = mesh_ # template mesh if(opts.multi_gpu): template_mesh_ = sr.Mesh(self.model.module.get_mean_shape(), self.faces[0]) else: template_mesh_ = sr.Mesh(self.model.get_mean_shape(), self.faces[0]) vis_dict['template_mesh'] = template_mesh_ return vis_dict
def forward(self, batch_size): base = torch.log(self.vertices.abs() / (1 - self.vertices.abs())) centroid = torch.tanh(self.center) vertices = torch.sigmoid(base + self.displace) * torch.sign( self.vertices) vertices = F.relu(vertices) * (1 - centroid) - F.relu(-vertices) * ( centroid + 1) vertices = vertices + centroid # apply Laplacian and flatten geometry constraints laplacian_loss = self.laplacian_loss(vertices).mean() flatten_loss = self.flatten_loss(vertices).mean() return sr.Mesh(vertices.repeat(batch_size, 1, 1), self.faces.repeat(batch_size, 1, 1)), laplacian_loss, flatten_loss
def save_as_obj(verts, faces, output_path, verbose=False): assert len(verts.shape) == 3 assert len(faces.shape) == 3 if torch.cuda.is_available() is not True: return None # soft renderer is only supported under cuda else: if verbose: print("saving as obj...") # prepare for output output_path = os.path.splitext( output_path)[0] + '.obj' # replace extention by .obj os.makedirs(os.path.dirname(output_path), exist_ok=True) # make output dir if verbose: print("output_path: {}".format(output_path)) # make mesh mesh = sr.Mesh(verts[0, :, :], faces[0, :, :]) mesh.save_obj(output_path)
def forward(self, vertices, faces, textures=None, sh_coeff=None, mode=None, texture_type='surface', maskout=True, return_mask=False, is_lighting=False): tex_bary_weights = textures[..., 3:].detach() textures = textures[..., :3] #.detach() mesh = sr.Mesh(vertices, faces, textures=textures, texture_type=texture_type) if is_lighting: mesh = self.lighting(mesh, sh_coeff, tex_bary_weights) self.set_texture_mode(mesh.texture_type) mesh = self.transform(mesh) images = self.rasterizer(mesh, mode) # fix mouth issue: mask out inner mouth part # Observation: the rendered vertices in mouth part actually belong to back head, which have negative normal directions # Method: 1. calculate normal direction of each vertex. 2. set color of negtive normal to 0, others 1 # transformed_vertices = torch.stack([mesh.vertices[b,mesh.faces.detach().long()[b,...], :] for b in range(faces.shape[0])]) # v10 = transformed_vertices[:, :, 0] - transformed_vertices[:, :, 1] # v12 = transformed_vertices[:, :, 2] - transformed_vertices[:, :, 1] # fnorm = F.normalize(torch.cross(v12, v10), p=2, dim=2, eps=1e-6) # # fmaskcolor = fnorm[:,:,2:].gt(0).float() # mask_textures = fmaskcolor[:,:,None,:].expand(textures.shape) # mesh_mask = sr.Mesh(mesh.vertices.detach(), faces, textures=mask_textures.detach(), texture_type=texture_type) # masks = self.rasterizer(mesh_mask, mode).detach() # if maskout: # images = images*masks if return_mask: return images, masks else: return images
def forward(self, batch_size): base = torch.log(self.vertices.abs() / (1 - self.vertices.abs())) centroid = torch.tanh(self.center) vertices = torch.sigmoid(base + self.displace) * torch.sign(self.vertices) vertices = F.relu(vertices) * (1 - centroid) - F.relu(-vertices) * (centroid + 1) vertices = vertices + centroid all_vertices = torch.cat([vertices, self.bg_vertices], dim=1) idx_offset = self.template_mesh.vertices.shape[1] all_faces = self.faces if OPTIMIZE_SILHOUETTE else torch.cat([self.faces, self.bg_faces + idx_offset], dim=1) # apply Laplacian and flatten geometry constraints laplacian_loss = self.laplacian_loss(vertices).mean() flatten_loss = self.flatten_loss(vertices).mean() return sr.Mesh(all_vertices.repeat(batch_size, 1, 1), all_faces.repeat(batch_size, 1, 1), textures=self.all_textures.repeat(batch_size, 1, 1), texture_type='vertex'), laplacian_loss, flatten_loss
def render_multiview(self, image_a, image_b, viewpoint_a, viewpoint_b): # [Ia, Ib] images = torch.cat((image_a, image_b), dim=0) # [Va, Va, Vb, Vb], set viewpoints viewpoints = torch.cat((viewpoint_a, viewpoint_a, viewpoint_b, viewpoint_b), dim=0) self.transform.set_eyes(viewpoints) vertices, faces = self.reconstruct(images) laplacian_loss = self.laplacian_loss(vertices) flatten_loss = self.flatten_loss(vertices) # [Ma, Mb, Ma, Mb] vertices = torch.cat((vertices, vertices), dim=0) faces = torch.cat((faces, faces), dim=0) # [Raa, Rba, Rab, Rbb], render for cross-view consistency mesh = sr.Mesh(vertices, faces) mesh = self.lighting(mesh) mesh = self.transform(mesh) silhouettes = self.rasterizer(mesh) return silhouettes.chunk(4, dim=0), laplacian_loss, flatten_loss
def forward(self, vertices, faces, textures=None, sh_coeff=None, maskout=False, return_mask=False, is_lighting=False, is_normal=False): mesh = sr.Mesh(vertices, faces, textures=textures, texture_type=self.texture_type) # import ipdb; ipdb.set_trace() if is_lighting: if self.light_mode == 'directional': lights = self.lighting(mesh) elif self.light_mode == 'spherical': lights = self.lighting(mesh, sh_coeff) mesh.textures = torch.cat([mesh.textures, lights], -1) if self.is_transform: mesh = self.transform(mesh) if is_normal: transformed_vertices = torch.stack([ mesh.vertices[b, mesh.faces.detach().long()[b, ...], :] for b in range(faces.shape[0]) ]) v10 = transformed_vertices[:, :, 0] - transformed_vertices[:, :, 1] v12 = transformed_vertices[:, :, 2] - transformed_vertices[:, :, 1] fnorm = F.normalize(torch.cross(v12, v10), p=2, dim=2, eps=1e-6) mesh.textures = torch.cat( [mesh.textures, fnorm[:, :, None, :].expand(-1, -1, 3, -1)], -1) images = self.rasterizer(mesh) # fnorm = F.normalize(torch.cross(v12, v10), p=2, dim=2, eps=1e-6) # return images
def render_mesh_rt(vertices, triangles, colors, rot=None, tran=None): ''' Apply the rt(rotation and translation) to a mesh then render it to an image Args: vertices: torch.Tensor, shape = [nver, 3] (on gpu) triangles: torch.Tensor, shape = [ntri, 3] (on gpu) colors: torch.Tensor, shape = [nver, 3] (on gpu) r: torch.Tensor, shape=[3,3] (on gpu) t: torch.Tensor, shape=[3,1] (on gpu) Return: image: torch.Tensor, shape=[4, width, height] (on gpu) ''' # vertices = torch.Tensor(vertices).cuda() # triangles = torch.Tensor(triangles).cuda().int() # colors = torch.Tensor(colors).cuda() # rot = torch.Tensor(rot).cuda() # tran = torch.Tensor(trans).cuda() # apply r and t if type(rot) == torch.Tensor and type(tran) == torch.Tensor: # new_vertices = (rot @ vertices.transpose(1,0) + tran).transpose(1,0) new_vertices = (rot.T @ (vertices.transpose(1,0) - tran)).transpose(1,0) else: new_vertices = vertices # prepare the mesh face_mesh = sr.Mesh(new_vertices, triangles, colors, texture_type="vertex") # prepare the renderer # renderer = sr.SoftRenderer(camera_mode="look_at", far=10000, image_size=512, camera_direction=[0,1,0], light_directions=[0,0,1], dist_func="hard") # the dist_func can be 'hard', 'barycentric' or 'euclidean' # renderer.transform.set_eyes_from_angles(-400, 0, 0) # this three number can be trainable parameters as well renderer = sr.SoftRenderer(camera_mode="look", far=10000, image_size=224, viewing_angle=15, camera_direction=[0,0,-1], camera_up=[-0.3,1,0], light_intensity_ambient=1, light_color_ambient=[1,1,1], dist_func="hard") renderer.transform.set_eyes([0,-50,680]) # do the rendering images = renderer.render_mesh(face_mesh) image = images[0] image = torch.flip(image, [2]) return image
def vis_single(video_path, key_id, save_name): overlay = True #key_id = 79 # #video_path = '/home/cxu-serve/p1/lchen63/voxceleb/unzip/test_video/id04276/k0zLls_oen0/00341.mp4' reference_img_path = video_path[:-4] + '_%05d.png'%key_id reference_prnet_lmark_path = video_path[:-4] +'_prnet.npy' original_obj_path = video_path[:-4] + '_original.obj' rt_path = video_path[:-4] + '_sRT.npy' lmark_path = video_path[:-4] +'_front.npy' # extract the frontal facial landmarks for key frame lmk3d_all = np.load(lmark_path) lmk3d_target = lmk3d_all[key_id] # load the 3D facial landmarks on the PRNet 3D reconstructed face lmk3d_origin = np.load(reference_prnet_lmark_path) # lmk3d_origin[:,1] = res - lmk3d_origin[:,1] # load RTs rots, trans = recover(np.load(rt_path)) # calculate the affine transformation between PRNet 3D face and the frotal face landmarks lmk3d_origin_homo = np.hstack((lmk3d_origin, np.ones([lmk3d_origin.shape[0],1]))) # 68x4 p_affine = np.linalg.lstsq(lmk3d_origin_homo, lmk3d_target, rcond=1)[0].T # Affine matrix. 3 x 4 pr = p_affine[:,:3] # 3x3 pt = p_affine[:,3:] # 3x1 # load the original 3D face mesh then transform it to align frontal face landmarks vertices_org, triangles, colors = load_obj(original_obj_path) # get unfrontalized vertices position vertices_origin_affine = (pr @ (vertices_org.T) + pt).T # aligned vertices # set up the renderer renderer = setup_renderer() # generate animation if os.path.exists('./tempo1'): shutil.rmtree('./tempo1') os.mkdir('./tempo1') if overlay: real_video = mmcv.VideoReader(video_path) for i in range(rots.shape[0]): t = time.time() # get rendered frame vertices = (rots[i].T @ (vertices_origin_affine.T - trans[i])).T face_mesh = sr.Mesh(vertices, triangles, colors, texture_type="vertex") image_render = get_np_uint8_image(face_mesh, renderer) # RGBA, (224,224,3), np.uint8 print (image_render.shape) print (image_render.max()) print (image_render.min()) #save rgba image as bgr in cv2 rgb_frame = (image_render ).astype(int)[:,:,:-1][...,::-1] overla_frame = (0.5* rgb_frame + 0.5 * real_video[i]).astype(int) cv2.imwrite("./tempo1/%05d.png"%i, overla_frame) print (time.time() - t) # writer.append_data((255*warped_image).astype(np.uint8)) print("[{}/{}]".format(i+1, rots.shape[0])) # if i == 5: # breakT t = time.time() ani_mp4_file_name = save_name # './f**k.mp4' command = 'ffmpeg -framerate 25 -i ./tempo1/%5d.png -c:v libx264 -y -vf format=yuv420p ' + ani_mp4_file_name os.system(command) print (time.time() - t)
def get_lrs(batch = 0 ): # key_id = 58 # # model_id = "00025" root = '/home/cxu-serve/p1/common/lrs3/lrs3_v0.4' _file = open(os.path.join(root, 'pickle','test2_lmark2img.pkl'), "rb") # _file = open(os.path.join(root, 'txt', "front_rt.pkl"), "rb") data = pickle.load(_file) _file.close() flage = False print (len(data)) for k, v_id in enumerate(data): key_id = v_id[-1] video_path = os.path.join(root, 'test', v_id[0] , v_id[1][:5] + '_crop.mp4' ) ani_path = os.path.join(root, 'test', v_id[0] , v_id[1][:5] + '_ani.mp4' ) reference_img_path = os.path.join(root, 'test', v_id[0] , v_id[1][:5] + '_%05d.png'%key_id ) reference_prnet_lmark_path = os.path.join(root, 'test', v_id[0] , v_id[1][:5] + '_prnet.npy') original_obj_path = os.path.join(root, 'test', v_id[0] , v_id[1][:5] + '_original.obj') rt_path = os.path.join(root, 'test', v_id[0] , v_id[1][:5] +'_rt.npy' ) lmark_path = os.path.join(root, 'test', v_id[0] , v_id[1][:5] +'_front.npy' ) if os.path.exists( ani_path): print (ani_path) print ('=====') continue if not os.path.exists(original_obj_path) or not os.path.exists(reference_prnet_lmark_path) or not os.path.exists(lmark_path) or not os.path.exists(rt_path): print (os.path.exists(original_obj_path) , os.path.exists(reference_prnet_lmark_path), os.path.exists(lmark_path), os.path.exists(rt_path)) print (original_obj_path) print ('++++') continue # try: # extract the frontal facial landmarks for key frame lmk3d_all = np.load(lmark_path) lmk3d_target = lmk3d_all[key_id] # load the 3D facial landmarks on the PRNet 3D reconstructed face lmk3d_origin = np.load(reference_prnet_lmark_path) # lmk3d_origin[:,1] = lmk3d_origin[:,1] + 32 # load RTs rots, trans = recover(np.load(rt_path)) # calculate the affine transformation between PRNet 3D face and the frotal face landmarks lmk3d_origin_homo = np.hstack((lmk3d_origin, np.ones([lmk3d_origin.shape[0],1]))) # 68x4 p_affine = np.linalg.lstsq(lmk3d_origin_homo, lmk3d_target, rcond=1)[0].T # Affine matrix. 3 x 4 pr = p_affine[:,:3] # 3x3 pt = p_affine[:,3:] # 3x1 # load the original 3D face mesh then transform it to align frontal face landmarks vertices_org, triangles, colors = load_obj(original_obj_path) # get unfrontalized vertices position vertices_origin_affine = (pr @ (vertices_org.T) + pt).T # aligned vertices # set up the renderer renderer = setup_renderer() # generate animation temp_path = './tempp_%05d'%batch # generate animation if os.path.exists(temp_path): shutil.rmtree(temp_path) os.mkdir(temp_path) # writer = imageio.get_writer('rotation.gif', mode='I') for i in range(rots.shape[0]): # get rendered frame vertices = (rots[i].T @ (vertices_origin_affine.T - trans[i])).T face_mesh = sr.Mesh(vertices, triangles, colors, texture_type="vertex") image_render = get_np_uint8_image(face_mesh, renderer) # RGBA, (224,224,3), np.uint8 #save rgba image as bgr in cv2 rgb_frame = (image_render).astype(int)[:,:,:-1][...,::-1] cv2.imwrite( temp_path + "/%05d.png"%i, rgb_frame) command = 'ffmpeg -framerate 25 -i ' + temp_path + '/%5d.png -c:v libx264 -y -vf format=yuv420p ' + ani_path os.system(command) print (command)
def demo_obama(): key_id = 405 # index of the frame used to do the 3D face reconstruction (key frame) itvl = 1000.0/25.0 # 25fps overlay = True # extract the frontal facial landmarks for key frame' lmk3d_all = np.load("/home/cxu-serve/p1/common/demo/00025_aligned_front.npy") lmk3d_target = lmk3d_all[key_id] # load the 3D facial landmarks on the PRNet 3D reconstructed face lmk3d_origin = np.load('/home/cxu-serve/p1/common/voxceleb2/unzip/test_video/id04094/2sjuXzB2I1M/00025_prnet.npy') lmk3d_origin[:,1] = lmk3d_origin[:,1] # load RTs for all frame rots, trans = recover(np.load("/home/cxu-serve/p1/common/demo/00025_aligned_rt.npy")) # calculate the affine transformation between PRNet 3D face and the frotal face landmarks lmk3d_origin_homo = np.hstack((lmk3d_origin, np.ones([lmk3d_origin.shape[0],1]))) # 68x4 p_affine = np.linalg.lstsq(lmk3d_origin_homo, lmk3d_target, rcond=1)[0].T # Affine matrix. 3 x 4 pr = p_affine[:,:3] # 3x3 pt = p_affine[:,3:] # 3x1 # load the original 3D face mesh then transform it to align frontal face landmarks vertices_org, triangles, colors = load_obj("/home/cxu-serve/p1/common/Obama/video/3_3__original2.obj") # get unfrontalized vertices position vertices_origin_affine = (pr @ (vertices_org.T) + pt).T # aligned vertices # set up the renderer renderer = setup_renderer() ani_path ='/home/cxu-serve/p1/common/demo/demo_00025_3_3__ani2.mp4' if overlay: real_video = mmcv.VideoReader('/home/cxu-serve/p1/common/Obama/video/3_3__crop2.mp4') fig = plt.figure() ims = [] temp_path = './tempp_00005' # if os.path.exists(temp_path): # shutil.rmtree(temp_path) # os.mkdir(temp_path) face_mesh = sr.Mesh(vertices_org, triangles, colors, texture_type="vertex") image_render = get_np_uint8_image(face_mesh, renderer) # RGBA, (224,224,3), np.uint8 #####save rgba image as bgr in cv2 # rgb_frame = (image_render).astype(int)[:,:,:-1][...,::-1] # # flipBoth = cv2.flip(rgb_frame, 1) # flipBoth = cv2.flip(rgb_frame, 0) # cv2.imwrite( temp_path + "/demo_front.png", flipBoth) for i in range(rots.shape[0]): # if i == 20: # break # get rendered frame vertices = (rots[i].T @ (vertices_origin_affine.T - trans[i])).T face_mesh = sr.Mesh(vertices, triangles, colors, texture_type="vertex") image_render = get_np_uint8_image(face_mesh, renderer) # RGBA, (224,224,3), np.uint8 #save rgba image as bgr in cv2 rgb_frame = (image_render).astype(int)[:,:,:-1][...,::-1] cv2.imwrite( temp_path + "/%05d.png"%i, rgb_frame) command = 'ffmpeg -framerate 25 -i ' + temp_path + '/%5d.png -c:v libx264 -y -vf format=yuv420p ' + ani_path os.system(command) print (command)
def demo_single_video(config, front_lmark_path = None , key_id = None, front_img_path=None, prnet_lmark_path=None, ref_lmark_path = None ): itvl = 1000.0/25.0 # 25fps overlay = False # extract the frontal facial landmarks for key frame' lmk3d_all = np.load(front_lmark_path) if config.same: print('fdjfkdjklfj===') lmk3d_target = lmk3d_all[key_id] else: lmk3d_target = np.load(ref_lmark_path) print(ref_lmark_path) print(lmk3d_target.shape) # load the 3D facial landmarks on the PRNet 3D reconstructed face lmk3d_origin = np.load(prnet_lmark_path) print (lmk3d_target.shape, lmk3d_origin.shape,'+++++++') # load RTs for all frame rots, trans = recover(np.load( front_lmark_path[:-9] + "rt.npy")) # calculate the affine transformation between PRNet 3D face and the frotal face landmarks lmk3d_origin_homo = np.hstack((lmk3d_origin, np.ones([lmk3d_origin.shape[0],1]))) # 68x4 p_affine = np.linalg.lstsq(lmk3d_origin_homo, lmk3d_target, rcond=1)[0].T # Affine matrix. 3 x 4 pr = p_affine[:,:3] # 3x3 pt = p_affine[:,3:] # 3x1 # load the original 3D face mesh then transform it to align frontal face landmarks vertices_org, triangles, colors = load_obj(front_lmark_path[:-9] +"original.obj") # get unfrontalized vertices position vertices_origin_affine = (pr @ (vertices_org.T) + pt).T # aligned vertices # set up the renderer renderer = setup_renderer() ani_path =front_lmark_path[:-9] +'ani.mp4' if overlay: real_video = mmcv.VideoReader('/home/cxu-serve/p1/common/voxceleb2/unzip/test_video/id00419/S8fiWqrZEew/00216_aligned.mp4') fig = plt.figure() ims = [] temp_path = './tempp_00005' if os.path.exists(temp_path): shutil.rmtree(temp_path) os.mkdir(temp_path) face_mesh = sr.Mesh(vertices_org, triangles, colors, texture_type="vertex") image_render = get_np_uint8_image(face_mesh, renderer) # RGBA, (224,224,3), np.uint8 # #####save rgba image as bgr in cv2 # rgb_frame = (image_render).astype(int)[:,:,:-1][...,::-1] # # flipBoth = cv2.flip(rgb_frame, 1) # flipBoth = cv2.flip(rgb_frame, 0) # cv2.imwrite( temp_path + "/" + id + "_original.png", flipBoth) for i in range(rots.shape[0]): # if i == 510: # break # get rendered frame vertices = (rots[i].T @ (vertices_origin_affine.T - trans[i])).T face_mesh = sr.Mesh(vertices, triangles, colors, texture_type="vertex") image_render = get_np_uint8_image(face_mesh, renderer) # RGBA, (224,224,3), np.uint8 #save rgba image as bgr in cv2 rgb_frame = (image_render).astype(int)[:,:,:-1][...,::-1] cv2.imwrite( temp_path + "/%05d.png"%i, rgb_frame) command = 'ffmpeg -framerate 25 -i ' + temp_path + '/%5d.png -c:v libx264 -y -vf format=yuv420p ' + ani_path os.system(command) print (command)
def compute_uv(self): dataloader = iter(self.dataloader) self.best_shape = None best_idx = 0 best_mask_loss = 100000.0 self.best_uv = None print(tf_visualizer.green("Search for the examplar instance.")) for i, batch in tqdm(enumerate(self.dataloader)): self.set_input(batch) with torch.no_grad(): outputs = self.model(self.input_imgs) # shape delta_v = outputs['delta_v'] if (opts.symmetric): delta_v = self.model.symmetrize(delta_v) self.mean_shape = self.model.get_mean_shape() pred_vs = self.mean_shape + delta_v # camera proj_cam = outputs['cam'] faces = self.faces.repeat(delta_v.size(0), 1, 1) ori_flip = torch.FloatTensor([1, -1, 1, 1, 1, -1, -1]).view(1, -1).cuda() proj_cam = proj_cam * ori_flip pred_seen, _, _ = self.renderer.forward( pred_vs, faces, proj_cam) rgb_pred_seen = pred_seen[:, 0:3, :, :] mask_pred_seen = pred_seen[:, 3, :, :] # flip images flip_imgs = torch.flip(self.imgs, [3]) flip_masks = torch.flip(self.masks, [2]) texture_loss = self.texture_loss_fn(rgb_pred_seen, flip_imgs, flip_masks, mask_pred_seen, avg=False) # the new best shape should also be close to the old one, # because we don't want the template to change all the time. mean_shape = self.mean_shape.unsqueeze(0).repeat( pred_vs.size(0), 1, 1) dist = torch.nn.MSELoss(reduction='none')(pred_vs, mean_shape) # scale is used to make sure these two losses are comparable scale = texture_loss.mean() / torch.sum(dist, dim=(1, 2)).mean() texture_loss += (torch.sum(dist, dim=(1, 2))) * scale # select a shape that has both low mask loss and close to the current best_shape min_mask, min_idx = torch.min(texture_loss, 0) if (best_mask_loss > min_mask): best_idx = min_idx best_mask_loss = min_mask self.best_shape = pred_vs[best_idx].unsqueeze(0) uvimage_pred = outputs['uvimage_pred'] uv_parts = torch.nn.functional.grid_sample( self.part_segs.cuda(), uvimage_pred.permute(0, 2, 3, 1)) self.best_uv = uv_parts[best_idx].unsqueeze(0) # visulize semantic texture tex_flow = outputs['tex_flow'][best_idx].unsqueeze(0) tex = geom_utils.sample_textures( tex_flow, self.part_segs[best_idx].unsqueeze(0).cuda()) best_tex = tex.contiguous() bs, fs, ts, _, cs = best_tex.size() best_tex = best_tex.view(bs, fs, -1, cs) best_tex = torch.argmax(best_tex, dim=-1) self.best_tex = self.batch_colorize(best_tex.cpu()) self.best_tex = self.best_tex.permute(0, 2, 3, 1) self.best_shape = self.best_shape.repeat(opts.batch_size, 1, 1) self.best_uv = self.best_uv.repeat(opts.batch_size, 1, 1, 1) print(tf_visualizer.green("Start to compute semantic template.")) counter = 0 avg_uv_parts = None for i, batch in tqdm(enumerate(self.dataloader)): self.set_input(batch) with torch.no_grad(): outputs = self.model(self.input_imgs) self.uvimage_pred = outputs['uvimage_pred'] uv_parts = torch.nn.functional.grid_sample( self.part_segs.cuda(), self.uvimage_pred.permute(0, 2, 3, 1)) uv_parts_ch = uv_parts.clone() best_uv_ch = self.best_uv.clone() dist = torch.nn.MSELoss(reduction='none')(uv_parts_ch, best_uv_ch) dist = torch.sum(dist, dim=(1, 2, 3)) _, idx = torch.topk(dist, k=5, largest=False) if (avg_uv_parts is None): avg_uv_parts = torch.sum(uv_parts[idx, :, :, :], dim=0).unsqueeze(0) else: avg_uv_parts += torch.sum(uv_parts[idx, :, :, :], dim=0).unsqueeze(0) counter += idx.size(0) avg_prob = avg_uv_parts / counter avg_prob = avg_prob.cpu().squeeze().numpy() avg_prob = avg_prob.transpose(1, 2, 0) uv_path = osp.join(opts.out_dir, "semantic_prob.npy") np.save(uv_path, avg_prob) avg_prob = np.asarray(np.argmax(avg_prob, axis=2), dtype=np.int) pil_image = Image.fromarray(avg_prob.astype(dtype=np.uint8)) pil_image.save(osp.join(opts.out_dir, "semantic_seg.png"), 'PNG') color_vis = self.colorize(avg_prob) color_vis = torch.from_numpy(color_vis).float() # wrap the uv map onto template uv_sampler = self.model.uv_sampler[0].unsqueeze(0) tex = torch.nn.functional.grid_sample( color_vis.unsqueeze(0).cuda().float(), uv_sampler) tex = tex.view(tex.size(0), -1, tex.size(2), opts.tex_size, opts.tex_size).permute(0, 2, 3, 4, 1) tex_left = tex[:, -self.model.texture_predictor.num_sym_faces:] tex = torch.cat([tex, tex_left], 1) tex = tex.view(tex.size(0), tex.size(1), -1, 3) mean_v = self.model.get_mean_shape() mesh_ = sr.Mesh(mean_v, self.faces, tex) mesh_path = osp.join(opts.out_dir, "mean_template.obj") mesh_.save_obj(mesh_path, save_texture=True) # compute vertices/faces belong to each part uv_label = np.load(uv_path) uv_label = torch.from_numpy(uv_label).float().unsqueeze(0).permute( 0, 3, 1, 2) uv_label = uv_label.cuda() tex_seg = torch.nn.functional.grid_sample(uv_label, uv_sampler) tex_seg = tex_seg.view(tex_seg.size(0), -1, tex_seg.size(2), opts.tex_size, opts.tex_size).permute(0, 2, 3, 4, 1) tex_left = tex_seg[:, -self.model.texture_predictor.num_sym_faces:] tex_seg = torch.cat([tex_seg, tex_left], 1) tex_seg = tex_seg.view(tex_seg.size(0), tex_seg.size(1), -1, (self.opts.num_parts + 1)) tex_seg = torch.argmax(tex_seg, dim=-1) # obtain vertex label through face label tex_seg = self.most_freq(tex_seg.squeeze()) tex_seg = tex_seg.float() face = self.faces[0] parts = [] for cnt in range(opts.num_parts): parts.append([]) # go through all vertices and compute their label # for sanity check vert_tex = [] for cnt in range(face.max() + 1): v0 = (face[:, 0] == cnt) * 1 v1 = (face[:, 1] == cnt) * 1 v2 = (face[:, 2] == cnt) * 1 v = v0 + v1 + v2 # which faces relate to this vertex idxes = torch.nonzero(v).squeeze() labels = tex_seg[idxes].long().view(1, idxes.size(0)) label = self.most_freq(labels) if (label > 0): parts[label - 1].append(cnt) vert_tex.append(label) np.save(osp.join(opts.out_dir, "head_vertices.npy"), parts[0]) np.save(osp.join(opts.out_dir, "neck_vertices.npy"), parts[1]) np.save(osp.join(opts.out_dir, "back_vertices.npy"), parts[2]) np.save(osp.join(opts.out_dir, "belly_vertices.npy"), parts[3]) # visualize part label for each vertex vert_tex = torch.stack(vert_tex) vert_tex = self.colorize(vert_tex.view(642, 1).cpu().numpy()) vert_tex = torch.from_numpy(vert_tex).float().squeeze() vert_tex = vert_tex.permute(1, 0) mesh_ = sr.Mesh(mean_v, self.faces, vert_tex.view(1, 642, 3), texture_type='vertex') mesh_path = osp.join(opts.out_dir, "vertex_label.obj") mesh_.save_obj(mesh_path, save_texture=True) print( tf_visualizer.green("Semantic template saved at {}.".format( opts.out_dir)))
def main(): filename_input = os.path.join(data_dir, 'banana.obj') filename_output = os.path.join(output_directory, 'example1.gif') ########################### # camera settings ########################### camera_distance = 2 elevation = 30 ########################### # load object ########################### mesh = TriangleMesh.from_obj(filename_input) vertices = mesh.vertices faces = mesh.faces.int() uvs = torch.FloatTensor(get_spherical_coords_x(vertices.data.numpy())) face_textures = (faces).clone() vertices = vertices[None, :, :].cuda() faces = faces[None, :, :].cuda() uvs = uvs[None, :, :].cuda() face_textures[None, :, :].cuda() ########################### # normalize verts ########################### vertices_max = vertices.max() vertices_min = vertices.min() vertices_middle = (vertices_max + vertices_min) / 2. vertices = vertices - vertices_middle coef = 5 vertices = vertices * coef ########################### # NMR ########################### textures = torch.ones(1, faces.shape[1], 2, 2, 2, 3, dtype=torch.float32).cuda() renderer = nr.Renderer(camera_mode='look_at') loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing NMR') writer = imageio.get_writer(os.path.join(output_directory_nmr, 'rotation.gif'), mode='I') for num, azimuth in enumerate(loop): renderer.eye = nr.get_points_from_angles(camera_distance, elevation, azimuth) images, _, _ = renderer(vertices, faces, textures) image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) writer.close() ########################### # Soft Rasterizer ########################### textures = torch.ones(1, faces.shape[1], 2, 3, dtype=torch.float32).cuda() mesh = sr.Mesh(vertices, faces, textures) renderer = sr.SoftRenderer(camera_mode='look_at') loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing SR') writer = imageio.get_writer(os.path.join(output_directory_sr, 'rotation.gif'), mode='I') for azimuth in loop: mesh.reset_() renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) images = renderer.render_mesh(mesh) image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) writer.close() ################################ # Dib-Renderer - Vertex Colours ################################ renderer = Dib_Renderer(256, 256, mode='VertexColor') textures = torch.ones(1, vertices.shape[1], 3).cuda() loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing Dib_Renderer VertexColor') writer = imageio.get_writer(os.path.join(output_directory_dib, 'rotation_VertexColor.gif'), mode='I') for azimuth in loop: renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) predictions, _, _ = renderer.forward( points=[vertices, faces[0].long()], colors=[textures]) image = predictions.detach().cpu().numpy()[0] writer.append_data((image * 255).astype(np.uint8)) writer.close() ################################ # Dib-Renderer - Lambertian ################################ renderer = Dib_Renderer(256, 256, mode='Lambertian') textures = torch.ones(1, 3, 256, 256).cuda() loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing Dib_Renderer Lambertian') writer = imageio.get_writer(os.path.join(output_directory_dib, 'rotation_Lambertian.gif'), mode='I') for azimuth in loop: renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) predictions, _, _ = renderer.forward(points=[vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), textures]) image = predictions.detach().cpu().numpy()[0] writer.append_data((image * 255).astype(np.uint8)) writer.close() ################################ # Dib-Renderer - Phong ################################ renderer = Dib_Renderer(256, 256, mode='Phong') textures = torch.ones(1, 3, 256, 256).cuda() ### Lighting info ### material = np.array([[0.1, 0.1, 0.1], [1.0, 1.0, 1.0], [0.4, 0.4, 0.4]], dtype=np.float32).reshape(-1, 3, 3) material = torch.from_numpy(material).repeat(1, 1, 1).cuda() shininess = np.array([100], dtype=np.float32).reshape(-1, 1) shininess = torch.from_numpy(shininess).repeat(1, 1).cuda() lightdirect = 2 * np.random.rand(1, 3).astype(np.float32) - 1 lightdirect[:, 2] += 2 lightdirect = torch.from_numpy(lightdirect).cuda() loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing Dib_Renderer Phong') writer = imageio.get_writer(os.path.join(output_directory_dib, 'rotation_Phong.gif'), mode='I') for azimuth in loop: renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) predictions, _, _ = renderer.forward(points=[vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), textures],\ light= lightdirect, \ material=material, \ shininess=shininess ) image = predictions.detach().cpu().numpy()[0] writer.append_data((image * 255).astype(np.uint8)) writer.close() ################################ # Dib-Renderer - SH ################################ renderer = Dib_Renderer(256, 256, mode='SphericalHarmonics') textures = torch.ones(1, 3, 256, 256).cuda() ### Lighting info ### lightparam = np.random.rand(1, 9).astype(np.float32) lightparam[:, 0] += 2 lightparam = torch.from_numpy(lightparam).cuda() loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing Dib_Renderer SH') writer = imageio.get_writer(os.path.join(output_directory_dib, 'rotation_SH.gif'), mode='I') for azimuth in loop: renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) predictions, _, _ = renderer.forward(points=[vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), textures],\ light=lightparam) image = predictions.detach().cpu().numpy()[0] writer.append_data((image * 255).astype(np.uint8)) writer.close()
def generatert(): itvl = 1000.0 / 25.0 # 25fps # extract the frontal facial landmarks for key frame lmk3d_target = np.load("/home/cxu-serve/p1/common/demo/lisa2_original.npy") # load the 3D facial landmarks on the PRNet 3D reconstructed face lmk3d_origin = np.load( "/home/cxu-serve/p1/common/demo/lisa2_crop_prnet.npy") # calculate the affine transformation between PRNet 3D face and the frotal face landmarks lmk3d_origin_homo = np.hstack( (lmk3d_origin, np.ones([lmk3d_origin.shape[0], 1]))) # 68x4 p_affine = np.linalg.lstsq(lmk3d_origin_homo, lmk3d_target, rcond=1)[0].T # Affine matrix. 3 x 4 pr = p_affine[:, :3] # 3x3 pt = p_affine[:, 3:] # 3x1 # load the original 3D face mesh then transform it to align frontal face landmarks vertices_org, triangles, colors = load_obj( "/home/cxu-serve/p1/common/demo/lisa2_crop_original.obj" ) # get unfrontalized vertices position vertices_origin_affine = (pr @ (vertices_org.T) + pt).T # aligned vertices n_frames = 200 # number of RT you want rots = np.zeros((n_frames, 3, 3)) trans = np.zeros((n_frames, 3, 1)) t = vertices_origin_affine.mean(axis=0).reshape(3, 1) for i in range(n_frames): rots[i, :, :] = R.from_euler( "xyz", [0, np.pi * 0.5 - np.pi / n_frames * i, 0]).as_dcm() trans[i, :, :] = t - rots[i, :, :] @ t rt = rots.copy() for j in range(200): ret_R = rt[j] r = Rotation.from_dcm(ret_R) vec = r.as_rotvec() RTs[j, :3] = vec RTs[j, 3:] = np.squeeze(np.asarray(ret_t)) # save the RT now np.save("/home/cxu-serve/p1/common/demo/lisa2_rt.npy", rots) ani_path = "/home/cxu-serve/p1/common/demo/lisa2_crop_ani.mp4" # lele: XXXXXX # set up the renderer renderer = setup_renderer() fig = plt.figure() ims = [] s for i in range(rots.shape[0]): # get rendered frame vertices = (rots[i].T @ (vertices_origin_affine.T - trans[i])).T # vertices = vertices_origin_affine face_mesh = sr.Mesh(vertices, triangles, colors, texture_type="vertex") image_render = get_np_uint8_image( face_mesh, renderer) # RGBA, (224,224,3), np.uint8 # im = plt.imshow(image_render, animated=True) # ims.append([im]) # print("[{}/{}]".format(i+1, rots.shape[0])) #save rgba image as bgr in cv2 rgb_frame = (image_render).astype(int)[:, :, :-1][..., ::-1] cv2.imwrite(temp_path + "/%05d.png" % i, rgb_frame) command = 'ffmpeg -framerate 25 -i ' + temp_path + '/%5d.png -c:v libx264 -y -vf format=yuv420p ' + ani_path os.system(command)
expression_id = "01" video_path = source_path + '01-01-%s-01-01-01-01.mp4' % expression_id mesh_path = source_path + "3dmesh" mesh_id = "01/01-01-%s-01-01-01-01" % expression_id renderer = setup_renderer() if overlay: real_video = mmcv.VideoReader(video_path) fig = plt.figure() ims = [] num_frames = count_frames(video_path) for i in range(num_frames): vertices, triangles, colors = load_obj("{}/{}_%05d_original.obj".format(mesh_path, mesh_id) % i) face_mesh = sr.Mesh(vertices, triangles, colors, texture_type='vertex') image_render = get_np_uint8_image(face_mesh, renderer) if overlay: frame = mmcv.bgr2rgb(real_video[i]) # RGB, (224,224,3), np.uint8 if not overlay: im = plt.imshow(image_render, animated=True) else: im = plt.imshow((frame[:, :, :3] * 0.5 + image_render[:, :, :3] * 0.5).astype(np.uint8), animated=True) ims.append([im]) ani = animation.ArtistAnimation(fig, ims, interval=itvl, blit=True, repeat_delay=1000) if not overlay: ani.save('{}/{}_render.mp4'.format(mesh_path, mesh_id))
def get_current_visuals(self): vis_dict = {} nb, nf, _, nc = self.tex.size() tex = self.tex.detach().view(nb, nf, opts.tex_size, opts.tex_size, nc).unsqueeze(4).repeat( 1, 1, 1, 1, opts.tex_size, 1) vis_dict['image_pred'] = self.vis_renderer(self.pred_vs.detach(), self.faces, self.proj_cam.detach(), tex) vis_dict['mask_pred'] = self.vis_renderer( self.pred_vs.detach(), self.faces, self.proj_cam.detach()).unsqueeze(1) uv_flows = self.uvimage_pred.detach() uv_flows = uv_flows.permute(0, 2, 3, 1) uv_images = torch.nn.functional.grid_sample(self.imgs, uv_flows) vis_dict['uv_images'] = uv_images vis_dict['mask_gt'] = self.masks.unsqueeze(1) vis_dict['image_gt'] = self.imgs for cnt in range(5): mesh_ = sr.Mesh( self.pred_vs[cnt].detach(), self.faces[cnt], self.tex[cnt].view(self.faces.size(1), -1, 3).detach()) vis_dict['mesh_' + str(cnt)] = mesh_ # visualize part projects for cnt in range(len(self.part_projs)): vis_dict['part_render_' + str(cnt)] = self.part_projs[cnt] # visualize gt part projects for cnt in range(4): vis_dict['gt_part_render_' + str(cnt)] = self.part_segs[:, cnt + 1].unsqueeze(1) mean_shape = self.mean_shape.unsqueeze(0).repeat( self.pred_vs.size(0), 1, 1) _, vert2ds = self.corr_loss_fn(self.head_points, self.belly_points, self.back_points, self.neck_points, mean_shape, self.proj_cam.detach()) vert2ds = (vert2ds + 1.0) / 2 * opts.image_size self.center_imgs = torch.zeros(self.imgs.size()) for cnt in range(self.imgs.size(0)): img_ = self.imgs[cnt] img_ = img_.cpu().permute(1, 2, 0).numpy() img_ = img_ * 255.0 img_ = np.array(img_, dtype=np.uint8).copy() vert2d = vert2ds[cnt] for ccnt in range(vert2d.size(0)): x = int(vert2d[ccnt, 0].cpu().detach().numpy()) y = int(vert2d[ccnt, 1].cpu().detach().numpy()) head_num = self.corr_loss_fn.head_num belly_num = self.corr_loss_fn.belly_num # we only visualize points from the head and belly parts # since they're most informative. if (ccnt < head_num): cv2.drawMarker(img_, (x, y), (0, 0, 255), markerType=cv2.MARKER_CROSS, markerSize=5, thickness=2) elif (ccnt < head_num + belly_num): cv2.drawMarker(img_, (x, y), (0, 255, 0), markerType=cv2.MARKER_CROSS, markerSize=5, thickness=2) head_points = (self.head_points[cnt] + 1) / 2.0 * 256.0 belly_points = (self.belly_points[cnt] + 1) / 2.0 * 256.0 for ccnt in range(head_points.size(0)): x = int(head_points[ccnt, 0].cpu().numpy()) y = int(head_points[ccnt, 1].cpu().numpy()) cv2.drawMarker(img_, (x, y), (0, 255, 255), markerType=cv2.MARKER_CROSS, markerSize=5, thickness=2) for ccnt in range(belly_points.size(0)): x = int(belly_points[ccnt, 0].cpu().numpy()) y = int(belly_points[ccnt, 1].cpu().numpy()) cv2.drawMarker(img_, (x, y), (255, 255, 0), markerType=cv2.MARKER_CROSS, markerSize=5, thickness=2) self.center_imgs[cnt] = torch.from_numpy(img_ / 255.0).permute( 2, 0, 1).unsqueeze(0).float() vis_dict['part_cons_vis'] = self.center_imgs return vis_dict
def run_on_image(self, image, save_name): """ Args: image (np.ndarray): an image of shape (H, W, C) (in BGR order). This is the format used by OpenCV. Returns: predictions (dict): the output of the model. vis_output (VisImage): the visualized image output. """ vis_output = None predictions = self.predictor(image) # Convert image from OpenCV BGR format to Matplotlib RGB format. image = image[:, :, ::-1] mask = predictions['instances'].raw_masks.squeeze(1).data.cpu().numpy( ) if predictions['instances'].has("raw_masks") else None visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode) if "panoptic_seg" in predictions: panoptic_seg, segments_info = predictions["panoptic_seg"] vis_output = visualizer.draw_panoptic_seg_predictions( panoptic_seg.to(self.cpu_device), segments_info) else: if "sem_seg" in predictions: vis_output = visualizer.draw_sem_seg( predictions["sem_seg"].argmax(dim=0).to(self.cpu_device)) if "instances" in predictions: instances = predictions["instances"].to(self.cpu_device) pred_classes = torch.ones(instances.pred_classes.shape) # uncomment to open nms between different classes ''' res = batched_nms(instances.pred_boxes.tensor, instances.scores, pred_classes, 0.5) print('res:', res) print('res:', res.size()[0]) #instances.num_instances = res.size()[0] instances.pred_boxes.tensor = instances.pred_boxes.tensor[res] instances.pred_classes = instances.pred_classes[res] instances.scores = instances.scores[res] instances.pred_keypoints = instances.pred_keypoints[res] instances.predict_trans = instances.predict_trans[res] instances.predict_rotation = instances.predict_rotation[res] instances.predict_vertices = instances.predict_vertices[res] print('pred trans shape:', instances.predict_trans.shape) ''' vis_output = visualizer.draw_instance_predictions( predictions=instances) output_trans_dir = './inference_val_translation/' output_rotation_dir = './inference_val_rotation/' output_mesh_dir = './inference_val_mesh/' output_cls_dir = './inference_val_cls/' output_score_dir = './inference_val_score/' save_name = save_name.split('/')[1] template_path = './merge_mean_car_shape/' faces = sr.Mesh.from_obj(template_path + 'merge_mean_car_model_0.obj').faces for directory in [ output_trans_dir, output_rotation_dir, output_mesh_dir, output_cls_dir, output_score_dir ]: if not os.path.exists(directory): os.makedirs(directory) for index in range(instances.predict_trans.shape[0]): with open( os.path.join( output_trans_dir, save_name + '_' + str(index) + '.json'), 'w') as f: data = {} data['translation'] = list( instances.predict_trans[index].cpu().detach( ).numpy().astype(float)) json.dump(data, f) for index in range(instances.predict_rotation.shape[0]): with open( os.path.join( output_rotation_dir, save_name + '_' + str(index) + '.json'), 'w') as f: data = {} data['rotation'] = list( instances.predict_rotation[index].cpu().detach( ).numpy().astype(float)) json.dump(data, f) for index in range(instances.pred_classes.shape[0]): with open( os.path.join( output_cls_dir, save_name + '_' + str(index) + '.json'), 'w') as f: data = {} data['car_id'] = int(instances.pred_classes[index].cpu( ).detach().numpy().astype(float)) json.dump(data, f) for index in range(instances.scores.shape[0]): with open( os.path.join( output_score_dir, save_name + '_' + str(index) + '.json'), 'w') as f: data = {} data['score'] = float(instances.scores[index].cpu(). detach().numpy().astype(float)) json.dump(data, f) for index in range(instances.predict_vertices.shape[0]): vertices = instances.predict_vertices[index].unsqueeze(0) sr.Mesh(vertices, faces).save_obj(os.path.join( output_mesh_dir, save_name + '_' + str(index) + '.obj'), save_texture=False) return predictions, vis_output