From bbbb45c460b1d45e1e6838cdd8872f38e6ee2bde Mon Sep 17 00:00:00 2001 From: Jia Guo Date: Mon, 28 Nov 2022 21:05:17 +0800 Subject: [PATCH] add inswapper, to pip 0.7 --- examples/in_swapper/inswapper_main.py | 42 +++++++ python-package/insightface/__init__.py | 2 +- .../insightface/model_zoo/inswapper.py | 105 ++++++++++++++++++ .../insightface/model_zoo/model_zoo.py | 27 +++-- .../insightface/utils/face_align.py | 75 +++---------- python-package/insightface/utils/storage.py | 22 ++-- 6 files changed, 200 insertions(+), 73 deletions(-) create mode 100644 examples/in_swapper/inswapper_main.py create mode 100644 python-package/insightface/model_zoo/inswapper.py diff --git a/examples/in_swapper/inswapper_main.py b/examples/in_swapper/inswapper_main.py new file mode 100644 index 0000000..828c244 --- /dev/null +++ b/examples/in_swapper/inswapper_main.py @@ -0,0 +1,42 @@ +import datetime +import numpy as np +import os +import os.path as osp +import glob +import cv2 +import insightface +from insightface.app import FaceAnalysis +from insightface.data import get_image as ins_get_image + + +assert insightface.__version__>='0.7' + +def detect_person(img, detector): + bboxes, kpss = detector.detect(img) + bboxes = np.round(bboxes[:,:4]).astype(np.int) + kpss = np.round(kpss).astype(np.int) + kpss[:,:,0] = np.clip(kpss[:,:,0], 0, img.shape[1]) + kpss[:,:,1] = np.clip(kpss[:,:,1], 0, img.shape[0]) + vbboxes = bboxes.copy() + vbboxes[:,0] = kpss[:, 0, 0] + vbboxes[:,1] = kpss[:, 0, 1] + vbboxes[:,2] = kpss[:, 4, 0] + vbboxes[:,3] = kpss[:, 4, 1] + return bboxes, vbboxes + +if __name__ == '__main__': + app = FaceAnalysis(name='buffalo_l') + app.prepare(ctx_id=0, det_size=(640, 640)) + swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True) + + + img = ins_get_image('t1') + faces = app.get(img) + faces = sorted(faces, key = lambda x : x.bbox[0]) + assert len(faces)==6 + source_face = faces[2] + for face in faces: + img = swapper.get(img, face, source_face, paste_back=True) + cv2.imwrite("./t1_swapped.jpg", img) + + diff --git a/python-package/insightface/__init__.py b/python-package/insightface/__init__.py index 4acb7fb..b7a2e2b 100644 --- a/python-package/insightface/__init__.py +++ b/python-package/insightface/__init__.py @@ -11,7 +11,7 @@ except ImportError: "Unable to import dependency onnxruntime. " ) -__version__ = '0.6.3' +__version__ = '0.7' from . import model_zoo from . import utils diff --git a/python-package/insightface/model_zoo/inswapper.py b/python-package/insightface/model_zoo/inswapper.py new file mode 100644 index 0000000..dcaceb1 --- /dev/null +++ b/python-package/insightface/model_zoo/inswapper.py @@ -0,0 +1,105 @@ +import time +import numpy as np +import onnxruntime +import cv2 +import onnx +from onnx import numpy_helper +from ..utils import face_align + + + + +class INSwapper(): + def __init__(self, model_file=None, session=None): + self.model_file = model_file + self.session = session + model = onnx.load(self.model_file) + graph = model.graph + self.emap = numpy_helper.to_array(graph.initializer[-1]) + self.input_mean = 0.0 + self.input_std = 255.0 + #print('input mean and std:', model_file, self.input_mean, self.input_std) + if self.session is None: + self.session = onnxruntime.InferenceSession(self.model_file, None) + inputs = self.session.get_inputs() + self.input_names = [] + for inp in inputs: + self.input_names.append(inp.name) + outputs = self.session.get_outputs() + output_names = [] + for out in outputs: + output_names.append(out.name) + self.output_names = output_names + assert len(self.output_names)==1 + output_shape = outputs[0].shape + input_cfg = inputs[0] + input_shape = input_cfg.shape + self.input_shape = input_shape + print('inswapper-shape:', self.input_shape) + self.input_size = tuple(input_shape[2:4][::-1]) + + def forward(self, img, latent): + img = (img - self.input_mean) / self.input_std + pred = self.session.run(self.output_names, {self.input_names[0]: img, self.input_names[1]: latent})[0] + return pred + + def get(self, img, target_face, source_face, paste_back=True): + aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0]) + blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size, + (self.input_mean, self.input_mean, self.input_mean), swapRB=True) + latent = source_face.normed_embedding.reshape((1,-1)) + latent = np.dot(latent, self.emap) + latent /= np.linalg.norm(latent) + pred = self.session.run(self.output_names, {self.input_names[0]: blob, self.input_names[1]: latent})[0] + #print(latent.shape, latent.dtype, pred.shape) + img_fake = pred.transpose((0,2,3,1))[0] + bgr_fake = np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:,:,::-1] + if not paste_back: + return bgr_fake, M + else: + target_img = img + fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32) + fake_diff = np.abs(fake_diff).mean(axis=2) + fake_diff[:2,:] = 0 + fake_diff[-2:,:] = 0 + fake_diff[:,:2] = 0 + fake_diff[:,-2:] = 0 + IM = cv2.invertAffineTransform(M) + img_white = np.full((aimg.shape[0],aimg.shape[1]), 255, dtype=np.float32) + bgr_fake = cv2.warpAffine(bgr_fake, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0) + img_white = cv2.warpAffine(img_white, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0) + fake_diff = cv2.warpAffine(fake_diff, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0) + img_white[img_white>20] = 255 + fthresh = 10 + fake_diff[fake_diff=fthresh] = 255 + img_mask = img_white + mask_h_inds, mask_w_inds = np.where(img_mask==255) + mask_h = np.max(mask_h_inds) - np.min(mask_h_inds) + mask_w = np.max(mask_w_inds) - np.min(mask_w_inds) + mask_size = int(np.sqrt(mask_h*mask_w)) + k = max(mask_size//10, 10) + #k = max(mask_size//20, 6) + #k = 6 + kernel = np.ones((k,k),np.uint8) + img_mask = cv2.erode(img_mask,kernel,iterations = 1) + kernel = np.ones((2,2),np.uint8) + fake_diff = cv2.dilate(fake_diff,kernel,iterations = 1) + k = max(mask_size//20, 5) + #k = 3 + #k = 3 + kernel_size = (k, k) + blur_size = tuple(2*i+1 for i in kernel_size) + img_mask = cv2.GaussianBlur(img_mask, blur_size, 0) + k = 5 + kernel_size = (k, k) + blur_size = tuple(2*i+1 for i in kernel_size) + fake_diff = cv2.GaussianBlur(fake_diff, blur_size, 0) + img_mask /= 255 + fake_diff /= 255 + #img_mask = fake_diff + img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1]) + fake_merged = img_mask * bgr_fake + (1-img_mask) * target_img.astype(np.float32) + fake_merged = fake_merged.astype(np.uint8) + return fake_merged + diff --git a/python-package/insightface/model_zoo/model_zoo.py b/python-package/insightface/model_zoo/model_zoo.py index 1429b72..fc62831 100644 --- a/python-package/insightface/model_zoo/model_zoo.py +++ b/python-package/insightface/model_zoo/model_zoo.py @@ -13,6 +13,7 @@ from .retinaface import * #from .scrfd import * from .landmark import * from .attribute import Attribute +from .inswapper import INSwapper from ..utils import download_onnx __all__ = ['get_model'] @@ -38,18 +39,21 @@ class ModelRouter: def get_model(self, **kwargs): session = PickableInferenceSession(self.onnx_file, **kwargs) print(f'Applied providers: {session._providers}, with options: {session._provider_options}') - input_cfg = session.get_inputs()[0] + inputs = session.get_inputs() + input_cfg = inputs[0] input_shape = input_cfg.shape outputs = session.get_outputs() if len(outputs)>=5: return RetinaFace(model_file=self.onnx_file, session=session) - elif input_shape[2]==112 and input_shape[3]==112: - return ArcFaceONNX(model_file=self.onnx_file, session=session) elif input_shape[2]==192 and input_shape[3]==192: return Landmark(model_file=self.onnx_file, session=session) elif input_shape[2]==96 and input_shape[3]==96: return Attribute(model_file=self.onnx_file, session=session) + elif len(inputs)==2 and input_shape[2]==128 and input_shape[3]==128: + return INSwapper(model_file=self.onnx_file, session=session) + elif input_shape[2]==input_shape[3] and input_shape[2]>=112 and input_shape[2]%16==0: + return ArcFaceONNX(model_file=self.onnx_file, session=session) else: #raise RuntimeError('error on model routing') return None @@ -63,11 +67,18 @@ def find_onnx_file(dir_path): paths = sorted(paths) return paths[-1] +def get_default_providers(): + return ['CUDAExecutionProvider', 'CPUExecutionProvider'] + +def get_default_provider_options(): + return None + def get_model(name, **kwargs): root = kwargs.get('root', '~/.insightface') root = os.path.expanduser(root) model_root = osp.join(root, 'models') allow_download = kwargs.get('download', False) + download_zip = kwargs.get('download_zip', False) if not name.endswith('.onnx'): model_dir = os.path.join(model_root, name) model_file = find_onnx_file(model_dir) @@ -76,10 +87,12 @@ def get_model(name, **kwargs): else: model_file = name if not osp.exists(model_file) and allow_download: - model_file = download_onnx('models', model_file, root=root) - assert osp.exists(model_file), 'model_file should exist' - assert osp.isfile(model_file), 'model_file should be file' + model_file = download_onnx('models', model_file, root=root, download_zip=download_zip) + assert osp.exists(model_file), 'model_file %s should exist'%model_file + assert osp.isfile(model_file), 'model_file %s should be a file'%model_file router = ModelRouter(model_file) - model = router.get_model(providers=kwargs.get('providers'), provider_options=kwargs.get('provider_options')) + providers = kwargs.get('providers', get_default_providers()) + provider_options = kwargs.get('provider_options', get_default_provider_options()) + model = router.get_model(providers=providers, provider_options=provider_options) return model diff --git a/python-package/insightface/utils/face_align.py b/python-package/insightface/utils/face_align.py index c4e9a7c..57f90f9 100644 --- a/python-package/insightface/utils/face_align.py +++ b/python-package/insightface/utils/face_align.py @@ -2,76 +2,35 @@ import cv2 import numpy as np from skimage import transform as trans -src1 = np.array([[51.642, 50.115], [57.617, 49.990], [35.740, 69.007], - [51.157, 89.050], [57.025, 89.702]], - dtype=np.float32) -#<--left -src2 = np.array([[45.031, 50.118], [65.568, 50.872], [39.677, 68.111], - [45.177, 86.190], [64.246, 86.758]], - dtype=np.float32) -#---frontal -src3 = np.array([[39.730, 51.138], [72.270, 51.138], [56.000, 68.493], - [42.463, 87.010], [69.537, 87.010]], - dtype=np.float32) - -#-->right -src4 = np.array([[46.845, 50.872], [67.382, 50.118], [72.737, 68.111], - [48.167, 86.758], [67.236, 86.190]], - dtype=np.float32) - -#-->right profile -src5 = np.array([[54.796, 49.990], [60.771, 50.115], [76.673, 69.007], - [55.388, 89.702], [61.257, 89.050]], - dtype=np.float32) - -src = np.array([src1, src2, src3, src4, src5]) -src_map = {112: src, 224: src * 2} - -arcface_src = np.array( +arcface_dst = np.array( [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.7299, 92.2041]], dtype=np.float32) -arcface_src = np.expand_dims(arcface_src, axis=0) - -# In[66]: - - -# lmk is prediction; src is template -def estimate_norm(lmk, image_size=112, mode='arcface'): +def estimate_norm(lmk, image_size=112,mode='arcface'): assert lmk.shape == (5, 2) - tform = trans.SimilarityTransform() - lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1) - min_M = [] - min_index = [] - min_error = float('inf') - if mode == 'arcface': - if image_size == 112: - src = arcface_src - else: - src = float(image_size) / 112 * arcface_src + assert image_size%112==0 or image_size%128==0 + if image_size%112==0: + ratio = float(image_size)/112.0 else: - src = src_map[image_size] - for i in np.arange(src.shape[0]): - tform.estimate(lmk, src[i]) - M = tform.params[0:2, :] - results = np.dot(M, lmk_tran.T) - results = results.T - error = np.sum(np.sqrt(np.sum((results - src[i])**2, axis=1))) - # print(error) - if error < min_error: - min_error = error - min_M = M - min_index = i - return min_M, min_index - + ratio = float(image_size)/128.0 + dst = arcface_dst * ratio + tform = trans.SimilarityTransform() + tform.estimate(lmk, dst) + M = tform.params[0:2, :] + return M def norm_crop(img, landmark, image_size=112, mode='arcface'): - M, pose_index = estimate_norm(landmark, image_size, mode) + M = estimate_norm(landmark, image_size, mode) warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0) return warped +def norm_crop2(img, landmark, image_size=112, mode='arcface'): + M = estimate_norm(landmark, image_size, mode) + warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0) + return warped, M + def square_crop(im, S): if im.shape[0] > im.shape[1]: height = S diff --git a/python-package/insightface/utils/storage.py b/python-package/insightface/utils/storage.py index de6f6be..b03fe27 100644 --- a/python-package/insightface/utils/storage.py +++ b/python-package/insightface/utils/storage.py @@ -29,7 +29,7 @@ def download(sub_dir, name, force=False, root='~/.insightface'): def ensure_available(sub_dir, name, root='~/.insightface'): return download(sub_dir, name, force=False, root=root) -def download_onnx(sub_dir, model_file, force=False, root='~/.insightface'): +def download_onnx(sub_dir, model_file, force=False, root='~/.insightface', download_zip=False): _root = os.path.expanduser(root) model_root = osp.join(_root, sub_dir) new_model_file = osp.join(model_root, model_file) @@ -38,9 +38,17 @@ def download_onnx(sub_dir, model_file, force=False, root='~/.insightface'): if not osp.exists(model_root): os.makedirs(model_root) print('download_path:', new_model_file) - model_url = "%s/%s/%s"%(BASE_REPO_URL, sub_dir, model_file) - #model_url = "%s/%s"%(BASE_REPO_URL, model_file) - download_file(model_url, - path=new_model_file, - overwrite=True) - return new_model_file + if not download_zip: + model_url = "%s/%s/%s"%(BASE_REPO_URL, sub_dir, model_file) + download_file(model_url, + path=new_model_file, + overwrite=True) + else: + model_url = "%s/%s/%s.zip"%(BASE_REPO_URL, sub_dir, model_file) + zip_file_path = new_model_file+".zip" + download_file(model_url, + path=zip_file_path, + overwrite=True) + with zipfile.ZipFile(zip_file_path) as zf: + zf.extractall(model_root) + return new_model_file