chore: import upstream snapshot with attribution

contrib/omnihub/src/omnihub/frameworks/huggingface.py

@@ -0,0 +1,96 @@
+import os
+
+from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
+
+from omnihub.model_hub import omnihub_dir, ModelHub
+framework_name = 'huggingface'
+framework_dir = os.path.join(omnihub_dir, framework_name)
+BASE_URL = 'https://huggingface.co'
+from transformers import AutoTokenizer, AutoModel, TFAutoModel
+import tensorflow as tf
+import torch
+
+class HuggingFaceModelHub(ModelHub):
+    def __init__(self):
+        """
+        Note that when downloading models for usage from huggingface the URLs should take a very specific format.
+        Since huggingface spaces uses git LFS it uses branch names. Huggingface spaces defaults to the main branch.
+        Typically, the URL formula is: https://huggingface.co + the repo name followed by resolve/main/file_name
+        This file name should be a path to a raw model file. For specific framework tools, feel free to reuse
+        code in this repository
+        """
+        super().__init__(framework_name, BASE_URL)
+
+    def resolve_url(self, repo_path, file_name, branch_name='main'):
+        """
+        Resolve the file name for downloading from huggingface hub.
+        This creates a path using a branch name with a default of main
+        for downloading models from the hub.
+        :param repo_path: repo path to download from. This is usually
+        the namespace after  huggingface.co
+        :param file_name:  the file name to download, this should be a model file
+        :param branch_name: the branch name (defaults to main)
+        :return:  the real url to use for downloading the target file
+        """
+        return f'{repo_path}/resolve/{branch_name}/{file_name}'
+
+    def _download_tf(self,model_path):
+        output_model = TFAutoModel.from_pretrained(model_path)
+        return output_model
+    def _download_pytorch(self,model_path):
+        output_model = AutoModel.from_pretrained(model_path)
+        dummy_inputs = output_model.dummy_inputs
+        inputs_ordered = []
+        non_main_inputs = []
+        for name, array in dummy_inputs.items():
+            if name == output_model.main_input_name:
+                inputs_ordered.append(array)
+            else:
+                non_main_inputs.append(array)
+            ordred_dummy_inputs = inputs_ordered + non_main_inputs
+        return output_model, tuple(ordred_dummy_inputs)
+
+    def download_model(self, model_path, **kwargs) -> str:
+        """
+        Download the model for the given path.
+        A framework_name kwarg is required in order to
+        put the model in the proper location.
+        Due to the nature of huggingface repos being multi framework
+        it's up to the user to specify where a file should go.
+        Valid frameworks are:
+        onnx
+        pytorch
+        tensorflow
+        keras
+        :param model_path: the path to the model to download
+        :param kwargs:  a kwargs containing framework_name as described above
+        :return:  the path to the model
+        """
+        assert 'framework_name' in kwargs
+        model_name = model_path.split('/')[-1]
+        framework_name = kwargs.pop('framework_name')
+
+        if framework_name == 'keras' or framework_name == 'tensorflow':
+            output_model = TFAutoModel.from_pretrained(model_path)
+            callable = tf.function(output_model.call)
+            concrete_function = callable.get_concrete_function(output_model.dummy_inputs)
+            frozen = convert_variables_to_constants_v2(concrete_function)
+            graph_def = frozen.graph.as_graph_def()
+            tf.io.write_graph(graph_def, os.path.join(omnihub_dir,'tensorflow'), f'{model_name}.pb', as_text=False)
+        else:
+            download_function = kwargs.get('download_function', self._download_pytorch)
+            if 'download_function' in kwargs:
+                kwargs.pop('download_function')
+            output_model,dummy_inputs = download_function(model_path)
+            torch.onnx.export(output_model,
+                              dummy_inputs,
+                              f'{os.path.join(omnihub_dir,framework_name)}/{model_name}.onnx',
+                              export_params=True,
+                              do_constant_folding=False,
+                              opset_version=13,
+                              **kwargs)
+
+
+
+    def stage_model(self, model_path: str, model_name: str):
+        super().stage_model(model_path, model_name)

contrib/omnihub/src/omnihub/frameworks/keras.py

@@ -0,0 +1,111 @@
+import os
+
+from tensorflow.python.keras.applications.mobilenet_v2 import MobileNetV2
+from tensorflow.python.keras.applications.mobilenet_v3 import MobileNetV3
+
+from omnihub.model_hub import omnihub_dir, ModelHub
+from tensorflow.keras.applications.vgg19 import VGG19
+from tensorflow.keras.applications.vgg16 import VGG16
+from tensorflow.keras.applications.resnet import ResNet50, ResNet101, ResNet152
+from tensorflow.keras.applications.resnet_v2 import ResNet50V2, ResNet101V2, ResNet152V2
+from tensorflow.keras.applications.densenet import DenseNet121, DenseNet169, DenseNet201
+from tensorflow.keras.applications.inception_resnet_v2 import InceptionResNetV2
+from tensorflow.keras.applications.efficientnet import EfficientNetB0, EfficientNetB1, EfficientNetB2, EfficientNetB3, \
+    EfficientNetB4, EfficientNetB5, EfficientNetB6, EfficientNetB7
+from tensorflow.keras.applications.mobilenet import MobileNet
+from tensorflow.keras.applications.inception_v3 import InceptionV3
+from tensorflow.keras.applications.nasnet import NASNetLarge, NASNetMobile
+from tensorflow.keras.applications.xception import Xception
+framework_name = 'keras'
+framework_dir = os.path.join(omnihub_dir, framework_name)
+BASE_URL = 'https://storage.googleapis.com/tensorflow/keras-applications'
+keras_path = os.path.join(os.path.expanduser('~'), '.keras','models')
+
+
+class KerasModelHub(ModelHub):
+    def __init__(self):
+        super().__init__(framework_name, BASE_URL)
+
+    def download_model(self, model_path, **kwargs) -> str:
+        """
+        Download the model and return the model path on the file system
+        :param model_path:  the model path for the URL
+        :param kwargs: various kwargs for customizing the underlying behavior
+        :return:  the local file path
+        """
+        model_path = self.download_for_url(model_path,**kwargs)
+        return model_path
+
+    def stage_model(self, model_path: str, model_name: str):
+        super().stage_model(model_path, model_name)
+
+    def download_for_url(self, path: str,**kwargs):
+        """
+        Download the file at the given URL
+        :param path:  the path to download
+        :param kwargs:  various kwargs for customizing the underlying behavior of
+        the model download and setup
+        :return: the absolute path to the model
+        """
+        path_split = path.split('/')
+        type = path_split[0]
+        weights_file = path_split[1]
+        include_top = 'no_top' in weights_file
+        if type == 'vgg19':
+            ret = VGG19(include_top=include_top, **kwargs)
+        elif type == 'vgg16':
+            ret = VGG16(include_top=include_top, **kwargs)
+        elif type == 'resnet50':
+            ret = ResNet50(include_top=include_top, **kwargs)
+        elif type == 'resnet101':
+            ret = ResNet101(include_top=include_top, **kwargs)
+        elif type == 'resnet152':
+            ret = ResNet152(include_top=include_top, **kwargs)
+        elif type == 'resnet50v2':
+            ret = ResNet50V2(include_top=include_top, **kwargs)
+        elif type == 'resnet101v2':
+            ret = ResNet101V2(include_top=include_top, **kwargs)
+        elif type == 'resnet152v2':
+            ret = ResNet152V2(include_top=include_top, **kwargs)
+        elif type == 'densenet121':
+            ret = DenseNet121(include_top=include_top)
+        elif type == 'densenet169':
+            ret = DenseNet169(include_top=include_top, **kwargs)
+        elif type == 'densenet201':
+            ret = DenseNet201(include_top=include_top, **kwargs)
+        elif type == 'inceptionresnetv2':
+            ret = InceptionResNetV2(include_top=include_top, **kwargs)
+        elif type == 'efficientnetb0':
+            ret = EfficientNetB0(include_top=include_top, **kwargs)
+        elif type == 'efficientnetb1':
+            ret = EfficientNetB1(include_top=include_top, **kwargs)
+        elif type == 'efficientnetb2':
+            ret = EfficientNetB2(include_top=include_top, **kwargs)
+        elif type == 'efficientnetb3':
+            ret = EfficientNetB3(include_top=include_top, **kwargs)
+        elif type == 'efficientnetb4':
+            ret = EfficientNetB4(include_top=include_top, **kwargs)
+        elif type == 'efficientnetb5':
+            ret = EfficientNetB5(include_top=include_top, **kwargs)
+        elif type == 'efficientnetb6':
+            ret = EfficientNetB6(include_top=include_top, **kwargs)
+        elif type == 'efficientnetb7':
+            efficient_net = EfficientNetB7(include_top=include_top, **kwargs)
+        elif type == 'mobilenet':
+            ret = MobileNet(include_top=include_top, **kwargs)
+        elif type == 'mobilenetv2':
+            ret = MobileNetV2(include_top=include_top)
+        #  MobileNetV3() missing 2 required positional arguments: 'stack_fn' and 'last_point_ch'
+        #elif type == 'mobilenetv3':
+        #    mobile_net = MobileNetV3(include_top=include_top, **kwargs)
+        elif type == 'inceptionv3':
+            ret = InceptionV3(include_top=include_top, **kwargs)
+        elif type == 'nasnet':
+            ret = NASNetLarge(include_top=include_top, **kwargs)
+        elif type == 'nasnet_mobile':
+            ret = NASNetMobile(include_top=include_top, **kwargs)
+        elif type == 'xception':
+            ret = Xception(include_top=include_top, **kwargs)
+        model_path = os.path.join(keras_path, weights_file)
+        ret.save(model_path)
+        return model_path

contrib/omnihub/src/omnihub/frameworks/onnx.py

@@ -0,0 +1,11 @@
+import os
+from omnihub.model_hub import omnihub_dir, ModelHub
+
+framework_name = 'onnx'
+framework_dir = os.path.join(omnihub_dir, framework_name)
+BASE_URL = 'https://media.githubusercontent.com/media/onnx/models/master'
+
+
+class OnnxModelHub(ModelHub):
+    def __init__(self):
+        super().__init__(framework_name, BASE_URL)

contrib/omnihub/src/omnihub/frameworks/pytorch.py

@@ -0,0 +1,84 @@
+import os
+from omnihub.model_hub import omnihub_dir, ModelHub
+import torch
+from torchvision import models
+import numpy as np
+
+framework_name = 'pytorch'
+framework_dir = os.path.join(omnihub_dir, framework_name)
+BASE_URL = 'https://s3.amazonaws.com/pytorch/models'
+
+# models with default 224 x 224 height,width
+MODEL_224_DEFAULTS = ['resnet18', 'vgg16', 'shufflenet_v2_x1_0', 'resnext50_32x4d', 'wide_resnet50_2', 'mnasnet1_0']
+# models with default 256 x 256 height,width
+MODEL_256_DEFAULTS = ['alexnet', 'squeezenet1_0', 'densenet161', 'googlenet', 'inception_v3', 'fasterrcnn', 'ssd',
+                      'retinanet', 'maskrcnn', 'keypointrcnn']
+# models in pytorch's model.detection
+detection_models = ['fasterrcnn', 'ssd', 'retinanet', 'maskrcnn', 'keypointrcnn', 'retinanet']
+# misc defaults and base dictionary for storing heights,widths
+MODEL_DEFAULTS = {
+    'mobilenet_v2': {
+        'height': 32,
+        'width': 32
+    },
+    'mobilenet_v3_large': {
+        'height': 320,
+        'width': 320
+    },
+    'mobilenet_v3_small': {
+        'height': 320,
+        'width': 320
+    },
+    'retinanet': {
+        'height': 512,
+        'width': 512
+    },
+
+}
+
+# efficient_b0 to 7 has height,width default 256 x 256
+for i in range(0, 8):
+    MODEL_256_DEFAULTS.append(f'efficientnet_b{i}')
+
+#regnet_x/y_sizes has default height,width 256 x 256
+regnet_suffix_sizes = ['400mf', '800mf', '1_6gf', '3_2gf', '8gf', '16gf', '32gf']
+for suffix in ['x', 'y']:
+    for size in regnet_suffix_sizes:
+        MODEL_256_DEFAULTS.append(f'regnet_{suffix}_{size}')
+
+for model_224 in MODEL_224_DEFAULTS:
+    MODEL_DEFAULTS[model_224] = {
+        'height': 224,
+        'width': 224
+    }
+
+for model_256 in MODEL_256_DEFAULTS:
+    MODEL_DEFAULTS[model_256] = {
+        'height': 256,
+        'width': 256
+    }
+
+
+class PytorchModelHub(ModelHub):
+    def __init__(self):
+        super().__init__(framework_name, BASE_URL)
+
+    def download_model(self, model_path, **kwargs) -> str:
+        model = None
+        height = kwargs.get('height', MODEL_DEFAULTS[model_path]['height'])
+        width = kwargs.get('width', MODEL_DEFAULTS[model_path]['width'])
+        x = torch.from_numpy(np.ones((1, 3, height, width), dtype=np.float32))
+        if model_path in detection_models:
+            model = models.detection[model_path](pretrained=True, **kwargs)
+        else:
+            model = models.__dict__[model_path](pretrained=True, **kwargs)
+        torch.onnx.export(model,
+                          x,
+                          f'{framework_dir}/{model_path}.onnx',
+                          export_params=True,
+                          do_constant_folding=False,
+                          opset_version=13,
+                          **kwargs)
+
+    def stage_model(self, model_path: str, model_name: str):
+        super().stage_model(model_path, model_name)

contrib/omnihub/src/omnihub/frameworks/tensorflow.py

@@ -0,0 +1,45 @@
+import os
+
+from tensorflow.core.framework.graph_pb2 import GraphDef
+
+from omnihub.model_hub import omnihub_dir, ModelHub
+import tarfile
+from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
+import tensorflow as tf
+import tempfile
+
+framework_name = 'tensorflow'
+framework_dir = os.path.join(omnihub_dir, framework_name)
+BASE_URL = 'https://tfhub.dev'
+
+
+def convert_saved_model(saved_model_dir) -> GraphDef:
+    """
+    Convert the saved model (expanded as a directory)
+    to a frozen graph def
+    :param saved_model_dir: the input model directory
+    :return:  the loaded graph def with all parameters in the model
+    """
+    saved_model = tf.saved_model.load(saved_model_dir)
+    graph_def = saved_model.signatures['serving_default']
+    frozen = convert_variables_to_constants_v2(graph_def)
+    return frozen.graph.as_graph_def()
+
+
+class TensorflowModelHub(ModelHub):
+    def __init__(self):
+        super().__init__(framework_name, BASE_URL)
+
+    def download_model(self, model_path, **kwargs):
+        final_name = model_path.split('/')[-2]
+        model_path = super().download_model(model_path + '?tf-hub-format=compressed')
+        if not tarfile.is_tarfile(model_path):
+            raise Exception(f'Unable to open tar file at path {model_path}')
+
+        mode = kwargs.get('mode', 'r:gz')
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with tarfile.open(model_path, mode=mode) as downloaded:
+                downloaded.extractall(tmpdir)
+                tf.io.write_graph(convert_saved_model(tmpdir), framework_dir, f'{final_name}.pb', as_text=False)
+        # remove extra tar file
+        os.remove(model_path)

contrib/omnihub/src/omnihub/model_hub.py

@@ -0,0 +1,70 @@
+import os
+import shutil
+from typing import IO
+
+import requests
+
+OMNIHUB_HOME = 'OMNIHUB_HOME'
+if os.environ.__contains__(OMNIHUB_HOME):
+    omnihub_dir = os.environ[OMNIHUB_HOME]
+else:
+    omnihub_dir = os.path.join(os.path.expanduser('~'), '.omnihub')
+
+if not os.path.exists(omnihub_dir):
+    os.mkdir(omnihub_dir)
+
+
+class ModelHub(object):
+    def __init__(self, framework_name: str, base_url: str):
+        self.framework_name = framework_name
+        self.stage_model_dir = os.path.join(omnihub_dir, self.framework_name)
+        if not os.path.exists(self.stage_model_dir):
+            os.mkdir(self.stage_model_dir)
+        self.base_url = base_url
+
+    def _download_file(self, url: str,**kwargs):
+        local_filename = os.path.join(self.stage_model_dir, url.split('/')[-1])
+        # NOTE the stream=True parameter below
+        with requests.get(url, stream=True) as r:
+            r.raise_for_status()
+            with open(local_filename, 'wb') as f:
+                for chunk in r.iter_content(chunk_size=8192):
+                    # If you have chunk encoded response uncomment if
+                    # and set chunk_size parameter to None.
+                    # if chunk:
+                    f.write(chunk)
+        return local_filename
+
+    def download_model(self, model_path,**kwargs) -> str:
+        """
+        Meant to be overridden by sub classes.
+        Handles downloading a model with the target URL
+        at the path specified.
+        :param model_path:  the path to the model from the base URL of the web service
+        :return: the path to the original model
+        """
+        model_path = self._download_file(f'{self.base_url}/{model_path}')
+        return model_path
+
+    def stage_model(self, model_path: str, model_name: str):
+        """
+        Copy the model from its original path to the target
+        directory under self.stage_model_dir
+        :param model_path: the original path to the model downloaded
+        by the underlying framework
+        :param model_name: the name of the model file to save as
+        :return:
+        """
+        shutil.copy(model_path, os.path.join(self.stage_model_dir, model_name))
+
+    def stage_model_stream(self, model_path: IO, model_name: str):
+        """
+        Copy the model from its original path to the target
+        directory under self.stage_model_dir
+        :param model_path: the original path to the model downloaded
+        by the underlying framework
+        :param model_name: the name of the model file to save as
+        :return:
+        """
+        with open(os.path.join(self.stage_model_dir, model_name), 'wb+') as f:
+            shutil.copyfileobj(model_path, f)

contrib/omnihub/src/tests/omnihub/test_frameworks.py

@@ -0,0 +1,38 @@
+from omnihub.frameworks.huggingface import HuggingFaceModelHub
+from omnihub.frameworks.keras import KerasModelHub
+from omnihub.frameworks.onnx import OnnxModelHub
+from omnihub.frameworks.pytorch import PytorchModelHub
+from omnihub.frameworks.tensorflow import TensorflowModelHub
+import os
+from omnihub.model_hub import omnihub_dir
+
+
+def test_keras():
+    keras_model_hub = KerasModelHub()
+    model_path = keras_model_hub.download_model('vgg19/vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5')
+    keras_model_hub.stage_model(model_path, 'vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5')
+    assert os.path.exists(os.path.join(omnihub_dir, 'keras', 'vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5'))
+
+
+def test_onnx():
+    onnx_model_hub = OnnxModelHub()
+    onnx_model_hub.download_model('vision/body_analysis/age_gender/models/age_googlenet.onnx')
+    assert os.path.exists(os.path.join(omnihub_dir, 'onnx', 'age_googlenet.onnx'))
+
+
+def test_tensorflow():
+    # https://tfhub.dev/emilutz/vgg19-block4-conv2-unpooling-decoder/1?tf-hub-format=compressed
+    tensorflow_model_hub = TensorflowModelHub()
+    tensorflow_model_hub.download_model('emilutz/vgg19-block4-conv2-unpooling-decoder/1')
+
+
+def test_pytorch():
+    pytorch_model_hub = PytorchModelHub()
+    pytorch_model_hub.download_model('resnet18')
+    assert os.path.exists(os.path.join(omnihub_dir, 'pytorch', 'resnet18.onnx'))
+
+
+def test_huggingface():
+    huggingface_model_hub = HuggingFaceModelHub()
+    huggingface_model_hub.download_model('gpt2',framework_name='pytorch')
+    #assert os.path.exists(os.path.join(omnihub_dir, 'huggingface', 'tf_model.h5'))