support coreml and tflite model testing

generate_thresholds.py

@@ -13,7 +13,7 @@
 from sklearn.metrics import precision_recall_curve
 import warnings
 from lib.model_taxonomy_dataframe import ModelTaxonomyDataframe
-from lib.tf_gp_elev_model import TFGeoPriorModelElev
+from lib.geo_inferrer_tf import TFGeoPriorModelElev
 
 
 def ignore_shapely_deprecation_warning(message, category, filename, lineno, file=None, line=None):

lib/geo_inferrer.py

@@ -0,0 +1,48 @@
+from abc import ABC, abstractmethod
+import math
+
+import numpy as np
+import tensorflow as tf
+
+
+class GeoInferrer(ABC):
+    @abstractmethod
+    def __init__(self, model_path: str):
+        """Subclasses must implement this constructor."""
+        pass
+
+    @abstractmethod
+    def predict(
+        self, latitude: float, longitude: float, elevation: float
+    ) -> np.ndarray:
+        """
+        given a location, calculate geo results
+
+        Subclasses must implement this method.
+        """
+        pass
+
+    @staticmethod
+    def encode_loc(latitude, longitude, elevation):
+        latitude = np.array(latitude)
+        longitude = np.array(longitude)
+        elevation = np.array(elevation)
+        elevation = elevation.astype("float32")
+        grid_lon = longitude.astype("float32") / 180.0
+        grid_lat = latitude.astype("float32") / 90.0
+
+        elevation[elevation > 0] = elevation[elevation > 0] / 6574.0
+        elevation[elevation < 0] = elevation[elevation < 0] / 32768.0
+        norm_elev = elevation
+
+        norm_loc = tf.stack([grid_lon, grid_lat], axis=1)
+
+        encoded_loc = tf.concat(
+            [
+                tf.sin(norm_loc * math.pi),
+                tf.cos(norm_loc * math.pi),
+                tf.expand_dims(norm_elev, axis=1),
+            ],
+            axis=1,
+        )
+        return encoded_loc

lib/geo_inferrer_coreml.py

@@ -0,0 +1,21 @@
+import coremltools as ct
+import numpy as np
+
+from lib.geo_inferrer import GeoInferrer
+
+
+class CoremlGeoPriorModelElev(GeoInferrer):
+
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.gpmodel = ct.models.MLModel(self.model_path)
+
+    def predict(
+        self, latitude: float, longitude: float, elevation: float
+    ) -> np.ndarray:
+        encoded_loc = GeoInferrer.encode_loc(
+            [latitude], [longitude], [elevation]
+        ).numpy()
+        out_dict = self.gpmodel.predict({"input_1": encoded_loc})
+        preds = out_dict["Identity"][0]
+        return preds

lib/geo_inferrer_factory.py

@@ -0,0 +1,20 @@
+from sys import platform
+
+from lib.geo_inferrer import GeoInferrer
+from lib.geo_inferrer_coreml import CoremlGeoPriorModelElev
+from lib.geo_inferrer_tflite import TFLiteGeoPriorModelElev
+from lib.geo_inferrer_tf import TFGeoPriorModelElev
+
+
+class GeoInferrerFactory:
+    @staticmethod
+    def create(model_path: str) -> GeoInferrer:
+        if "mlmodel" in model_path:
+            assert platform == "darwin", "CoreML models can only be used on macOS"
+            return CoremlGeoPriorModelElev(model_path)
+        elif "tflite" in model_path:
+            return TFLiteGeoPriorModelElev(model_path)
+        elif "h5" in model_path:
+            return TFGeoPriorModelElev(model_path)
+        else:
+            raise ValueError(f"Unsupported model format in path: {model_path}")

lib/tf_gp_elev_model.py → lib/geo_inferrer_tf.py

@@ -1,31 +1,34 @@
+import os
+
 import tensorflow as tf
 import numpy as np
-import math
-import os
+
 from lib.res_layer import ResLayer
+from lib.geo_inferrer import GeoInferrer
 
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 
 
-class TFGeoPriorModelElev:
+class TFGeoPriorModelElev(GeoInferrer):
 
-    def __init__(self, model_path):
+    def __init__(self, model_path: str):
         # initialize the geo model for inference
         tf.config.set_visible_devices([], "GPU")
         visible_devices = tf.config.get_visible_devices()
         for device in visible_devices:
             assert device.device_type != "GPU"
         self.gpmodel = tf.keras.models.load_model(
             model_path,
-            custom_objects={"ResLayer": ResLayer},
+            custom_objects={"ResLayer": ResLayer}, 
             compile=False
         )
 
-    def predict(self, latitude, longitude, elevation):
-        encoded_loc = TFGeoPriorModelElev.encode_loc([latitude], [longitude], [elevation])
-        return self.gpmodel(tf.convert_to_tensor(
+    def predict(self, latitude: float, longitude: float, elevation: float) -> np.ndarray:
+        encoded_loc = GeoInferrer.encode_loc([latitude], [longitude], [elevation])
+        output = self.gpmodel(tf.convert_to_tensor(
             tf.expand_dims(encoded_loc[0], axis=0)
         ), training=False)[0]
+        return output
 
     def features_for_one_class_elevation(self, latitude, longitude, elevation):
         """Evalutes the model for a single class and multiple locations
@@ -60,26 +63,3 @@ def eval_one_class_elevation_from_features(self, features, class_of_interest):
                 transpose_b=True
             )
         ).numpy()
-
-    @staticmethod
-    def encode_loc(latitude, longitude, elevation):
-        latitude = np.array(latitude)
-        longitude = np.array(longitude)
-        elevation = np.array(elevation)
-        elevation = elevation.astype("float32")
-        grid_lon = longitude.astype("float32") / 180.0
-        grid_lat = latitude.astype("float32") / 90.0
-
-        elevation[elevation > 0] = elevation[elevation > 0] / 6574.0
-        elevation[elevation < 0] = elevation[elevation < 0] / 32768.0
-        norm_elev = elevation
-
-        norm_loc = tf.stack([grid_lon, grid_lat], axis=1)
-
-        encoded_loc = tf.concat([
-            tf.sin(norm_loc * math.pi),
-            tf.cos(norm_loc * math.pi),
-            tf.expand_dims(norm_elev, axis=1),
-
-        ], axis=1)
-        return encoded_loc

lib/geo_inferrer_tflite.py

@@ -0,0 +1,33 @@
+import numpy as np
+import tensorflow as tf
+
+from lib.geo_inferrer import GeoInferrer
+
+
+class TFLiteGeoPriorModelElev(GeoInferrer):
+
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.interpreter = tf.lite.Interpreter(model_path=self.model_path)
+        self.interpreter.allocate_tensors()
+
+    def predict(
+        self, latitude: float, longitude: float, elevation: float
+    ) -> np.ndarray:
+        encoded_loc = GeoInferrer.encode_loc(
+            [latitude], [longitude], [elevation]
+        ).numpy()
+
+        input_details = self.interpreter.get_input_details()
+        output_details = self.interpreter.get_output_details()
+
+        input_dtype = input_details[0]["dtype"]
+        encoded_loc = encoded_loc.astype(input_dtype)
+
+        self.interpreter.set_tensor(
+            input_details[0]["index"],
+            encoded_loc,
+        )
+        self.interpreter.invoke()
+        output_data = self.interpreter.get_tensor(output_details[0]["index"])
+        return output_data[0]

lib/inat_inferrer.py

@@ -16,8 +16,8 @@
 import asyncio
 
 from PIL import Image
-from lib.tf_gp_elev_model import TFGeoPriorModelElev
-from lib.vision_inferrer import VisionInferrer
+from lib.geo_inferrer_factory import GeoInferrerFactory
+from lib.vision_inferrer_factory import VisionInferrerFactory
 from lib.model_taxonomy_dataframe import ModelTaxonomyDataframe
 
 pd.options.mode.copy_on_write = True
@@ -140,7 +140,7 @@ def setup_synonym_taxonomy(self):
         self.taxonomy = synonym_taxonomy
 
     def setup_vision_model(self):
-        self.vision_inferrer = VisionInferrer(
+        self.vision_inferrer = VisionInferrerFactory.create(
             self.config["vision_model_path"]
         )
 
@@ -184,13 +184,18 @@ def setup_geo_model(self):
         if self.geo_elevation_cells is None:
             return
 
-        self.geo_elevation_model = TFGeoPriorModelElev(self.config["tf_geo_elevation_model_path"])
-        self.geo_model_features = self.geo_elevation_model.features_for_one_class_elevation(
-            latitude=list(self.geo_elevation_cells.lat),
-            longitude=list(self.geo_elevation_cells.lng),
-            elevation=list(self.geo_elevation_cells.elevation)
+        self.geo_elevation_model = GeoInferrerFactory.create(
+            self.config["tf_geo_elevation_model_path"]
         )
 
+        if hasattr(self.geo_elevation_model, "features_for_one_class_elevation"):
+            self.geo_model_features = self.geo_elevation_model.features_for_one_class_elevation(
+                latitude=list(self.geo_elevation_cells.lat),
+                longitude=list(self.geo_elevation_cells.lng),
+                elevation=list(self.geo_elevation_cells.elevation)
+            )
+
+
     def vision_predict(self, image, debug=False):
         if debug:
             start_time = time.time()

lib/vision_inferrer.py

@@ -1,35 +1,36 @@
+from abc import ABC, abstractmethod
+from typing import Optional, TypedDict
+
+import numpy as np
 import tensorflow as tf
 
 
-class VisionInferrer:
-
-    def __init__(self, model_path):
-        self.model_path = model_path
-        self.prepare_tf_model()
-
-    # initialize the TF model given the configured path
-    def prepare_tf_model(self):
-        # disable GPU processing
-        tf.config.set_visible_devices([], "GPU")
-        visible_devices = tf.config.get_visible_devices()
-        for device in visible_devices:
-            assert device.device_type != "GPU"
-
-        full_model = tf.keras.models.load_model(self.model_path, compile=False)
-        self.layered_model = tf.keras.Model(
-            inputs=full_model.inputs,
-            outputs=[
-                full_model.layers[4].output,
-                full_model.layers[2].output
-            ]
-        )
-        self.layered_model.compile()
-
-    # given an image object (usually coming from prepare_image_for_inference),
-    # calculate vision results for the image
-    def process_image(self, image):
-        layer_results = self.layered_model(tf.convert_to_tensor(image), training=False)
-        return {
-            "predictions": layer_results[0][0],
-            "features": layer_results[1][0],
-        }
+class VisionResults(TypedDict):
+    predictions: np.ndarray
+    features: Optional[np.ndarray]
+
+
+class VisionInferrer(ABC):
+    @abstractmethod
+    def __init__(self, model_path: str):
+        """Subclasses must implement this constructor."""
+        pass
+
+    @abstractmethod
+    def prepare_model(self):
+        """
+        Initialize the model.
+
+        Subclasses must implement this method.
+        """
+        pass
+
+    @abstractmethod
+    def process_image(self, image: tf.Tensor) -> VisionResults:
+        """
+        given an image object (usually coming from prepare_image_for_inference),
+        calculate vision results for the image
+
+        Subclasses must implement this method.
+        """
+        pass

lib/vision_inferrer_coreml.py

@@ -0,0 +1,42 @@
+import coremltools as ct
+from PIL import Image
+import tensorflow as tf
+
+from lib.vision_inferrer import VisionInferrer, VisionResults
+
+
+class VisionInferrerCoreML(VisionInferrer):
+    """Vision Inferrer for the CoreML variant of iNat vision models.
+    Our implementation expects a single PIL image in the range [0, 255).
+    """
+
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.prepare_model()
+
+    def prepare_model(self):
+        """initialize the CoreML model given the configured path"""
+        self.model = ct.models.MLModel(self.model_path)
+        spec = self.model.get_spec()
+        self.input_name = spec.description.input[0].name
+
+    def process_image(self, image_tensor: tf.Tensor) -> VisionResults:
+        """given an image object (coming from prepare_image_for_inference),
+        calculate & return vision results for the image."""
+        # coreml expects a PIL image so we have to convert from tf
+        # first we convert from floats [0, 1) to ints [0, 255)
+        image = tf.image.convert_image_dtype(image_tensor, dtype=tf.uint8)
+
+        # Remove batch dimension if present and convert to NumPy array
+        image_numpy = image.numpy()
+        if image_numpy.ndim == 4:
+            image_numpy = image_numpy[0]
+
+        # Create PIL Image from NumPy array
+        image_pil = Image.fromarray(image_numpy)
+
+        out_dict = self.model.predict({self.input_name: image_pil})
+        preds = out_dict["Identity"][0]
+
+        # don't return features, not relevant for coreml at this point
+        return {"predictions": preds, "features": None}

lib/vision_inferrer_factory.py

@@ -0,0 +1,20 @@
+from sys import platform
+
+from lib.vision_inferrer import VisionInferrer
+from lib.vision_inferrer_coreml import VisionInferrerCoreML
+from lib.vision_inferrer_tflite import VisionInferrerTFLite
+from lib.vision_inferrer_tf import VisionInferrerTF
+
+
+class VisionInferrerFactory:
+    @staticmethod
+    def create(model_path: str) -> VisionInferrer:
+        if "mlmodel" in model_path:
+            assert platform == "darwin", "CoreML models can only be used on macOS"
+            return VisionInferrerCoreML(model_path)
+        elif "tflite" in model_path:
+            return VisionInferrerTFLite(model_path)
+        elif "h5" in model_path:
+            return VisionInferrerTF(model_path)
+        else:
+            raise ValueError(f"Unsupported model format in path: {model_path}")