diff --git a/reactor_patcher.py b/reactor_patcher.py deleted file mode 100644 index efa3c9e..0000000 --- a/reactor_patcher.py +++ /dev/null @@ -1,161 +0,0 @@ -import os.path as osp -import glob -import logging -import insightface -from insightface.model_zoo.model_zoo import ModelRouter, PickableInferenceSession -from insightface.model_zoo.retinaface import RetinaFace -from insightface.model_zoo.landmark import Landmark -from insightface.model_zoo.attribute import Attribute -from insightface.model_zoo.inswapper import INSwapper -from insightface.model_zoo.arcface_onnx import ArcFaceONNX -from insightface.app import FaceAnalysis -from insightface.utils import DEFAULT_MP_NAME, ensure_available -from insightface.model_zoo import model_zoo -import onnxruntime -import onnx -from onnx import numpy_helper -from scripts.reactor_logger import logger - - -def patched_get_model_log(self, **kwargs): - session = PickableInferenceSession(self.onnx_file, **kwargs) - print(f'Applied providers: {session._providers}, with options: {session._provider_options}') - 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] == 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 len(inputs) == 2 and input_shape[2] == 256 and input_shape[3] == 256: - 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: - return None - -def patched_get_model(self, **kwargs): - session = PickableInferenceSession(self.onnx_file, **kwargs) - 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] == 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 len(inputs) == 2 and input_shape[2] == 256 and input_shape[3] == 256: - 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: - return None - - -def patched_faceanalysis_init(self, name=DEFAULT_MP_NAME, root='~/.insightface', allowed_modules=None, **kwargs): - onnxruntime.set_default_logger_severity(3) - self.models = {} - self.model_dir = ensure_available('models', name, root=root) - onnx_files = glob.glob(osp.join(self.model_dir, '*.onnx')) - onnx_files = sorted(onnx_files) - for onnx_file in onnx_files: - model = model_zoo.get_model(onnx_file, **kwargs) - if model is None: - print('model not recognized:', onnx_file) - elif allowed_modules is not None and model.taskname not in allowed_modules: - print('model ignore:', onnx_file, model.taskname) - del model - elif model.taskname not in self.models and (allowed_modules is None or model.taskname in allowed_modules): - self.models[model.taskname] = model - else: - print('duplicated model task type, ignore:', onnx_file, model.taskname) - del model - assert 'detection' in self.models - self.det_model = self.models['detection'] - - -def patched_faceanalysis_prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)): - self.det_thresh = det_thresh - assert det_size is not None - self.det_size = det_size - for taskname, model in self.models.items(): - if taskname == 'detection': - model.prepare(ctx_id, input_size=det_size, det_thresh=det_thresh) - else: - model.prepare(ctx_id) - - -def patched_inswapper_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 - 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 - input_cfg = inputs[0] - input_shape = input_cfg.shape - self.input_shape = input_shape - self.input_size = tuple(input_shape[2:4][::-1]) - - -def pathced_retinaface_prepare(self, ctx_id, **kwargs): - if ctx_id<0: - self.session.set_providers(['CPUExecutionProvider']) - nms_thresh = kwargs.get('nms_thresh', None) - if nms_thresh is not None: - self.nms_thresh = nms_thresh - det_thresh = kwargs.get('det_thresh', None) - if det_thresh is not None: - self.det_thresh = det_thresh - input_size = kwargs.get('input_size', None) - if input_size is not None and self.input_size is None: - self.input_size = input_size - - -def patch_insightface(get_model, faceanalysis_init, faceanalysis_prepare, inswapper_init, retinaface_prepare): - insightface.model_zoo.model_zoo.ModelRouter.get_model = get_model - insightface.app.FaceAnalysis.__init__ = faceanalysis_init - insightface.app.FaceAnalysis.prepare = faceanalysis_prepare - insightface.model_zoo.inswapper.INSwapper.__init__ = inswapper_init - insightface.model_zoo.retinaface.RetinaFace.prepare = retinaface_prepare - - -# original_functions = [ModelRouter.get_model, FaceAnalysis.__init__, FaceAnalysis.prepare, INSwapper.__init__, RetinaFace.prepare] -original_functions = [patched_get_model_log, FaceAnalysis.__init__, FaceAnalysis.prepare, INSwapper.__init__, RetinaFace.prepare] -patched_functions = [patched_get_model, patched_faceanalysis_init, patched_faceanalysis_prepare, patched_inswapper_init, pathced_retinaface_prepare] - - -def apply_patch(console_log_level): - if console_log_level == 0: - patch_insightface(*patched_functions) - logger.setLevel(logging.WARNING) - elif console_log_level == 1: - patch_insightface(*patched_functions) - logger.setLevel(logging.STATUS) - elif console_log_level == 2: - patch_insightface(*original_functions) - logger.setLevel(logging.INFO)