|
@@ -1,13 +1,91 @@
|
|
|
import cv2
|
|
|
import numpy as np
|
|
|
# import time
|
|
|
-from util.add_watermark import watermark_image
|
|
|
+import torch
|
|
|
+from torch.nn import functional as F
|
|
|
+import torch.nn as nn
|
|
|
|
|
|
-def reverse2wholeimage(swaped_imgs, mats, crop_size, oriimg, logoclass, save_path = '', no_simswaplogo = False):
|
|
|
+
|
|
|
+def encode_segmentation_rgb(segmentation, no_neck=True):
|
|
|
+ parse = segmentation
|
|
|
+
|
|
|
+ face_part_ids = [1, 2, 3, 4, 5, 6, 10, 12, 13] if no_neck else [1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 13, 14]
|
|
|
+ mouth_id = 11
|
|
|
+ hair_id = 17
|
|
|
+ face_map = np.zeros([parse.shape[0], parse.shape[1]])
|
|
|
+ mouth_map = np.zeros([parse.shape[0], parse.shape[1]])
|
|
|
+ hair_map = np.zeros([parse.shape[0], parse.shape[1]])
|
|
|
+
|
|
|
+ for valid_id in face_part_ids:
|
|
|
+ valid_index = np.where(parse==valid_id)
|
|
|
+ face_map[valid_index] = 255
|
|
|
+ valid_index = np.where(parse==mouth_id)
|
|
|
+ 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)
|
|
|
+
|
|
|
+
|
|
|
+class SoftErosion(nn.Module):
|
|
|
+ def __init__(self, kernel_size=15, threshold=0.6, iterations=1):
|
|
|
+ super(SoftErosion, self).__init__()
|
|
|
+ r = kernel_size // 2
|
|
|
+ self.padding = r
|
|
|
+ self.iterations = iterations
|
|
|
+ self.threshold = threshold
|
|
|
+
|
|
|
+ # Create kernel
|
|
|
+ y_indices, x_indices = torch.meshgrid(torch.arange(0., kernel_size), torch.arange(0., kernel_size))
|
|
|
+ dist = torch.sqrt((x_indices - r) ** 2 + (y_indices - r) ** 2)
|
|
|
+ kernel = dist.max() - dist
|
|
|
+ kernel /= kernel.sum()
|
|
|
+ kernel = kernel.view(1, 1, *kernel.shape)
|
|
|
+ self.register_buffer('weight', kernel)
|
|
|
+
|
|
|
+ def forward(self, x):
|
|
|
+ x = x.float()
|
|
|
+ for i in range(self.iterations - 1):
|
|
|
+ x = torch.min(x, F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding))
|
|
|
+ x = F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding)
|
|
|
+
|
|
|
+ mask = x >= self.threshold
|
|
|
+ x[mask] = 1.0
|
|
|
+ x[~mask] /= x[~mask].max()
|
|
|
+
|
|
|
+ return x, mask
|
|
|
+
|
|
|
+
|
|
|
+def postprocess(swapped_face, target, target_mask,smooth_mask):
|
|
|
+ # target_mask = cv2.resize(target_mask, (self.size, self.size))
|
|
|
+
|
|
|
+ mask_tensor = torch.from_numpy(target_mask.copy().transpose((2, 0, 1))).float().mul_(1/255.0).cuda()
|
|
|
+ face_mask_tensor = mask_tensor[0] + mask_tensor[1]
|
|
|
+
|
|
|
+ soft_face_mask_tensor, _ = smooth_mask(face_mask_tensor.unsqueeze_(0).unsqueeze_(0))
|
|
|
+ soft_face_mask_tensor.squeeze_()
|
|
|
+
|
|
|
+ soft_face_mask = soft_face_mask_tensor.cpu().numpy()
|
|
|
+ soft_face_mask = soft_face_mask[:, :, np.newaxis]
|
|
|
+
|
|
|
+ result = swapped_face * soft_face_mask + target * (1 - soft_face_mask)
|
|
|
+ result = result[:,:,::-1]# .astype(np.uint8)
|
|
|
+ return result
|
|
|
+
|
|
|
+def reverse2wholeimage(b_align_crop_tenor_list,swaped_imgs, mats, crop_size, oriimg, logoclass, save_path = '', \
|
|
|
+ no_simswaplogo = False,pasring_model =None,norm = None, use_mask = False):
|
|
|
|
|
|
target_image_list = []
|
|
|
img_mask_list = []
|
|
|
- for swaped_img, mat in zip(swaped_imgs, mats):
|
|
|
+ if use_mask:
|
|
|
+ smooth_mask = SoftErosion(kernel_size=17, threshold=0.9, iterations=7).cuda()
|
|
|
+ else:
|
|
|
+ pass
|
|
|
+
|
|
|
+ # print(len(swaped_imgs))
|
|
|
+ # print(mats)
|
|
|
+ # print(len(b_align_crop_tenor_list))
|
|
|
+ for swaped_img, mat ,source_img in zip(swaped_imgs, mats,b_align_crop_tenor_list):
|
|
|
swaped_img = swaped_img.cpu().detach().numpy().transpose((1, 2, 0))
|
|
|
img_white = np.full((crop_size,crop_size), 255, dtype=float)
|
|
|
|
|
@@ -23,7 +101,27 @@ def reverse2wholeimage(swaped_imgs, mats, crop_size, oriimg, logoclass, save_pat
|
|
|
mat_rev[1][2] = -(mat[0][2]*mat[1][0]-mat[0][0]*mat[1][2])/div2
|
|
|
|
|
|
orisize = (oriimg.shape[1], oriimg.shape[0])
|
|
|
- target_image = cv2.warpAffine(swaped_img, mat_rev, orisize)
|
|
|
+ if use_mask:
|
|
|
+ source_img_norm = norm(source_img)
|
|
|
+ source_img_512 = F.interpolate(source_img_norm,size=(512,512))
|
|
|
+ out = pasring_model(source_img_512)[0]
|
|
|
+ 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)
|
|
|
+
|
|
|
+ else:
|
|
|
+ target_image = cv2.warpAffine(swaped_img, mat_rev, orisize)
|
|
|
+ # source_image = cv2.warpAffine(source_img, mat_rev, orisize)
|
|
|
+
|
|
|
img_white = cv2.warpAffine(img_white, mat_rev, orisize)
|
|
|
|
|
|
|
|
@@ -31,16 +129,39 @@ def reverse2wholeimage(swaped_imgs, mats, crop_size, oriimg, logoclass, save_pat
|
|
|
|
|
|
img_mask = img_white
|
|
|
|
|
|
- kernel = np.ones((10,10),np.uint8)
|
|
|
- img_mask = cv2.erode(img_mask,kernel,iterations = 1)
|
|
|
+ if use_mask:
|
|
|
+ kernel = np.ones((10,10),np.uint8)
|
|
|
+ img_mask = cv2.erode(img_mask,kernel,iterations = 1)
|
|
|
+ else:
|
|
|
+ kernel = np.ones((40,40),np.uint8)
|
|
|
+ img_mask = cv2.erode(img_mask,kernel,iterations = 1)
|
|
|
+ kernel_size = (20, 20)
|
|
|
+ blur_size = tuple(2*i+1 for i in kernel_size)
|
|
|
+ img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
|
|
|
+
|
|
|
+ # kernel = np.ones((10,10),np.uint8)
|
|
|
+ # img_mask = cv2.erode(img_mask,kernel,iterations = 1)
|
|
|
+
|
|
|
+
|
|
|
|
|
|
img_mask /= 255
|
|
|
|
|
|
img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1])
|
|
|
- target_image = np.array(target_image, dtype=np.float)[..., ::-1] * 255
|
|
|
+
|
|
|
+ # pasing mask
|
|
|
+
|
|
|
+ # target_image_parsing = postprocess(target_image, source_image, tgt_mask)
|
|
|
+
|
|
|
+ if use_mask:
|
|
|
+ target_image = np.array(target_image_parsing, dtype=np.float) * 255
|
|
|
+ else:
|
|
|
+ target_image = np.array(target_image, dtype=np.float)[..., ::-1] * 255
|
|
|
+
|
|
|
|
|
|
img_mask_list.append(img_mask)
|
|
|
target_image_list.append(target_image)
|
|
|
+
|
|
|
+
|
|
|
# target_image /= 255
|
|
|
# target_image = 0
|
|
|
img = np.array(oriimg, dtype=np.float)
|
|
@@ -52,7 +173,6 @@ def reverse2wholeimage(swaped_imgs, mats, crop_size, oriimg, logoclass, save_pat
|
|
|
final_img = logoclass.apply_frames(final_img)
|
|
|
cv2.imwrite(save_path, final_img)
|
|
|
|
|
|
- # cv2.imwrite('E:\\lny\\SimSwap-main\\output\\img_div.jpg', img * 255)
|
|
|
- # cv2.imwrite('E:\\lny\\SimSwap-main\\output\\ori_img.jpg', oriimg)
|
|
|
+
|
|
|
|
|
|
|