Spaces:
Sleeping
Sleeping
File size: 3,308 Bytes
1a8eef2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
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
|