From 00e74125695ff43a591298f68ed49ae4225dd8d1 Mon Sep 17 00:00:00 2001 From: NNNNAI <844294823@qq.com> Date: Mon, 19 Jul 2021 21:01:03 +0800 Subject: [PATCH] Update reverse2original.py --- util/reverse2original.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/util/reverse2original.py b/util/reverse2original.py index cfa890e..0c5d29d 100644 --- a/util/reverse2original.py +++ b/util/reverse2original.py @@ -23,7 +23,7 @@ def encode_segmentation_rgb(segmentation, no_neck=True): mouth_map[valid_index] = 255 # valid_index = np.where(parse==hair_id) # hair_map[valid_index] = 255 - + #return np.stack([face_map, mouth_map,hair_map], axis=2) return np.stack([face_map, mouth_map], axis=2) @@ -108,7 +108,7 @@ 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) - if tgt_mask.sum() != 0: + if tgt_mask.sum() >= 5000: # face_mask_tensor = tgt_mask[...,0] + tgt_mask[...,1] target_mask = cv2.resize(tgt_mask, (224, 224)) # print(source_img)