Spaces:
Sleeping
Sleeping
| import tensorflow as tf | |
| import coremltools as ct | |
| import numpy as np | |
| import PIL | |
| from huggingface_hub import hf_hub_download | |
| from huggingface_hub import snapshot_download | |
| import os | |
| # Helper class to extract features from one model, and then feed those features into a classification head | |
| # Because coremltools will only perform inference on OSX, an alternative tensorflow inference pipeline uses | |
| # a tensorflow feature extractor and feeds the features into a dynamically created keras model based on the coreml classification head. | |
| class CoreMLPipeline: | |
| def __init__(self, config, auth_key, use_tf): | |
| self.config = config | |
| self.use_tf = use_tf | |
| if use_tf: | |
| extractor_path = snapshot_download(repo_id=config["tf_extractor_repoid"], use_auth_token = auth_key) | |
| else: | |
| extractor_path = hf_hub_download(repo_id=config["coreml_extractor_repoid"], | |
| filename=config["coreml_extractor_path"], use_auth_token = auth_key) | |
| classifier_path = hf_hub_download(repo_id=config["coreml_classifier_repoid"], filename=config["coreml_classifier_path"], | |
| use_auth_token = auth_key) | |
| print(f"Loading extractor...{extractor_path}") | |
| if use_tf: | |
| self.extractor = tf.saved_model.load(os.path.join(extractor_path, config["tf_extractor_path"])) | |
| else: | |
| self.extractor = ct.models.MLModel(extractor_path) | |
| print(f"Loading classifier...{classifier_path}") | |
| self.classifier = ct.models.MLModel(classifier_path) | |
| if use_tf: | |
| self.make_keras_model() | |
| #Only MacOS can run inference on CoreML models. Convert it to tensorflow to match the tf feature extractor | |
| def make_keras_model(self): | |
| spec = self.classifier.get_spec() | |
| nnClassifier = spec.neuralNetworkClassifier | |
| labels = nnClassifier.stringClassLabels.vector | |
| input = tf.keras.Input(shape = (1280)) | |
| activation = "sigmoid" if len(labels) == 1 else "softmax" | |
| x = tf.keras.layers.Dense(len(labels), activation = activation)(input) | |
| model = tf.keras.Model(input,x, trainable = False) | |
| weights = np.array(nnClassifier.layers[0].innerProduct.weights.floatValue) | |
| weights = weights.reshape((len(labels),1280)) | |
| weights = weights.T | |
| bias = np.array(nnClassifier.layers[0].innerProduct.bias.floatValue) | |
| model.set_weights([weights,bias]) | |
| self.tf_model = model | |
| self.labels = labels | |
| def classify(self,resized): | |
| if self.use_tf: | |
| image = tf.image.convert_image_dtype(resized, tf.float32) | |
| image = tf.expand_dims(image, 0) | |
| features = self.extractor.signatures['serving_default'](image) | |
| input = {"input_1":features["output_1"]} | |
| output = self.tf_model.predict(input) | |
| results = {} | |
| for i,label in enumerate(self.labels): | |
| results[label] = output[i] | |
| else: | |
| features = self.extractor.predict({"image":resized}) | |
| features = features["Identity"] | |
| output = self.classifier.predict({"features":features[0]}) | |
| results = output["Identity"] | |
| return results | |