def visualize_row(self, pipeline, batch, batch_idx, deg_range, axis=1, return_format=None):
if not isinstance(axis, list):
axis = [axis]
result = list()
idx = 0
if 1 in axis:
deg_range = -1 * deg_range
for theta_d in deg_range:
batch2 = batch.copy()
batch_size = batch2[DIK.R].shape[0]
theta = [0, 0, 0]
for i in range(len(axis)):
theta[axis[i]] = theta_d
R = theta2rotation_matrix(theta_all=theta).to(batch2[DIK.R].device).unsqueeze(0)
t = self.get_neutral_t(batch)
uv = self.bfm_uv_factory.getUV(R=R, t=t, s=batch2[DIK.SCALE], sp=batch2[DIK.COEFF_SHAPE],
ep=batch2[DIK.COEFF_EXPRESSION])
batch2[DIK.UV_RENDERED] = uv.expand_as(batch[DIK.UV_RENDERED]).to(batch[DIK.UV_RENDERED].device)
img_out = self.run(pipeline, batch2, idx)
result.append(img_out)
nrow = len(result) if 1 in axis else 1
vis = make_grid(result, nrow=nrow)
return self.format_output(vis, return_format=return_format)
def visualize_single(self, pipeline, batch, batch_idx, theta_all=None, R=None, forward_type='style2image'):
batch2 = self.detach(batch)
if R is None:
R = theta2rotation_matrix(theta_all=theta_all)
uv = self.bfm_uv_factory.getUV(sp=batch[DIK.COEFF_SHAPE], ep=batch[DIK.COEFF_EXPRESSION], R=R, t=batch[DIK.T])
batch2[DIK.UV_RENDERED] = uv.expand_as(batch2[DIK.UV_RENDERED]).to(batch2[DIK.UV_RENDERED].device)
out = self.run(pipeline, batch2, batch_idx, return_format='pil', forward_type=forward_type)
return out
def inference_posed_ffhq(opt, results_folder, n=3000):
"""
Runs different poses for each latent code extracted from the FFHQ dataset.
Uses shape and expression parameters from validation set.
"""
print("Running inference on FFHQ. Using the extracted face model parameters, various poses, and predicted latent codes. This can be slow.")
mkdir(results_folder)
visualizer = Visualizer(opt, return_format='pil')
dataset = NPYDataset(opt, augmentation=False, split='val')
dataloader = iter(DataLoader(dataset, batch_size=1, num_workers=0, shuffle=False))
model = get_model(opt)
visualizer.mask_key = DIK.SEGMENTATION_PREDICTED
n = min(n, 3000)
all_pose_axis = [[0], [1], [0, 1]]
uv_factory_bfm = BFMUVFactory(opt, use_bfm_gpu=True)
s = torch.Tensor([0.1])
t = torch.Tensor([[0, -2, -57]])
all_poses = np.arange(-45, 46, 15)
for i, batch_orig in tqdm(enumerate(dataloader)):
batch_orig = model.forward(batch_orig, i, std_multiplier=0)
file_id = batch_orig[DIK.FILENAME][0]
sp = batch_orig[DIK.COEFF_SHAPE].clone()
ep = batch_orig[DIK.COEFF_EXPRESSION].clone()
latent = batch_orig[DIK.STYLE_LATENT].clone()
for pose_axis in all_pose_axis:
for pose in all_poses:
theta = torch.Tensor([0, 0, 0])
theta[pose_axis] = pose
R = theta2rotation_matrix(theta_all=theta).unsqueeze(0)
uv_tensor = uv_factory_bfm.getUV(R, t, s, sp, ep, correct_translation=True)
batch = {
DIK.UV_RENDERED: uv_tensor,
DIK.R: R,
DIK.T: t,
DIK.SCALE: s,
DIK.STYLE_LATENT: latent
}
batch = {k: v.cuda() for k, v in batch.items()}
batch_out = model.forward_latent2image(batch, 0)
out = visualizer.tensor2image(batch_out[DIK.IMAGE_OUT][0], batch=batch_out, white_bg=False)
out = visualizer.format_output(out, return_format='pil')
folder_out = os.path.join(results_folder, "axis{}_theta{}".format(pose_axis, pose))
mkdir(folder_out)
out.save(os.path.join(results_folder, folder_out, "{}.png".format(file_id)))
if i >= n:
break
print("Done")
def get_posed_uvs(self, sp=torch.zeros((199,)), ep=torch.zeros((100,)), deg_range=torch.arange(-30, 31, 15),
t=torch.Tensor((0, 0, -35)), s=torch.Tensor((0.1,))):
uv_list = list()
for theta_y in deg_range:
for theta_x in deg_range:
theta = theta_x, theta_y, 0
R = theta2rotation_matrix(theta_all=theta)
uv = self.getUV(sp, ep, R, t, s)
uv_list.append(uv)
return uv_list
def inference_posed(opt, results_folder, path_latent, n=3000):
"""
Runs different poses for each latent code extracted from the FFHQ dataset.
Uses random shape and expression parameters.
"""
print("Running inference on FFHQ. Using random face model parameters, various poses, and predicted latent codes. This can be slow.")
visualizer = Visualizer(opt, return_format='pil')
visualizer.mask_key = DIK.SEGMENTATION_PREDICTED
# We use the latent codes extracted from FFHQ, but they could also be sampled from the extracted distributions.
latents = np.load(path_latent)[:, 0] # We want to use the distribution means directly (index 0)
mkdir(results_folder)
model = get_model(opt)
all_pose_axis = [[0], [1], [0, 1]] # List of all axis: 0 is pitch, 1 is yaw
uv_factory_bfm = BFMUVFactory(opt, use_bfm_gpu=True)
s = torch.Tensor([0.1])
t = torch.Tensor([[0, -2, -57]])
results = list()
all_poses = np.arange(-45, 46, 15)
with torch.no_grad():
for i in tqdm(range(n)):
sp = uv_factory_bfm.bfm.sample_shape(1)
ep = uv_factory_bfm.bfm.sample_expression(1)
latent = torch.Tensor([latents[i]])
for pose_axis in all_pose_axis:
for pose in all_poses:
theta = torch.Tensor([0, 0, 0])
theta[pose_axis] = pose
R = theta2rotation_matrix(theta_all=theta).unsqueeze(0)
uv_tensor = uv_factory_bfm.getUV(R, t, s, sp, ep, correct_translation=True)
batch = {
DIK.UV_RENDERED: uv_tensor,
DIK.R: R,
DIK.T: t,
DIK.SCALE: s,
DIK.STYLE_LATENT: latent
}
batch = {k: v.cuda() for k, v in batch.items()}
batch_out = model.forward_latent2image(batch, 0)
out = visualizer.tensor2image(batch_out[DIK.IMAGE_OUT][0], batch=batch_out, white_bg=False)
results.append(out)
out = visualizer.format_output(out, return_format='pil')
folder_out = os.path.join(results_folder, "axis{}_theta{}".format(pose_axis, pose))
mkdir(folder_out)
out.save(os.path.join(results_folder, folder_out, "{:05d}.png".format(i)))
def visualize_single(self, pipeline, batch, batch_idx, theta_all=None, R=None, forward_type='style2image',
correct_translation=True):
batch2 = batch
if R is None:
R = theta2rotation_matrix(theta_all=theta_all).to(self.device)
batch2[DIK.R] = R.unsqueeze(0)
uv = self.bfm_uv_factory.getUV(sp=batch[DIK.COEFF_SHAPE], ep=batch[DIK.COEFF_EXPRESSION], R=R, t=batch[DIK.T],
correct_translation=correct_translation)
batch2[DIK.UV_RENDERED] = uv.expand(batch2[DIK.COEFF_SHAPE].shape[0], -1, -1, -1).to(self.device)
batch2 = self.run(pipeline, batch2, batch_idx, return_format='pil', forward_type=forward_type)
return batch2
def visualize_grid(self, pipeline, batch, batch_idx, deg_range, return_format=None):
result = list()
idx = 0
for theta_y in deg_range:
for theta_x in deg_range:
batch2 = batch.copy()
theta = [theta_x, theta_y, 0]
R = theta2rotation_matrix(theta_all=theta).to(batch2[DIK.R].device).unsqueeze(0)
t = self.get_neutral_t(batch)
uv = self.bfm_uv_factory.getUV(R=R, t=t, s=batch2[DIK.SCALE], sp=batch2[DIK.COEFF_SHAPE],
ep=batch2[DIK.COEFF_EXPRESSION])
batch2[DIK.UV_RENDERED] = uv.expand_as(batch[DIK.UV_RENDERED]).to(batch[DIK.UV_RENDERED].device)
img_out = self.run(pipeline, batch2, idx)
result.append(img_out)
vis = make_grid(result, nrow=int(np.sqrt(len(result))))
return self.format_output(vis, return_format=return_format)
def visualize_show(self, pipeline, batch, batch_idx, n_y, n_x, pose_range_x=30, pose_range_y=30):
theta_x = torch.linspace(-pose_range_x, pose_range_x, n_x) if pose_range_x > 0 else [0]
theta_y = torch.linspace(-pose_range_y, pose_range_y, n_y) if pose_range_y > 0 else [0]
top = [[theta_x[0], theta_y[i], 0] for i in range(n_y - 1)] if n_y > 0 else []
down = [[theta_x[i], theta_y[-1], 0] for i in range(n_x - 1)] if n_x > 0 else []
bottom = [[theta_x[-1], theta_y[-i - 1], 0] for i in range(n_y - 1)] if n_y > 0 else []
up = [[theta_x[-i - 1], theta_y[0], 0] for i in range(n_x - 1)] if n_x > 0 else []
all = top + down + bottom + up
all = [theta2rotation_matrix(theta_all=theta) for theta in all]
result = list()
for i, R in enumerate(all):
batch2 = batch.copy()
t = batch2[DIK.T]
uv = self.bfm_uv_factory.getUV(sp=batch[DIK.COEFF_SHAPE], ep=batch[DIK.COEFF_EXPRESSION], R=R, t=t)
batch2[DIK.UV_RENDERED] = uv.expand_as(batch2[DIK.UV_RENDERED]).to(batch2[DIK.UV_RENDERED].device)
img_out = self.run(pipeline, batch2, batch_idx, return_format='np')
result.append(img_out)
return result