V0.6 Add Grad-CAM++ visualization for exported tf models, refactor backend base classes (#28)

* Add Grad-CAM++ visualization for exported tensorflow image classification models.

* Clean comments

* Batch-ify the ImageModel.visualize() function so you can pass a list of images.

* Bump Pillow version to latest

Author: mbeissinger

.gitignore

@@ -115,6 +115,7 @@ venv.bak/
 .v36
 .v37
 .v38
+.v39
 
 # Spyder project settings
 .spyderproject

README.md

@@ -1,7 +1,7 @@
 # Lobe Python API
 Code to run exported Lobe models in Python using the TensorFlow, TensorFlow Lite, or ONNX options.
 
-Works with Python 3.6, 3.7, and 3.8 untested for other versions.
+Works with Python 3.6, 3.7, 3.8, and 3.9 untested for other versions.
 
 ## Install
 ### Backend options with pip
@@ -94,6 +94,10 @@ print(result.prediction)
 for label, confidence in result.labels:
     print(f"{label}: {confidence*100}%")
 
+# Visualize the heatmap of the prediction on the image 
+# this shows where the model was looking to make its prediction.
+heatmap = model.visualize(img)
+heatmap.show()
 ```
 Note: model predict functions should be thread-safe. If you find bugs please file an issue.
 

RELEASE.md

@@ -1,3 +1,23 @@
+# Release 0.6.0
+___
+## Breaking Changes
+* Refactored the ML backends into sub-folders:
+  * `TFModel` class: `backends/backend_tf.py -> backends/tf/backend.py`
+  * `TFLiteModel` class: `backends/backend_tflite.py -> backends/tflite/backend.py`
+  * `ONNXModel` class: `backends/backend_onnx.py -> backends/onnx/backend.py`
+
+## Bug Fixes and Other Improvements
+* Added `Backend` and `ImageBackend` abstract base classes in `backends/backend.py`
+* Added ImageBackend classes for each ML backend:
+  * `TFImageModel` class: `backends/tf/image_backend.py`
+  * `TFLiteImageModel` class: `backends/tflite/image_backend.py`
+  * `ONNXImageModel` class: `backends/onnx/image_backend.py`
+* Added Grad-CAM++ implementation (`ImageBackend.gradcam_plusplus(image, label) -> np.ndarray`) for visualizing 
+convolutional neural network heatmaps for explaining why the model predicted a certain label. 
+_Note:_ Grad-CAM++ only implemented currently in `TFImageModel` for TensorFlow Lobe model exports.
+The visualization can be called from the top-level API of `ImageModel` -> `ImageModel.visualize(image)`
+
+
 # Release 0.5.0
 ___
 ## Breaking Changes

setup.py

@@ -20,8 +20,9 @@
     mac_version = None
 
 requirements = [
-    "pillow~=8.3.1",
-    "requests"
+    "pillow~=8.4.0",
+    "requests",
+    "matplotlib~=3.4.3",
 ]
 tf_req = "tensorflow~=2.5.0;platform_machine!='armv7l'"
 onnx_req = "onnxruntime~=1.8.1;platform_machine!='armv7l'"
@@ -68,7 +69,7 @@
 
 setup(
     name="lobe",
-    version="0.5.0",
+    version="0.6.0",
     description="Lobe Python SDK",
     long_description=readme,
     long_description_content_type="text/markdown",
@@ -81,6 +82,5 @@
         'all': [tf_req, onnx_req],
         'tf': [tf_req],
         'onnx': [onnx_req],
-        #'tflite': [tflite_req],
     }
 )

src/lobe/__init__.py

@@ -1,2 +1,2 @@
 from .signature import Signature
-from .model.image_model import ImageModel
+from .model.image_model import ImageModel, VizEnum

src/lobe/backends/backend.py

@@ -0,0 +1,37 @@
+"""
+Abstract for our backend implementations.
+"""
+from abc import ABC, abstractmethod
+
+from ..signature import Signature
+from ..results import BackendResult
+
+
+class Backend(ABC):
+	def __init__(self, signature: Signature):
+		self.signature = signature
+
+	@abstractmethod
+	def predict(self, data: any) -> BackendResult:
+		"""
+		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, ...}
+		"""
+		pass
+
+
+class ImageBackend(Backend):
+	def gradcam_plusplus(self, image, label: str = None):
+		"""
+		Return the heatmap from Grad-CAM++
+		https://arxiv.org/abs/1710.11063
+		Grad-CAM++: Improved Visual Explanations for Deep Convolutional Networks
+		Aditya Chattopadhyay, Anirban Sarkar, Prantik Howlader, Vineeth N Balasubramanian
+		"""
+		raise NotImplementedError(
+			f"Image backend {self.__class__.__name__} doesn't have a Grad-CAM++ implementation yet."
+		)

src/lobe/backends/onnx/__init__.py

@@ -1,16 +1,20 @@
 from threading import Lock
-from ..signature import Signature
-from ..signature_constants import TENSOR_NAME
-from ..utils import decode_dict_bytes_as_str
+from ...signature import Signature
+from ...signature_constants import TENSOR_NAME
+from ...utils import decode_dict_bytes_as_str
+
+ONNX_IMPORT_ERROR = """
+ERROR: This is an ONNX model and requires onnx runtime to be installed on this device. 
+Please install lobe-python with lobe[onnx] or lobe[all] options. 
+If that doesn't work, please go to https://www.onnxruntime.ai/ for install instructions.
+"""
 
 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 install lobe-python with lobe[onnx] or lobe[all] options. If that doesn't work, please go to https://www.onnxruntime.ai/ for install instructions."
-    )
+    raise ImportError(ONNX_IMPORT_ERROR)
 
 
 class ONNXModel(object):

src/lobe/backends/backend_onnx.py src/lobe/backends/onnx/backend.pysrc/lobe/backends/backend_onnx.py renamed to src/lobe/backends/onnx/backend.py

@@ -0,0 +1,11 @@
+from .backend import ONNXModel
+from ..backend import ImageBackend
+from ...signature import ImageClassificationSignature
+
+
+class ONNXImageModel(ONNXModel, ImageBackend):
+    def __init__(self, signature: ImageClassificationSignature):
+        super(ONNXModel, self).__init__(signature=signature)
+
+    def gradcam_plusplus(self, image, label=None):
+        super(ONNXImageModel, self).gradcam_plusplus(image=image, label=label)