pip 0.4, supports onnx auto-downloading

python-package/insightface/__init__.py

@@ -11,7 +11,7 @@ except ImportError:
         "Unable to import dependency onnxruntime. "
     )
 
-__version__ = '0.3.8'
+__version__ = '0.4'
 
 from . import model_zoo
 from . import utils

python-package/insightface/model_zoo/model_zoo.py

@@ -12,6 +12,7 @@ from .arcface_onnx import *
 from .scrfd import *
 from .landmark import *
 from .attribute import Attribute
+from ..utils import download_onnx
 
 #__all__ = ['get_model', 'get_model_list', 'get_arcface_onnx', 'get_scrfd']
 __all__ = ['get_model']
@@ -49,18 +50,22 @@ def find_onnx_file(dir_path):
     return paths[-1]
 
 def get_model(name, **kwargs):
-    root = kwargs.get('root', '~/.insightface/models')
+    root = kwargs.get('root', '~/.insightface')
     root = os.path.expanduser(root)
+    model_root = osp.join(root, 'models')
+    allow_download = kwargs.get('download', False)
     if not name.endswith('.onnx'):
-        model_dir = os.path.join(root, name)
+        model_dir = os.path.join(model_root, name)
         model_file = find_onnx_file(model_dir)
         if model_file is None:
             return None
     else:
         model_file = name
-    assert osp.isfile(model_file), 'model should be file'
+    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'
     router = ModelRouter(model_file)
     model = router.get_model()
-    #print('get-model for ', name,' : ', model.taskname)
     return model
 

python-package/insightface/utils/__init__.py

@@ -6,7 +6,7 @@ from __future__ import absolute_import
 #from . import metrics
 #from . import parallel
 
-from .storage import download, ensure_available
+from .storage import download, ensure_available, download_onnx
 from .filesystem import get_model_dir
 from .filesystem import makedirs, try_import_dali
 from .constant import *

python-package/insightface/utils/storage.py

@@ -27,3 +27,17 @@ 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'):
+    _root = os.path.expanduser(root)
+    model_root = osp.join(_root, sub_dir)
+    new_model_file = osp.join(model_root, model_file)
+    if osp.exists(new_model_file) and not force:
+        return new_model_file
+    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)
+    download_file(model_url,
+             path=new_model_file,
+             overwrite=True)
+    return new_model_file