ZF/insightface
1
0
Fork 0
mirror of https://github.com/deepinsight/insightface.git synced 2025-12-30 08:02:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

add inswapper, to pip 0.7

Browse Source
This commit is contained in:
Jia Guo
2022-11-28 21:05:17 +08:00
parent 6258ab9869
commit bbbb45c460
6 changed files with 200 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
examples/in_swapper/inswapper_main.py Normal file
View File

@@ -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)

2
python-package/insightface/__init__.py
View File

@@ -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

105
python-package/insightface/model_zoo/inswapper.py Normal file
View File

@@ -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] = 0
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

27
python-package/insightface/model_zoo/model_zoo.py
View File

@@ -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

75
python-package/insightface/utils/face_align.py
View File

@@ -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

22
python-package/insightface/utils/storage.py
View File

@@ -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