def inference(self, tgt_paths=None, tgt_smpls=None, cam_strategy="smooth",
output_dir="", use_selected_f2pts=False, visualizer=None, verbose=True):
assert (tgt_paths is not None) or (tgt_smpls is not None)
if tgt_paths is not None:
length = len(tgt_paths)
else:
length = len(tgt_smpls)
outputs = []
process_bar = tqdm(range(length)) if verbose else range(length)
self.first_cam = None
tgt_smpls = torch.tensor(tgt_smpls).float().cuda()
tgt_smpls = self.weak_cam_swapper.stabilize(tgt_smpls)
for t in process_bar:
tgt_smpl = tgt_smpls[t:t+1]
input_G_tsf, Tst, Ttt, temp_enc_outs, temp_res_outs, ref_info = self.make_inputs_for_tsf(
self.src_info, tgt_smpl, cam_strategy, t, use_selected_f2pts=use_selected_f2pts
)
preds, tsf_mask = self.forward(input_G_tsf[:, 0], Tst, temp_enc_outs, temp_res_outs, Ttt)
if t != 0 and self._opt.temporal:
prev_preds = self.temporal_fifo.temporal_preds_tensor()
preds_warp = F.grid_sample(prev_preds, Ttt.view(-1, self._opt.image_size, self._opt.image_size, 2))
else:
preds_warp = None
if visualizer is not None:
src_warp = F.grid_sample(self.src_info["img"].view(-1, 3, self._opt.image_size, self._opt.image_size),
Tst.view(-1, self._opt.image_size, self._opt.image_size, 2))
visualizer.vis_named_img("pred_" + cam_strategy, preds)
visualizer.vis_named_img("uv_warp", input_G_tsf[0, :, 0:3])
visualizer.vis_named_img("src_warp", src_warp)
if preds_warp is not None:
visualizer.vis_named_img("preds_warp", preds_warp)
if self._opt.temporal:
self.post_update(ref_info, preds)
if output_dir:
filename = "{:0>8}.png".format(t)
preds = preds[0].cpu().numpy()
file_path = os.path.join(output_dir, "pred_" + filename)
cv_utils.save_cv2_img(preds, file_path, normalize=True)
outputs.append(file_path)
else:
preds = preds[0].cpu().numpy()
outputs.append(preds)
# tsf_mask = tsf_mask[0, 0].cpu().numpy() * 255
# tsf_mask = tsf_mask.astype(np.uint8)
# cv_utils.save_cv2_img(tsf_mask, os.path.join(output_dir, "mask_" + filename), normalize=False)
return outputs
def inference(self,
tgt_smpls,
cam_strategy="smooth",
output_dir="",
visualizer=None,
verbose=True):
length = len(tgt_smpls)
outputs = []
process_bar = tqdm(range(length)) if verbose else range(length)
for t in process_bar:
tgt_smpl = torch.tensor(tgt_smpls[t]).float()[None].to(self.device)
input_G_tsf, Tst, Ttt, temp_enc_outs, temp_res_outs, ref_info = self.make_inputs_for_tsf(
self.src_info, tgt_smpl, cam_strategy, t)
preds, tsf_mask = self.forward(input_G_tsf[:, 0], Tst,
temp_enc_outs, temp_res_outs, Ttt)
if t != 0 and self._opt.temporal:
prev_preds = self.temporal_fifo.temporal_preds_tensor()
preds_warp = F.grid_sample(
prev_preds,
Ttt.view(-1, self._opt.image_size, self._opt.image_size,
2))
else:
preds_warp = None
if visualizer is not None:
src_warp = F.grid_sample(
self.src_info["img"].view(-1, 3, self._opt.image_size,
self._opt.image_size),
Tst.view(-1, self._opt.image_size, self._opt.image_size,
2))
visualizer.vis_named_img("pred_" + cam_strategy, preds)
visualizer.vis_named_img("uv_warp", input_G_tsf[0, :, 0:3])
visualizer.vis_named_img("src_warp", src_warp)
if preds_warp is not None:
visualizer.vis_named_img("preds_warp", preds_warp)
if self._opt.temporal:
self.post_update(ref_info, preds)
if output_dir:
filename = "pred_{:0>8}.png".format(t)
preds = preds[0].cpu().numpy()
file_path = os.path.join(output_dir, filename)
cv_utils.save_cv2_img(preds, file_path, normalize=True)
outputs.append(file_path)
return outputs
def display_current_results(self,
visuals,
it,
is_train,
save_visuals=False):
for label, image_numpy in visuals.items():
# ipdb.set_trace()
sum_name = "{}/{}".format("Train" if is_train else "Test", label)
self._writer.add_image(sum_name, image_numpy, it)
if save_visuals:
save_cv2_img(
image_numpy,
os.path.join(self._opt.checkpoints_dir, self._opt.name,
"event_imgs", sum_name, "%08d.png" % it))
self._writer.export_scalars_to_json(self._tb_path)
def render_mask(smpl, render, cams, poses, shapes, offsets, img_dir,
image_names, parse_out_dir, visual_path):
"""
Args:
smpl (SMPL):
render (SMPLRenderer):
cams:
poses:
shapes:
offsets:
img_dir:
image_names:
parse_out_dir:
visual_path:
Returns:
"""
global IMAGE_SIZE
length = cams.shape[0]
device = torch.device("cuda:0")
offsets = torch.tensor(offsets).float().to(device)
textures = render.color_textures().to(device)[None]
fourcc = cv2.VideoWriter_fourcc(*"XVID")
videoWriter = cv2.VideoWriter(visual_path, fourcc, 25,
(IMAGE_SIZE, IMAGE_SIZE))
print(f"Preprocessing {img_dir}")
for i in tqdm(range(length)):
img_name = image_names[i]
name = img_name.split(".")[0]
image = read_cv2_img(os.path.join(img_dir, img_name))
cam = torch.from_numpy(cams[i:i + 1]).to(device)
shape = torch.from_numpy(shapes[i:i + 1]).to(device)
pose = torch.from_numpy(poses[i:i + 1]).to(device)
verts, _, _ = smpl(shape, pose, offsets=offsets, get_skin=True)
rd_imgs, _ = render.render(cam, verts, textures)
mask = render.render_silhouettes(cam, verts)[0]
mask.unsqueeze_(-1)
# (h, w, 1)
mask = mask.cpu().numpy()
# (3, h, w)
rd_imgs = rd_imgs.cpu().numpy()[0]
# (h, w, 3)
rd_imgs = np.transpose(rd_imgs, (1, 2, 0))
rd_imgs = (rd_imgs + 1) / 2 * 255
rd_imgs = rd_imgs.astype(np.uint8)
overly_img = image * (1 - mask) + rd_imgs * mask
parse_path = os.path.join(parse_out_dir, name + "_alpha.png")
save_cv2_img(mask[:, :, 0] * 255, parse_path, transpose=False)
overly_img = overly_img.astype(np.uint8)
videoWriter.write(overly_img)
videoWriter.release()
def save_images(self, visuals):
for label, image_numpy in visuals.items():
image_name = "%s.png" % label
save_path = os.path.join(self._save_path, "samples", image_name)
save_cv2_img(image_numpy, save_path)