diff --git a/embedding-calculator/Dockerfile b/embedding-calculator/Dockerfile index d65249ba25..7af6640d79 100644 --- a/embedding-calculator/Dockerfile +++ b/embedding-calculator/Dockerfile @@ -24,12 +24,12 @@ ENV PYTHONUNBUFFERED=0 ENV JOBLIB_MULTIPROCESSING=0 # download ML models -ARG INTEL_OPTIMIZATION -ARG GPU_IDX +ARG INTEL_OPTIMIZATION=false +ARG GPU_IDX=-1 ENV GPU_IDX=$GPU_IDX INTEL_OPTIMIZATION=$INTEL_OPTIMIZATION -ARG FACE_DETECTION_PLUGIN -ARG CALCULATION_PLUGIN -ARG EXTRA_PLUGINS +ARG FACE_DETECTION_PLUGIN="facenet.FaceDetector" +ARG CALCULATION_PLUGIN="facenet.Calculator" +ARG EXTRA_PLUGINS="rude_carnie.AgeDetector,rude_carnie.GenderDetector" ENV FACE_DETECTION_PLUGIN=$FACE_DETECTION_PLUGIN CALCULATION_PLUGIN=$CALCULATION_PLUGIN \ EXTRA_PLUGINS=$EXTRA_PLUGINS COPY src src diff --git a/embedding-calculator/src/_endpoints.py b/embedding-calculator/src/_endpoints.py index d7f5f62d92..939ec1d500 100644 --- a/embedding-calculator/src/_endpoints.py +++ b/embedding-calculator/src/_endpoints.py @@ -30,12 +30,12 @@ def endpoints(app): @app.route('/status') def status_get(): - availiable_plugins = {p.slug: str(p) - for p in managers.plugin_manager.plugins} + available_plugins = {p.slug: str(p) + for p in managers.plugin_manager.plugins} calculator = managers.plugin_manager.calculator return jsonify(status='OK', build_version=ENV.BUILD_VERSION, calculator_version=str(calculator), - availiable_plugins=availiable_plugins) + available_plugins=available_plugins) @app.route('/find_faces', methods=['POST']) @needs_attached_file diff --git a/embedding-calculator/src/constants.py b/embedding-calculator/src/constants.py index 531c9f02e3..492013ffc9 100644 --- a/embedding-calculator/src/constants.py +++ b/embedding-calculator/src/constants.py @@ -14,7 +14,7 @@ import logging -from src.services.utils.pyutils import get_env, get_env_split, Constants +from src.services.utils.pyutils import get_env, get_env_split, get_env_bool, Constants _DEFAULT_SCANNER = 'Facenet2018' @@ -25,13 +25,14 @@ class ENV(Constants): FACE_DETECTION_PLUGIN = get_env('FACE_DETECTION_PLUGIN', 'facenet.FaceDetector') CALCULATION_PLUGIN = get_env('CALCULATION_PLUGIN', 'facenet.Calculator') - EXTRA_PLUGINS = get_env_split('EXTRA_PLUGINS', 'rude_carnie.AgeDetector,rude_carnie.GenderDetector') + EXTRA_PLUGINS = get_env_split('EXTRA_PLUGINS', '') LOGGING_LEVEL_NAME = get_env('LOGGING_LEVEL_NAME', 'debug').upper() IS_DEV_ENV = get_env('FLASK_ENV', 'production') == 'development' BUILD_VERSION = get_env('APP_VERSION_STRING', 'dev') GPU_IDX = int(get_env('GPU_IDX', '-1')) + INTEL_OPTIMIZATION = get_env_bool('INTEL_OPTIMIZATION') LOGGING_LEVEL = logging._nameToLevel[ENV.LOGGING_LEVEL_NAME] diff --git a/embedding-calculator/src/services/facescan/plugins/conftest.py b/embedding-calculator/src/services/facescan/plugins/conftest.py index 0e32756264..ce396b7346 100644 --- a/embedding-calculator/src/services/facescan/plugins/conftest.py +++ b/embedding-calculator/src/services/facescan/plugins/conftest.py @@ -17,7 +17,7 @@ modules_by_lib = { 'tensorflow': ('facenet', 'rude_carnie'), - 'mexnet': ('insightface',) + 'mxnet': ('insightface',) } modules_to_skip = [] for lib, modules in modules_by_lib.items(): @@ -27,6 +27,6 @@ def pytest_ignore_collect(path): _, tail = os.path.split(path) - for module in modules: + for module in modules_to_skip: if tail.startswith(module): return True diff --git a/embedding-calculator/src/services/facescan/plugins/insightface/__init__.py b/embedding-calculator/src/services/facescan/plugins/insightface/__init__.py index 588ffd2ea4..dba69f354c 100644 --- a/embedding-calculator/src/services/facescan/plugins/insightface/__init__.py +++ b/embedding-calculator/src/services/facescan/plugins/insightface/__init__.py @@ -14,17 +14,16 @@ from src.constants import ENV -from src.services.utils.pyutils import get_env, get_env_bool +from src.services.utils.pyutils import get_env def get_requirements(): cuda_version = get_env('CUDA', '').replace('.', '') - intel_optimization = get_env_bool('INTEL_OPTIMIZATION', False) mxnet_lib = 'mxnet-' if ENV.GPU_IDX > -1 and cuda_version: mxnet_lib += f"cu{cuda_version}" - if intel_optimization: + if ENV.INTEL_OPTIMIZATION: mxnet_lib += 'mkl' mxnet_lib = mxnet_lib.rstrip('-') return ( diff --git a/embedding-calculator/src/services/facescan/plugins/insightface/insightface.py b/embedding-calculator/src/services/facescan/plugins/insightface/insightface.py index defcdf210d..f1f732e543 100644 --- a/embedding-calculator/src/services/facescan/plugins/insightface/insightface.py +++ b/embedding-calculator/src/services/facescan/plugins/insightface/insightface.py @@ -54,7 +54,7 @@ def scaled(self, coefficient: float) -> 'InsightFaceBoundingBox': class InsightFaceMixin: - _CTX_ID = ENV.GPU_ID + _CTX_ID = ENV.GPU_IDX _NMS = 0.4 def get_model_file(self, ml_model: base.MLModel): @@ -122,7 +122,9 @@ def crop_face(self, img: Array3D, box: InsightFaceBoundingBox) -> Array3D: class Calculator(InsightFaceMixin, base.BaseCalculator): ml_models = ( ('arcface_r100_v1', 'http://insightface.ai/files/models/arcface_r100_v1.zip'), - ('arcface_mnet', 'https://drive.google.com/uc?id=1ejWgx_7Nd1PvFXPxCu_1_QNzKfl6v7Hb'), + ('arcface_resnet34', 'https://drive.google.com/uc?id=1Zlork5jpEgUv_4PsfYtJls5lPH2cRaR3'), + ('arcface_resnet50', 'https://drive.google.com/uc?id=1bBtiJr5oIaWmLuSYuKPw-z05VId7Nwci'), + ('arcface_mobilefacenet', 'https://drive.google.com/uc?id=1ejWgx_7Nd1PvFXPxCu_1_QNzKfl6v7Hb'), ) DIFFERENCE_THRESHOLD = 400 diff --git a/embedding-calculator/src/services/facescan/plugins/managers.py b/embedding-calculator/src/services/facescan/plugins/managers.py index 9de58a1fb1..f863c81764 100644 --- a/embedding-calculator/src/services/facescan/plugins/managers.py +++ b/embedding-calculator/src/services/facescan/plugins/managers.py @@ -47,11 +47,11 @@ def requirements(self): return requirements def get_plugins_names(self): - return [ + return list(filter(None, [ constants.ENV.FACE_DETECTION_PLUGIN, constants.ENV.CALCULATION_PLUGIN, *constants.ENV.EXTRA_PLUGINS - ] + ])) @cached_property def plugins(self):