|
|
import cv2 |
|
|
import onnxruntime as rt |
|
|
import sys |
|
|
from insightface.app import FaceAnalysis |
|
|
sys.path.insert(1, './recognition') |
|
|
from scrfd import SCRFD |
|
|
from arcface_onnx import ArcFaceONNX |
|
|
import os.path as osp |
|
|
import os |
|
|
from pathlib import Path |
|
|
from tqdm import tqdm |
|
|
import ffmpeg |
|
|
import random |
|
|
import multiprocessing as mp |
|
|
from concurrent.futures import ThreadPoolExecutor |
|
|
from insightface.model_zoo.inswapper import INSwapper |
|
|
import psutil |
|
|
from enum import Enum |
|
|
from insightface.app.common import Face |
|
|
from insightface.utils.storage import ensure_available |
|
|
import re |
|
|
import subprocess |
|
|
|
|
|
class RefacerMode(Enum): |
|
|
CPU, CUDA, COREML, TENSORRT = range(1, 5) |
|
|
|
|
|
class Refacer: |
|
|
def __init__(self,force_cpu=False,colab_performance=False): |
|
|
self.first_face = False |
|
|
self.force_cpu = force_cpu |
|
|
self.colab_performance = colab_performance |
|
|
self.__check_encoders() |
|
|
self.__check_providers() |
|
|
self.total_mem = psutil.virtual_memory().total |
|
|
self.__init_apps() |
|
|
|
|
|
def __check_providers(self): |
|
|
if self.force_cpu : |
|
|
self.providers = ['CPUExecutionProvider'] |
|
|
else: |
|
|
self.providers = rt.get_available_providers() |
|
|
rt.set_default_logger_severity(4) |
|
|
self.sess_options = rt.SessionOptions() |
|
|
self.sess_options.execution_mode = rt.ExecutionMode.ORT_SEQUENTIAL |
|
|
self.sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL |
|
|
|
|
|
if len(self.providers) == 1 and 'CPUExecutionProvider' in self.providers: |
|
|
self.mode = RefacerMode.CPU |
|
|
self.use_num_cpus = mp.cpu_count()-1 |
|
|
self.sess_options.intra_op_num_threads = int(self.use_num_cpus/3) |
|
|
print(f"CPU mode with providers {self.providers}") |
|
|
elif self.colab_performance: |
|
|
self.mode = RefacerMode.TENSORRT |
|
|
self.use_num_cpus = mp.cpu_count()-1 |
|
|
self.sess_options.intra_op_num_threads = int(self.use_num_cpus/3) |
|
|
print(f"TENSORRT mode with providers {self.providers}") |
|
|
elif 'CoreMLExecutionProvider' in self.providers: |
|
|
self.mode = RefacerMode.COREML |
|
|
self.use_num_cpus = mp.cpu_count()-1 |
|
|
self.sess_options.intra_op_num_threads = int(self.use_num_cpus/3) |
|
|
print(f"CoreML mode with providers {self.providers}") |
|
|
elif 'CUDAExecutionProvider' in self.providers: |
|
|
self.mode = RefacerMode.CUDA |
|
|
self.use_num_cpus = 2 |
|
|
self.sess_options.intra_op_num_threads = 1 |
|
|
if 'TensorrtExecutionProvider' in self.providers: |
|
|
self.providers.remove('TensorrtExecutionProvider') |
|
|
print(f"CUDA mode with providers {self.providers}") |
|
|
|
|
|
def __init_apps(self): |
|
|
assets_dir = ensure_available('models', 'buffalo_l', root='~/.insightface') |
|
|
|
|
|
model_path = os.path.join(assets_dir, 'det_10g.onnx') |
|
|
sess_face = rt.InferenceSession(model_path, self.sess_options, providers=self.providers) |
|
|
self.face_detector = SCRFD(model_path,sess_face) |
|
|
self.face_detector.prepare(0,input_size=(640, 640)) |
|
|
|
|
|
model_path = os.path.join(assets_dir , 'w600k_r50.onnx') |
|
|
sess_rec = rt.InferenceSession(model_path, self.sess_options, providers=self.providers) |
|
|
self.rec_app = ArcFaceONNX(model_path,sess_rec) |
|
|
self.rec_app.prepare(0) |
|
|
|
|
|
model_path = 'inswapper_128.onnx' |
|
|
sess_swap = rt.InferenceSession(model_path, self.sess_options, providers=self.providers) |
|
|
self.face_swapper = INSwapper(model_path,sess_swap) |
|
|
|
|
|
|
|
|
|
