diff --git a/README.md b/README.md index 38fdd2a..864a22b 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,9 @@ -# lobe-python -Code to run exported Lobe models in Python. +# Lobe Python API +Code to run exported Lobe models in Python using the TensorFlow, TensorFlow Lite, or ONNX options. ## Install -``` +### Linux +```shell script # Install Python3 sudo apt update sudo apt install -y python3-dev python3-pip @@ -13,15 +14,36 @@ sudo apt install -y \ libatlas-base-dev \ libopenjp2-7 \ libtiff5 \ - libjpeg62-turbo + libjpeg62-dev # Install lobe-python pip3 install setuptools git pip3 install git+https://github.com/lobe/lobe-python ``` -## Usage +_Note for Raspbian OS (Raspberry Pi)_: Please install `libjpeg62-turbo` instead of `libjpeg62-dev` + +### Mac/Windows +Use a virtual environment with Python 3.7 +```shell script +python3 -m venv .venv + +# Mac: +source .venv/bin/activate + +# Windows: +.venv\Scripts\activate +``` +Install the library +```shell script +# make sure pip is up to date +python -m pip install --upgrade pip +# install +pip install git+https://github.com/lobe/lobe-python ``` + +## Usage +```python from lobe import ImageModel model = ImageModel.load('path/to/exported/model') @@ -30,7 +52,7 @@ model = ImageModel.load('path/to/exported/model') result = model.predict_from_file('path/to/file.jpg') # OPTION 2: Predict from an image url -result = model.predict_from_url('http://path/to/file.jpg') +result = model.predict_from_url('http://url/to/file.jpg') # OPTION 3: Predict from Pillow image from PIL import Image @@ -41,21 +63,12 @@ result = model.predict(img) print(result.prediction) # Print all classes -for label, prop in result.labels: - print(f"{label}: {prop*100}%") +for label, confidence in result.labels: + print(f"{label}: {confidence*100}%") ``` +Note: model predict functions should be thread-safe. If you find bugs please file an issue. ## Resources -If you're running this on a Pi and having issues, and seeing this error: - -```bash -Could not install packages due to an EnvironmentError: 404 Client Error: Not Found for url: https://pypi.org/simple/tflite-runtime/ -``` - -running this may help: - -```bash -pip3 install https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl -``` +See the [Raspberry Pi Trash Classifier](https://github.com/microsoft/TrashClassifier) example, and its [Adafruit Tutorial](https://learn.adafruit.com/lobe-trash-classifier-machine-learning). diff --git a/RELEASE.md b/RELEASE.md new file mode 100644 index 0000000..a14e42a --- /dev/null +++ b/RELEASE.md @@ -0,0 +1,12 @@ +# Release 0.3.0 +___ +## Breaking Changes +* Previous use of Signature should be ImageClassificationSignature. `from lobe.signature import Signature` -> + `from lobe.signature import ImageClassificationSignature` + +## Bug Fixes and Other Improvements +* Update to TensorFlow 2.4 from 1.15.4 +* Add ONNX runtime backend +* Use requests instead of urllib +* Make backends thread-safe +* Added constants file for signature keys to enable backwards-compatibility \ No newline at end of file diff --git a/examples/basic_usage.py b/examples/basic_usage.py index 1577734..12a8b6e 100644 --- a/examples/basic_usage.py +++ b/examples/basic_usage.py @@ -7,7 +7,7 @@ result = model.predict_from_file('path/to/file.jpg') # Predict from an image url -result = model.predict_from_url('http://path/to/file.jpg') +result = model.predict_from_url('http://url/to/file.jpg') # Predict from Pillow image from PIL import Image @@ -18,5 +18,5 @@ print("Top prediction:", result.prediction) # Print all classes -for label, prop in result.labels: - print(f"{label}: {prop*100:.6f}%") +for label, confidence in result.labels: + print(f"{label}: {confidence*100:.6f}%") diff --git a/setup.py b/setup.py index 3769810..5c0b2c6 100644 --- a/setup.py +++ b/setup.py @@ -1,18 +1,55 @@ from setuptools import setup, find_packages +import sys +import platform + + +python_version = platform.python_version().rsplit('.', maxsplit=1)[0] + +requirements = [ + "pillow", + "requests", + "numpy==1.19.3", + "tensorflow==2.4;platform_machine!='armv7l'", + "onnxruntime==1.6.0;platform_machine!='armv7l'" +] + +# get the right TF Lite runtime packages based on OS and python version: https://www.tensorflow.org/lite/guide/python#install_just_the_tensorflow_lite_interpreter +tflite_python = None +tflite_machine = None + +# get the right python string for the version +if python_version == '3.5': + tflite_python = 'cp35-cp35m' +elif python_version == '3.6': + tflite_python = 'cp36-cp36m' +elif python_version == '3.7': + tflite_python = 'cp37-cp37m' +elif python_version == '3.8': + tflite_python = 'cp38-cp38' + +# get the right machine string +if sys.platform == 'win32': + tflite_machine = 'win_amd64' +elif sys.platform == 'darwin': + tflite_machine = 'macosx_10_15_x86_64' +elif sys.platform == 'linux': + if platform.machine() == 'x86_64': + tflite_machine = 'linux_x86_64' + elif platform.machine() == 'armv7l': + tflite_machine = 'linux_armv7l' + +# add it to the requirements, or print the location to find the version to install +if tflite_python and tflite_machine: + requirements.append(f"tflite_runtime @ https://github.com/google-coral/pycoral/releases/download/release-frogfish/tflite_runtime-2.5.0-{tflite_python}-{tflite_machine}.whl") +else: + print( + f"Couldn't find tflite_runtime for your platform {sys.platform}, machine {platform.machine()}, and python version {python_version}, please see the install guide for the right version: https://www.tensorflow.org/lite/guide/python#install_just_the_tensorflow_lite_interpreter" + ) + setup( name="lobe", - version="0.2.1", + version="0.3.0", packages=find_packages("src"), package_dir={"": "src"}, - install_requires=[ - "numpy", - "pillow", - "requests", - "tensorflow>=1.15.2,<2;platform_machine!='armv7l'", - "tflite_runtime ; platform_machine=='armv7l'" - ], - dependency_links=[ - "https://www.piwheels.org/simple/tensorflow", - "https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl" - ] + install_requires=requirements, ) diff --git a/src/lobe/ImageModel.py b/src/lobe/ImageModel.py deleted file mode 100644 index 5d40ff9..0000000 --- a/src/lobe/ImageModel.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Load a Lobe saved model -""" -from __future__ import annotations -import os -import json -from PIL import Image -from typing import Tuple - -from ._results import PredictionResult -from ._model import Model -from . import Signature -from . import image_utils - -def load_from_signature(signature: Signature) -> ImageModel: - # Select the appropriate backend - model_format = signature.format - if model_format == "tf": - from .backends import _backend_tf as backend - elif model_format == "tf_lite": - from .backends import _backend_tflite as backend - else: - raise ValueError("Model is an unsupported format") - - backend_predict = backend.ImageClassificationModel(signature) - return ImageModel(signature, backend_predict) - -def load(model_path: str) -> ImageModel: - # Load the signature - signature = Signature.load(model_path) - - return load_from_signature(signature) - -class ImageModel(Model): - def __init__(self, signature, backend): - super(ImageModel, self).__init__(signature) - self.__backend = backend - - def predict_from_url(self, url: str): - return self.predict(image_utils.get_image_from_url(url)) - - def predict_from_file(self, path: str): - return self.predict(image_utils.get_image_from_file(path)) - - def predict(self, image: Image.Image): - image_processed = image_utils.preprocess_image(image, self.signature.input_image_size) - return self.__backend.predict(image_processed) \ No newline at end of file diff --git a/src/lobe/Signature.py b/src/lobe/Signature.py deleted file mode 100644 index 9958f07..0000000 --- a/src/lobe/Signature.py +++ /dev/null @@ -1,89 +0,0 @@ -from __future__ import annotations -import os -import json -import pathlib -from typing import Tuple - -def load(model_path: str) -> Signature: - """ - Loads the Signature for the given PATH or PATH_AND_FILENAME provided: - - Use PATH When: Using Lobe-Python in its default config, the Signature and TensorFlow (and TFLite) models are expected to be - in one folder by themselves. Additional models may exist, but they too are expected to be in their own folders. - - Use PATH_AND_FILENAME When: Using Lobe-Python with multiple TensorFlow (and TFLite) models in the same folder, with - the Signature and Model files named uniquely. This allows you to store multiple TensorFlow/TFLite models and signatures in the same folder. - """ - model_path_real = os.path.realpath(os.path.expanduser(model_path)) - - #This could be a full_path to the signature file - if os.path.isfile(model_path_real): - filename, extension = os.path.splitext(model_path_real) - if (extension.lower() != ".json" ): #Signature file must end in "json" - raise ValueError(f"Model file provided is not valid: {model_path_real}") - signature_path = model_path_real #We have the signature file, so load the model - elif os.path.isdir(model_path_real): - #This is a directory with a single Signature File to load - signature_path = os.path.join(model_path_real, "signature.json") - if not os.path.isfile(signature_path): - raise ValueError(f"signature.json file not found at path: {model_path}") - else: - raise ValueError(f"Invalid Signature file or Model directory: {model_path}") - - return Signature(signature_path) - -class Signature: - def __init__(self, signature_path: str): - signature_path = pathlib.Path(signature_path) - self.__model_path = str(signature_path.parent) - - with open(signature_path, "r") as f: - self.__signature = json.load(f) - - inputs = self.__signature.get("inputs") - input_tensor_shape = inputs["Image"]["shape"] - assert len(input_tensor_shape) == 4 - self.__input_image_size = (input_tensor_shape[1], input_tensor_shape[2]) - - @property - def model_path(self) -> str: - return self.__model_path - - @property - def id(self) -> str: - return self.__signature.get("doc_id") - - @property - def name(self) -> str: - return self.__signature.get("doc_name") - - @property - def version(self) -> str: - return self.__signature.get("doc_version") - - @property - def format(self) -> str: - return self.__signature.get("format") - - @property - def filename(self) -> str: - return self.__signature.get("filename") - - @property - def input_tensor_shape(self) -> str: - return self.__input_image_size - - @property - def classes(self): - classes = [] - if self.__signature.get("classes", None): - classes = self.__signature.get("classes").get("Label") - return classes - - @property - def input_image_size(self) -> Tuple[int, int]: - return self.__input_image_size - - def as_dict(self): - return self.__signature - - def __str__(self): - return json.dumps(self.as_dict()) \ No newline at end of file diff --git a/src/lobe/__init__.py b/src/lobe/__init__.py index e69de29..9e5bab9 100644 --- a/src/lobe/__init__.py +++ b/src/lobe/__init__.py @@ -0,0 +1,2 @@ +from .signature import Signature +from .model.image_model import ImageModel diff --git a/src/lobe/_model.py b/src/lobe/_model.py deleted file mode 100644 index d60cb2f..0000000 --- a/src/lobe/_model.py +++ /dev/null @@ -1,24 +0,0 @@ -from __future__ import annotations -from abc import ABC, abstractmethod -from .Signature import Signature -from ._results import PredictionResult - -class Model(ABC): - def __init__(self, signature: Signature) -> Model: - self.__signature = signature - - @property - def signature(self) -> Signature: - return self.__signature - - @abstractmethod - def predict_from_url(self, url: str): - pass - - @abstractmethod - def predict_from_file(self, path: str): - pass - - @abstractmethod - def predict(self, input) -> PredictionResult: - pass \ No newline at end of file diff --git a/src/lobe/_results.py b/src/lobe/_results.py deleted file mode 100644 index 4a60bdf..0000000 --- a/src/lobe/_results.py +++ /dev/null @@ -1,50 +0,0 @@ -from __future__ import annotations -import json - - -class PredictionResult(): - @classmethod - def from_dict(cls, results) -> PredictionResult: - outputs = results.get('outputs') - labels = outputs.get('Labels') - prediction = outputs.get('Prediction')[0] - return cls(labels, prediction) - - @classmethod - def from_json(cls, json_str: str) -> PredictionResult: - results = json.loads(json_str) - return PredictionResult.from_dict(results) - - @staticmethod - def __sort_predictions(labels, confidences): - top_predictions = confidences.argsort()[-5:][::-1] - labels = labels - sorted_labels = [] - for i in top_predictions: - sorted_labels.append(labels[i]) - return sorted_labels, confidences[top_predictions] - - def __init__(self, labels, confidences, prediction: str): - sorted_labels, sorted_confidences = self.__sort_predictions(labels, confidences) - self.__labels = list(zip(sorted_labels, sorted_confidences)) - self.__prediction = prediction - - @property - def labels(self): - return self.__labels - - @property - def prediction(self) -> str: - return self.__prediction - - def as_dict(self): - output = { - "outputs": { - "Labels": self.labels, - "Prediction": [self.prediction], - } - } - return output - - def __str__(self) -> str: - return json.dumps(self.as_dict()) diff --git a/src/lobe/api_client.py b/src/lobe/api_client.py index 487c7d2..24a0120 100644 --- a/src/lobe/api_client.py +++ b/src/lobe/api_client.py @@ -1,15 +1,15 @@ #!/usr/bin/env python import requests from PIL import Image -from image_utils import image_to_base64 -from ._results import PredictionResult -def send_image_predict_request(image: Image.Image, predict_url: str, key: str) -> PredictionResult: +from .image_utils import image_to_base64 +from .results import ClassificationResult + + +def send_image_predict_request(image: Image.Image, predict_url: str) -> ClassificationResult: payload = { "inputs": {"Image": image_to_base64(image)}, - "key": key } - headers = {'authorization': f'Bearer {key}'} - response = requests.post(predict_url, json=payload, headers=headers) + response = requests.post(predict_url, json=payload) response.raise_for_status() - return PredictionResult.from_json(response.text) + return ClassificationResult.from_json(response.text) diff --git a/src/lobe/backends/_backend_tf.py b/src/lobe/backends/_backend_tf.py deleted file mode 100644 index 724ad0b..0000000 --- a/src/lobe/backends/_backend_tf.py +++ /dev/null @@ -1,42 +0,0 @@ -import numpy as np - -from PIL import Image -from .._results import PredictionResult - -try: - import tensorflow as tf - from tensorflow.contrib import predictor -except ImportError: - raise ImportError("ERROR: This is a TensorFlow model and requires tensorflow to be installed on this device. Please run\n\tpip install tensorflow\n") - -class ImageClassificationModel(): - __input_key_image = 'Image' - __output_key_confidences = 'Confidences' - __output_key_prediction = 'Prediction' - - def __init__(self, signature): - self.__model_path = signature.model_path - self.__tf_predict_fn = None - self.__labels = signature.classes - - def __load(self): - self.__tf_predict_fn = predictor.from_saved_model(self.__model_path) - - def predict(self, image: Image.Image) -> PredictionResult: - if self.__tf_predict_fn is None: - self.__load() - - # Convert all values to range [0,1] - np_image = np.asarray(image) / 255. - - # Add an extra axis onto the numpy array - np_image = np_image[np.newaxis, ...] - - predictions = self.__tf_predict_fn({ - self.__input_key_image: np_image - }) - - confidences = predictions[self.__output_key_confidences][0] - top_prediction = predictions[self.__output_key_prediction][0].decode('utf-8') - - return PredictionResult(labels=self.__labels, confidences=confidences, prediction=top_prediction) \ No newline at end of file diff --git a/src/lobe/backends/_backend_tflite.py b/src/lobe/backends/_backend_tflite.py deleted file mode 100644 index c92354c..0000000 --- a/src/lobe/backends/_backend_tflite.py +++ /dev/null @@ -1,67 +0,0 @@ -import numpy as np -from PIL import Image -from .._results import PredictionResult - -try: - import tflite_runtime.interpreter as tflite - -except ImportError: - # Needs better error text - raise ImportError( - "ERROR: This is a TensorFlow Lite model and requires TensorFlow Lite interpreter to be installed on this device. Please go to https://www.tensorflow.org/lite/guide/python and download the appropriate version for you device." - ) - - -class ImageClassificationModel: - __MAX_UINT8 = 255 - - def __init__(self, signature): - self.__model_path = "{}/{}".format( - signature.model_path, signature.filename - ) - self.__tflite_predict_fn = None - self.__labels = signature.classes - - - def __load(self): - self.__tflite_predict_fn = tflite.Interpreter( - model_path=self.__model_path - ) - - def predict(self, image: Image.Image) -> PredictionResult: - if self.__tflite_predict_fn is None: - self.__load() - - self.__tflite_predict_fn.allocate_tensors() - - # Add an extra axis onto the numpy array - np_image = np.expand_dims(image, axis=0) - - # Converts to floating point and standardize range from 0 to 1. - np_image = np.float32(np_image) / self.__MAX_UINT8 - - input_details = self.__tflite_predict_fn.get_input_details() - output_details = self.__tflite_predict_fn.get_output_details() - - self.__tflite_predict_fn.set_tensor( - input_details[0]["index"], - np_image - ) - - self.__tflite_predict_fn.invoke() - top_prediction_output = self.__tflite_predict_fn.get_tensor( - output_details[0]["index"] - ) - - confidences_output = self.__tflite_predict_fn.get_tensor( - output_details[1]["index"] - ) - - confidences = np.squeeze(confidences_output) - top_prediction = top_prediction_output.item().decode("utf-8") - - return PredictionResult( - labels=self.__labels, - confidences=confidences, - prediction=top_prediction, - ) diff --git a/src/lobe/backends/backend_onnx.py b/src/lobe/backends/backend_onnx.py new file mode 100644 index 0000000..38294b7 --- /dev/null +++ b/src/lobe/backends/backend_onnx.py @@ -0,0 +1,71 @@ +from threading import Lock +from ..signature import Signature +from ..signature_constants import TENSOR_NAME +from ..utils import decode_dict_bytes_as_str + +try: + import onnxruntime as rt + +except ImportError: + # Needs better error text + raise ImportError( + "ERROR: This is an ONNX model and requires onnx runtime to be installed on this device. Please go to https://www.onnxruntime.ai/ for install instructions." + ) + + +class ONNXModel(object): + """ + Generic wrapper for running an ONNX model exported from Lobe + """ + def __init__(self, signature: Signature): + model_path = "{}/{}".format( + signature.model_path, signature.filename + ) + self.signature = signature + + # load our onnx inference session + self.session = rt.InferenceSession(path_or_bytes=model_path) + + self.lock = Lock() + + def predict(self, data): + """ + Predict the outputs by running the data through the model. + + data: can be either a single input value (such as an image array), or a dictionary mapping the input + keys from the signature to the data they should be assigned + + Returns a dictionary in the form of the signature outputs {Name: value, ...} + """ + # make the predict function thread-safe + with self.lock: + # create the feed dictionary that is the input to the model + feed_dict = {} + # either map the input data names to the appropriate tensors from the signature inputs, or map to the first + # input if we are just given a non-dictionary input + if not isinstance(data, dict): + # if data isn't a dictionary, set the input to the supplied value + # throw an error if more than 1 input found and we are only supplied a non-dictionary input + if len(self.signature.inputs) > 1: + raise ValueError( + f"Found more than 1 model input: {list(self.signature.inputs.keys())}, while supplied data wasn't a dictionary: {data}" + ) + feed_dict[list(self.signature.inputs.values())[0].get(TENSOR_NAME)] = data + else: + # otherwise, assign data to inputs based on the dictionary + for input_name, input_sig in self.signature.inputs.items(): + if input_name not in data: + raise ValueError(f"Couldn't find input {input_name} in the supplied data {data}") + feed_dict[input_sig.get(TENSOR_NAME)] = data.get(input_name) + + # run the model! + # get the outputs + fetches = [(key, value.get("name")) for key, value in self.signature.outputs.items()] + outputs = self.session.run(output_names=[name for (_, name) in fetches], input_feed=feed_dict) + # make our return a dict from the list of outputs that correspond to the fetches + results = {} + for i, (key, _) in enumerate(fetches): + results[key] = outputs[i].tolist() + # postprocessing! convert any byte strings to normal strings with .decode() + decode_dict_bytes_as_str(results) + return results diff --git a/src/lobe/backends/backend_tf.py b/src/lobe/backends/backend_tf.py new file mode 100644 index 0000000..8dfdfb0 --- /dev/null +++ b/src/lobe/backends/backend_tf.py @@ -0,0 +1,60 @@ +from threading import Lock + +from ..signature import Signature +from ..utils import decode_dict_bytes_as_str + +try: + import tensorflow as tf +except ImportError: + raise ImportError("ERROR: This is a TensorFlow model and requires tensorflow to be installed on this device. Please run\n\tpip install tensorflow==2.4\n") + + +class TFModel(object): + """ + Generic wrapper for running a tensorflow model + """ + def __init__(self, signature: Signature): + self.lock = Lock() + self.signature = signature + + self.model = tf.saved_model.load(export_dir=self.signature.model_path, tags=self.signature.tags) + self.predict_fn = self.model.signatures['serving_default'] + + def predict(self, data): + """ + Predict the outputs by running the data through the model. + + data: can be either a single input value (such as an image array), or a dictionary mapping the input + keys from the signature to the data they should be assigned + + Returns a dictionary in the form of the signature outputs {Name: value, ...} + """ + with self.lock: + # create the feed dictionary that is the input to the model + feed_dict = {} + # either map the input data names to the appropriate tensors from the signature inputs, or map to the first + # input if we are just given a non-dictionary input + if not isinstance(data, dict): + # if data isn't a dictionary, set the input to the supplied value + # throw an error if more than 1 input found and we are only supplied a non-dictionary input + if len(self.signature.inputs) > 1: + raise ValueError( + f"Found more than 1 model input: {list(self.signature.inputs.keys())}, while supplied data wasn't a dictionary: {data}" + ) + feed_dict[list(self.signature.inputs.keys())[0]] = tf.convert_to_tensor(data) + else: + # otherwise, assign data to inputs based on the dictionary + for input_name in self.signature.inputs.keys(): + if input_name not in data: + raise ValueError(f"Couldn't find input {input_name} in the supplied data {data}") + feed_dict[input_name] = tf.convert_to_tensor(data.get(input_name)) + + # run the model! there will be as many outputs from session.run as you have in the fetches list + outputs = self.predict_fn(**feed_dict) + + # postprocessing! make our output dictionary and convert any byte strings to normal strings with .decode() + results = {} + for i, (key, tf_val) in enumerate(outputs.items()): + results[key] = tf_val.numpy().tolist() + decode_dict_bytes_as_str(results) + return results diff --git a/src/lobe/backends/backend_tflite.py b/src/lobe/backends/backend_tflite.py new file mode 100644 index 0000000..4cd1879 --- /dev/null +++ b/src/lobe/backends/backend_tflite.py @@ -0,0 +1,81 @@ +from threading import Lock +from ..signature import Signature +from ..signature_constants import TENSOR_NAME +from ..utils import decode_dict_bytes_as_str + +try: + import tflite_runtime.interpreter as tflite + +except ImportError: + # Needs better error text + raise ImportError( + "ERROR: This is a TensorFlow Lite model and requires TensorFlow Lite interpreter to be installed on this device. Please go to https://www.tensorflow.org/lite/guide/python and download the appropriate version for you device." + ) + + +class TFLiteModel(object): + """ + Generic wrapper for running a TF Lite model exported from Lobe + """ + def __init__(self, signature: Signature): + model_path = "{}/{}".format( + signature.model_path, signature.filename + ) + self.signature = signature + self.interpreter = tflite.Interpreter(model_path=model_path) + self.interpreter.allocate_tensors() + + # Combine the information about the inputs and outputs from the signature.json file + # with the Interpreter runtime details + input_details = {detail.get("name"): detail for detail in self.interpreter.get_input_details()} + self.model_inputs = { + key: {**sig, **input_details.get(sig.get(TENSOR_NAME))} + for key, sig in self.signature.inputs.items() + } + output_details = {detail.get("name"): detail for detail in self.interpreter.get_output_details()} + self.model_outputs = { + key: {**sig, **output_details.get(sig.get(TENSOR_NAME))} + for key, sig in self.signature.outputs.items() + } + self.lock = Lock() + + def predict(self, data): + """ + Predict the outputs by running the data through the model. + + data: can be either a single input value (such as an image array), or a dictionary mapping the input + keys from the signature to the data they should be assigned + + Returns a dictionary in the form of the signature outputs {Name: value, ...} + """ + # make the predict function thread-safe + with self.lock: + # set the model inputs with our supplied data + if not isinstance(data, dict): + # if data isn't a dictionary, set the input to the supplied value + # throw an error if more than 1 input found and we are only supplied a non-dictionary input + if len(self.model_inputs) > 1: + raise ValueError( + f"Found more than 1 model input: {list(self.model_inputs.keys())}, while supplied data wasn't a dictionary: {data}" + ) + self.interpreter.set_tensor(list(self.model_inputs.values())[0].get("index"), data) + else: + # otherwise, assign data to inputs based on the dictionary + for input_name, input_detail in self.model_inputs.items(): + if input_name not in data: + raise ValueError(f"Couldn't find input {input_name} in the supplied data {data}") + self.interpreter.set_tensor(input_detail.get("index"), data.get(input_name)) + + # invoke the interpreter -- runs the model with the set inputs + self.interpreter.invoke() + + # grab our desired outputs from the interpreter + # convert to normal python types with tolist() + outputs = { + key: self.interpreter.get_tensor(value.get("index")).tolist() + for key, value in self.model_outputs.items() + } + + # postprocessing! convert any byte strings to normal strings with .decode() + decode_dict_bytes_as_str(outputs) + return outputs diff --git a/src/lobe/image_utils.py b/src/lobe/image_utils.py index db64ec4..7f1774e 100644 --- a/src/lobe/image_utils.py +++ b/src/lobe/image_utils.py @@ -1,10 +1,11 @@ -import sys from io import BytesIO from PIL import Image +import numpy as np from typing import Tuple -import urllib +import requests import base64 + def crop_center(image: Image.Image, size: Tuple[int, int]) -> Image.Image: crop_width, crop_height = size width, height = image.size @@ -14,12 +15,14 @@ def crop_center(image: Image.Image, size: Tuple[int, int]) -> Image.Image: bottom = min(height, top + crop_height) return image.crop((left, top, right, bottom)) + def crop_center_square(image: Image.Image, size: Tuple[int, int]=None) -> Image.Image: if not size: width, height = image.size size = min(width, height) return crop_center(image, (size, size)) + def resize_uniform_to_fill(image: Image.Image, size: Tuple[int, int]) -> Image.Image: width, height = image.size min_w, min_h = size @@ -29,6 +32,7 @@ def resize_uniform_to_fill(image: Image.Image, size: Tuple[int, int]) -> Image.I new_size = (round(scale * width), round(scale * height)) return image.resize(new_size) + def resize_uniform_to_fit(image: Image.Image, size: Tuple[int, int]) -> Image.Image: width, height = image.size max_w, max_h = size @@ -38,6 +42,7 @@ def resize_uniform_to_fit(image: Image.Image, size: Tuple[int, int]) -> Image.Im new_size = (round(scale * width), round(scale * height)) return image.resize(new_size) + def update_orientation(image: Image.Image) -> Image.Image: exif_orientation_tag = 0x0112 if hasattr(image, '_getexif'): @@ -54,21 +59,28 @@ def update_orientation(image: Image.Image) -> Image.Image: image = image.transpose(Image.FLIP_LEFT_RIGHT) return image + def ensure_rgb_format(image: Image.Image) -> Image.Image: return image.convert("RGB") + def image_to_base64(image: Image.Image) -> str: buffer = BytesIO() ensure_rgb_format(image).save(buffer,format="JPEG") return base64.b64encode(buffer.getvalue()).decode("utf-8") + def get_image_from_url(url: str) -> Image.Image: - path = BytesIO(urllib.request.urlopen(url).read()) - return Image.open(path) + response = requests.get(url) + response.raise_for_status() + image = Image.open(BytesIO(response.content)) + return image + def get_image_from_file(path: str) -> Image.Image: return Image.open(path) + def preprocess_image(image: Image.Image, size: Tuple[int, int]) -> Image.Image: image_processed = update_orientation(image) @@ -76,4 +88,12 @@ def preprocess_image(image: Image.Image, size: Tuple[int, int]) -> Image.Image: image_processed = ensure_rgb_format(image) image_processed = resize_uniform_to_fill(image_processed, size) image_processed = crop_center(image_processed, size) - return image_processed \ No newline at end of file + return image_processed + + +def image_to_array(image: Image.Image) -> np.ndarray: + # make 0-1 float instead of 0-255 int (that PIL Image loads by default) + image = np.asarray(image) / 255.0 + # pad with an extra batch dimension + image = np.expand_dims(image, axis=0).astype(np.float32) + return image diff --git a/src/lobe/model/__init__.py b/src/lobe/model/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/lobe/model/image_model.py b/src/lobe/model/image_model.py new file mode 100644 index 0000000..0ae5d91 --- /dev/null +++ b/src/lobe/model/image_model.py @@ -0,0 +1,52 @@ +""" +Load a Lobe saved model for image classification +""" +from __future__ import annotations +from PIL import Image + +from .model import Model +from ..signature import ImageClassificationSignature +from .. import image_utils +from ..results import ClassificationResult + + +class ImageModel(Model): + signature: ImageClassificationSignature + + @classmethod + def load_from_signature(cls, signature: ImageClassificationSignature) -> ImageModel: + # Select the appropriate backend + model_format = signature.format + if model_format == "tf": + from ..backends.backend_tf import TFModel + return cls(signature, TFModel(signature)) + elif model_format == "tf_lite": + from ..backends.backend_tflite import TFLiteModel + return cls(signature, TFLiteModel(signature)) + elif model_format == "onnx": + from ..backends.backend_onnx import ONNXModel + return cls(signature, ONNXModel(signature)) + else: + raise ValueError(f"Model is an unsupported format: {model_format}") + + @classmethod + def load(cls, model_path: str) -> ImageModel: + # Load the signature + return cls.load_from_signature(ImageClassificationSignature(model_path)) + + def __init__(self, signature: ImageClassificationSignature, backend): + super(ImageModel, self).__init__(signature) + self.backend = backend + + def predict_from_url(self, url: str): + return self.predict(image_utils.get_image_from_url(url)) + + def predict_from_file(self, path: str): + return self.predict(image_utils.get_image_from_file(path)) + + def predict(self, image: Image.Image) -> ClassificationResult: + image_processed = image_utils.preprocess_image(image, self.signature.input_image_size) + image_array = image_utils.image_to_array(image_processed) + results = self.backend.predict(image_array) + classification_results = ClassificationResult(results=results, labels=self.signature.classes) + return classification_results diff --git a/src/lobe/model/model.py b/src/lobe/model/model.py new file mode 100644 index 0000000..cd1e8a0 --- /dev/null +++ b/src/lobe/model/model.py @@ -0,0 +1,20 @@ +from __future__ import annotations +from abc import ABC, abstractmethod +from ..signature import Signature + + +class Model(ABC): + def __init__(self, signature: Signature): + self.signature = signature + + @abstractmethod + def predict_from_url(self, url: str): + pass + + @abstractmethod + def predict_from_file(self, path: str): + pass + + @abstractmethod + def predict(self, input): + pass diff --git a/src/lobe/results.py b/src/lobe/results.py new file mode 100644 index 0000000..ecd136a --- /dev/null +++ b/src/lobe/results.py @@ -0,0 +1,113 @@ +from __future__ import annotations +import json +from typing import List, Dict + +from .signature_constants import ( + PREDICTED_LABEL, PREDICTED_LABEL_COMPAT, LABEL_CONFIDENCES, LABEL_CONFIDENCES_COMPAT, OUTPUTS +) +from .utils import dict_get_compat, list_or_tuple + + +class ClassificationResult: + """ + Data structure to expose the classification predicted result from running a Lobe model. + Exposes the top predicted label, as well as a list of tuples (label, confidence) sorted by highest confidence to lowest. + + Sorted list of predictions (label, confidence): ClassificationResult.labels + Top predicted label: ClassificationResult.prediction + + These can be batched and contain the results for many examples. + """ + + @classmethod + def from_json(cls, json_str: str, labels: List[str] = None) -> ClassificationResult: + """ + Parse the classification results from a json string + """ + results = json.loads(json_str) + return cls(results=results, labels=labels) + + def __init__(self, results: Dict[str, any], labels: List[str] = None): + """ + Parse the classification results from a dictionary in the form {Name: val} for each output in the signature + + Labels need to be provided to map our confidences, but in the case of the local API they are already returned + with the prediction. + """ + # first check if this results dictionary started with the 'outputs' top-level key and navigate inside if true + outputs_dict = results.get(OUTPUTS) + if outputs_dict is not None: + results = outputs_dict + + # grab the list of confidences -- this may or may not include labels + confidences, _ = dict_get_compat(in_dict=results, current_key=LABEL_CONFIDENCES, + compat_keys=LABEL_CONFIDENCES_COMPAT, default=[]) + # zip the labels and confidences together. labels can be None (they should already exist in confidences) + labels_and_confidences = [] + # if confidences were unbatched and already contain (label, confidence) pairs (as when returned from local API) + # just sort the confidences + if _is_label_conf_pair(confidences): + labels_and_confidences = sorted(confidences, key=lambda pair: pair[1], reverse=True) + else: + # otherwise, it could be batched -- go deeper! + for row in confidences: + # if this row is a list of numbers, zip it with the labels to make (label, confidence) pairs + if list_or_tuple(row) and len(row) > 0 and isinstance(row[0], float): + if not labels: + raise ValueError(f"Needed labels to assign the confidences returned. Confidences: {confidences}") + label_conf_pairs = list(zip(labels, row)) + # sort them by confidence + label_conf_pairs = sorted(label_conf_pairs, key=lambda pair: pair[1], reverse=True) + labels_and_confidences.append(label_conf_pairs) + else: + # if this row is (label, confidence) pairs already, sort them + if list_or_tuple(row) and len(row) > 0 and list_or_tuple(row[0]) and len(row[0]) == 2 and isinstance(row[0][0], str) and isinstance(row[0][1], float): + label_conf_pairs = sorted(row, key=lambda pair: pair[1], reverse=True) + labels_and_confidences.append(label_conf_pairs) + else: + raise ValueError(f"Found unexpected confidence return: {confidences}") + + # grab the predicted class if it exists + prediction, _ = dict_get_compat(in_dict=results, current_key=PREDICTED_LABEL, + compat_keys=PREDICTED_LABEL_COMPAT) + # if there was no prediction, grab the label with the highest confidence from labels_and_confidences + if prediction is None: + new_prediction = [] + for row in labels_and_confidences: + new_prediction.append(row[0][0]) + prediction = new_prediction + + # un-batch if this is a batch size of 1, so that the return is just the value for the single image + labels_and_confidences = _un_batch(labels_and_confidences) + + # un-batch if this is a batch size of 1, so that the return is just the value for the single image + prediction = _un_batch(prediction) + + self.labels = labels_and_confidences + self.prediction = prediction + + def as_dict(self): + return { + "Labels": self.labels, + "Prediction": self.prediction, + } + + def __str__(self) -> str: + return json.dumps(self.as_dict()) + + +def _un_batch(item): + """ + Given an arbitrary input, if it is a list with exactly one item then return that first item + """ + if isinstance(item, list) and len(item) == 1: + item = item[0] + return item + + +def _is_label_conf_pair(row): + return ( + list_or_tuple(row) and len(row) > 0 and + list_or_tuple(row[0]) and len(row[0]) == 2 and + isinstance(row[0][0], str) and isinstance(row[0][1], float) + ) diff --git a/src/lobe/signature.py b/src/lobe/signature.py new file mode 100644 index 0000000..4ccf980 --- /dev/null +++ b/src/lobe/signature.py @@ -0,0 +1,74 @@ +from __future__ import annotations +import os +import json +import pathlib +from typing import List, Dict +from .signature_constants import ( + ID, NAME, VERSION, FORMAT, FILENAME, TAGS, CLASSES_KEY, LABELS_LIST, INPUTS, OUTPUTS, IMAGE_INPUT, TENSOR_SHAPE +) + + +def get_signature_path(model_or_sig_path: str): + model_or_sig_path = os.path.realpath(os.path.expanduser(model_or_sig_path)) + + # This could be a full_path to the signature file + if os.path.isfile(model_or_sig_path): + filename, extension = os.path.splitext(model_or_sig_path) + if (extension.lower() != ".json"): # Signature file must end in "json" + raise ValueError(f"Model file provided is not valid: {model_or_sig_path}") + signature_path = model_or_sig_path # We have the signature file, so load the model + elif os.path.isdir(model_or_sig_path): + # This is a directory with a single Signature File to load + signature_path = os.path.join(model_or_sig_path, "signature.json") + else: + raise ValueError(f"Invalid Signature file or Model directory: {model_or_sig_path}") + + if not os.path.isfile(signature_path): + raise ValueError(f"signature.json file not found at path: {model_or_sig_path}") + + return pathlib.Path(signature_path) + + +class Signature(object): + def __init__(self, model_or_sig_path: str): + """ + Loads the Signature for the given path to a Lobe signature.json file, or the exported model directory. + + - Use model path when: Using Lobe-Python in its default config, the Signature and TensorFlow (and TFLite) models are expected to be + in one folder by themselves. Additional models may exist, but they too are expected to be in their own folders. + - Use signature filepath when: Using Lobe-Python with multiple TensorFlow (and TFLite) models in the same folder, with + the Signature and Model files named uniquely. This allows you to store multiple TensorFlow/TFLite models and signatures in the same folder. + """ + # get the signature.json path from the input model or signature path + signature_path = get_signature_path(model_or_sig_path) + self.model_path = str(signature_path.parent) + + with open(signature_path, "r") as f: + self._signature = json.load(f) + + self.id: str = self._signature.get(ID) + self.name: str = self._signature.get(NAME) + self.version: str = self._signature.get(VERSION) + self.format: str = self._signature.get(FORMAT) + self.filename: str = self._signature.get(FILENAME) + self.tags: List[str] = self._signature.get(TAGS) + + self.inputs: Dict[any, any] = self._signature.get(INPUTS) + self.outputs: Dict[any, any] = self._signature.get(OUTPUTS) + + def as_dict(self): + return self._signature + + def __str__(self): + return json.dumps(self.as_dict()) + + +class ImageClassificationSignature(Signature): + def __init__(self, model_or_sig_path: str): + super(ImageClassificationSignature, self).__init__(model_or_sig_path) + + input_tensor_shape: List[int] = self.inputs[IMAGE_INPUT][TENSOR_SHAPE] + assert len(input_tensor_shape) == 4 + self.input_image_size = (input_tensor_shape[1], input_tensor_shape[2]) + + self.classes: List[str] = self._signature.get(CLASSES_KEY, {}).get(LABELS_LIST) diff --git a/src/lobe/signature_constants.py b/src/lobe/signature_constants.py new file mode 100644 index 0000000..dec4ae5 --- /dev/null +++ b/src/lobe/signature_constants.py @@ -0,0 +1,35 @@ +""" +The constants used for Lobe model exports +""" + +# predicted label from the 'outputs' top-level +PREDICTED_LABEL = 'Prediction' +PREDICTED_LABEL_COMPAT = ['Value'] # list of other previous keys to be backwards-compatible +# predicted confidence values from the 'outputs' top-level +LABEL_CONFIDENCES = 'Confidences' +LABEL_CONFIDENCES_COMPAT = ['Labels'] # list of other previous keys to be backwards-compatible + +# 'classes' top-level key +CLASSES_KEY = 'classes' + +# labels from signature.json 'classes' top-level key +LABELS_LIST = 'Label' + +# inputs top-level key +INPUTS = 'inputs' + +# outputs top-level return key +OUTPUTS = 'outputs' + +ID = 'doc_id' +NAME = 'doc_name' +VERSION = 'doc_version' +FORMAT = 'format' +FILENAME = 'filename' +TAGS = 'tags' + +# Input or output properties +TENSOR_NAME = 'name' +TENSOR_SHAPE = 'shape' + +IMAGE_INPUT = 'Image' diff --git a/src/lobe/utils.py b/src/lobe/utils.py new file mode 100644 index 0000000..1465880 --- /dev/null +++ b/src/lobe/utils.py @@ -0,0 +1,54 @@ +""" +Some utility functions +""" +from typing import Dict, List, Tuple, Optional, Union + + +def dict_get_compat(in_dict: Dict[str, any], current_key: str, compat_keys: List[str], default: any = None) -> Tuple[any, Optional[str]]: + """ + Given a dictionary, an expected key, and a list of older compatibility keys, return the first found value + from the dict and which key it was. + """ + test_keys = [current_key] + compat_keys + for key in test_keys: + if key in in_dict: + return in_dict[key], key + return default, None + + +def decode_dict_bytes_as_str(in_dict: Dict[any, any], encoding="utf-8"): + """ + Recursively decode any byte values in the dict as strings with .decode() + """ + # modifies the dict in place + for key, val in in_dict.items(): + if isinstance(val, bytes): + in_dict[key] = val.decode(encoding) + elif isinstance(val, list) or isinstance(val, tuple): + in_dict[key] = decode_list_bytes_as_str(val, encoding=encoding) + elif isinstance(val, dict): + decode_dict_bytes_as_str(val) + + +def decode_list_bytes_as_str(in_list: Union[List[any], Tuple[any]], encoding="utf-8"): + """ + Recursively decodes any byte values in the list as strings with .decode() + """ + # make the return the same type as this one + decoded_list = [] + for item in in_list: + if isinstance(item, bytes): + decoded_list.append(item.decode(encoding)) + elif isinstance(item, list) or isinstance(item, tuple): + decoded_list.append(decode_list_bytes_as_str(item)) + else: + decoded_list.append(item) + + return type(in_list)(decoded_list) + + +def list_or_tuple(item): + """ + returns true if the item is a list or tuple + """ + return isinstance(item, list) or isinstance(item, tuple)