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
