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)