From a2c3e958d26f14f8d26fabac9e625f7d6a251e22 Mon Sep 17 00:00:00 2001 From: NNNNAI <844294823@qq.com> Date: Mon, 19 Jul 2021 14:57:13 +0800 Subject: [PATCH] Update reverse2original.py --- util/reverse2original.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/util/reverse2original.py b/util/reverse2original.py index fdb729d..25f06b5 100644 --- a/util/reverse2original.py +++ b/util/reverse2original.py @@ -108,16 +108,17 @@ def reverse2wholeimage(b_align_crop_tenor_list,swaped_imgs, mats, crop_size, ori parsing = out.squeeze(0).detach().cpu().numpy().argmax(0) vis_parsing_anno = parsing.copy().astype(np.uint8) tgt_mask = encode_segmentation_rgb(vis_parsing_anno) - # face_mask_tensor = tgt_mask[...,0] + tgt_mask[...,1] - target_mask = cv2.resize(tgt_mask, (224, 224)) - - # print(source_img) - target_image_parsing = postprocess(swaped_img, source_img[0].cpu().detach().numpy().transpose((1, 2, 0)), target_mask,smooth_mask) - - - target_image_parsing = cv2.warpAffine(target_image_parsing, mat_rev, orisize) - # target_image_parsing = cv2.warpAffine(swaped_img, mat_rev, orisize) + if tgt_mask.sum() != 0: + # face_mask_tensor = tgt_mask[...,0] + tgt_mask[...,1] + target_mask = cv2.resize(tgt_mask, (224, 224)) + # print(source_img) + target_image_parsing = postprocess(swaped_img, source_img[0].cpu().detach().numpy().transpose((1, 2, 0)), target_mask,smooth_mask) + + target_image = cv2.warpAffine(target_image_parsing, mat_rev, orisize) + # target_image_parsing = cv2.warpAffine(swaped_img, mat_rev, orisize) + else: + target_image = cv2.warpAffine(swaped_img, mat_rev, orisize)[..., ::-1] else: target_image = cv2.warpAffine(swaped_img, mat_rev, orisize) # source_image = cv2.warpAffine(source_img, mat_rev, orisize) @@ -153,7 +154,7 @@ def reverse2wholeimage(b_align_crop_tenor_list,swaped_imgs, mats, crop_size, ori # target_image_parsing = postprocess(target_image, source_image, tgt_mask) if use_mask: - target_image = np.array(target_image_parsing, dtype=np.float) * 255 + target_image = np.array(target_image, dtype=np.float) * 255 else: target_image = np.array(target_image, dtype=np.float)[..., ::-1] * 255