Update reverse2original.py

2021-07-19 21:01:03 +08:00
parent 0b4e6deecb
commit 00e7412569
1 changed files with 2 additions and 2 deletions
@@ -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)