Refactor for cleanup and add more samples

app.py

@@ -1,71 +1,28 @@
-import tensorflow as tf
 import gradio as gr
-import coremltools as ct
-import numpy as np
-import requests
-import huggingface_hub as hf
-from huggingface_hub import hf_hub_download
-from huggingface_hub import snapshot_download
-from huggingface_hub import login
 import os
-import PIL
+import platform
+from helper import CoreMLPipeline
 
-#login()
+force_tf = os.environ.get('FORCE_TF', False)
+auth_key = os.environ.get('HF_TOKEN', True)
 
+config = { "coreml_extractor_repoid":"crossprism/efficientnetv2-21k-fv-m",
+           "coreml_extractor_path":"efficientnetV2M21kExtractor.mlmodel",
+           "tf_extractor_repoid":"crossprism/efficientnetv2-21k-fv-m-tf",
+           "tf_extractor_path":"efficientnetv2-21k-fv-m",
+           "coreml_classifier_repoid":"crossprism/tesla_sentry_dings",
+           "coreml_classifier_path":"tesla_sentry_door_ding.mlpackage/Data/com.apple.CoreML/tesla_door_dings.mlmodel"
+           }
+use_tf = force_tf or (platform.system() != 'Darwin')
 
-read_key = os.environ.get('HF_TOKEN', True)
-#extractor2_path = hf_hub_download(repo_id="crossprism/efficientnetv221k-M", filename="efficientnetV2M21kExtractor.mlmodel", use_auth_token  = read_key)
-extractor_path = snapshot_download(repo_id="crossprism/efficientnetv2-21k-fv-m", use_auth_token  = read_key)
-classifier_path = hf_hub_download(repo_id="crossprism/tesla_sentry_dings", filename="tesla_sentry_door_ding.mlpackage/Data/com.apple.CoreML/tesla_door_dings.mlmodel", use_auth_token = read_key)
-
-
-print(f"Loading extractor...{extractor_path}")
-extractor = tf.saved_model.load(extractor_path+"/efficientnetv2-21k-fv-m")
-
-#extractor2 = ct.models.MLModel(extractor2_path)
-print(f"Loading classifier...{classifier_path}")
-classifier = ct.models.MLModel(classifier_path)
-spec = classifier.get_spec()
-labels = spec.neuralNetworkClassifier.stringClassLabels.vector
-image = PIL.Image.open('test.jpg')
-
-def makeKerasModel(labels, coreml_classifier):
-    input = tf.keras.Input(shape = (1280))
-    x = tf.keras.layers.Dense(len(labels), activation = "sigmoid")(input)
-    model = tf.keras.Model(input,x, trainable = False)
-    weights = np.array(coreml_classifier.layers[0].innerProduct.weights.floatValue)
-    weights = weights.reshape((len(labels),1280))
-    #weights = weights.reshape((1280,len(labels)))
-    weights = weights.T
-    bias = np.array(coreml_classifier.layers[0].innerProduct.bias.floatValue)    
-    model.set_weights([weights,bias])    
-    return model
-
-#Only MacOS can run inference on CoreML models. Convert it to tensorflow to match the tf feature extractor
-tf_classifier = makeKerasModel(labels, spec.neuralNetworkClassifier)
+helper = CoreMLPipeline(config, auth_key, use_tf)
 
 def classify_image(image):
     resized = image.resize((480,480))
-    image = tf.image.convert_image_dtype(resized, tf.float32)
-    image = tf.reshape(image, [1,480,480,3])
-    features = extractor.signatures['serving_default'](image)
-    #features2 = extractor2.predict({"image":resized})
-    #print(features)
-    #print(features2)
-    #features2 = features2["Identity"]
-    #isDing = classifier.predict({"features":features2[0]})
-    #isDing = isDing["Identity"]
-    input = {"input_1":features["output_1"]}
-    p = tf_classifier.predict(input)
-    #print(p)
-    isDing = {}
-    for i,label in enumerate(labels):
-        isDing[label] = p[i]        
-    #print(isDing)
-    return {'ding': isDing["ding"]}
-#classify_image(image)
+    return helper.classify(resized)
+
 image = gr.Image(type='pil')
 label = gr.Label(num_top_classes=3)
 
-gr.Interface(fn=classify_image, inputs=image, outputs=label, examples = [["test.jpg","test2.jpg","test3.jpg"]]).launch()
+gr.Interface(fn=classify_image, inputs=image, outputs=label, examples = [["test.jpg"],["test2.jpg"],["test3.jpg"]]).launch()
 

helper.py

@@ -0,0 +1,70 @@
+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
+