Spaces:
Build error
Build error
File size: 1,473 Bytes
fa8453f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import time
import onnx
import cv2
import onnxruntime
import numpy as np
from onnx import numpy_helper
from numpy.linalg import norm as l2norm
from utils.face_alignment import norm_crop2
class Inswapper():
def __init__(self, model_file=None, provider=['CPUExecutionProvider'], session_options=None):
self.model_file = model_file
model = onnx.load(self.model_file)
graph = model.graph
self.emap = numpy_helper.to_array(graph.initializer[-1])
self.session_options = session_options
if self.session_options is None:
self.session_options = onnxruntime.SessionOptions()
self.session = onnxruntime.InferenceSession(self.model_file, sess_options=self.session_options, providers=provider)
def forward(self, frame, target, source, n_pass=1):
trg, matrix = norm_crop2(frame, target['kps'], 128)
latent = source['embedding'].reshape((1, -1))
latent = np.dot(latent, self.emap)
latent /= np.linalg.norm(latent)
blob = trg.astype('float32') / 255
blob = blob[:, :, ::-1]
blob = np.expand_dims(blob, axis=0).transpose(0, 3, 1, 2)
for _ in range(max(int(n_pass),1)):
blob = self.session.run(['output'], {'target': blob, 'source': latent})[0]
out = blob[0].transpose((1, 2, 0))
out = (out * 255).clip(0,255)
out = out.astype('uint8')[:, :, ::-1]
del blob, latent
return trg, out, matrix
|