Add Grad-CAM++ visualization for exported tf models

.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/backend_onnx.py → src/lobe/backends/onnx/backend.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/onnx/image_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)

src/lobe/backends/backend_tf.py → src/lobe/backends/tf/backend.py

@@ -1,23 +1,31 @@
 from threading import Lock
 
-from ..signature import Signature
-from ..utils import decode_dict_bytes_as_str
+from ..backend import Backend
+from ...signature import Signature
+from ...utils import decode_dict_bytes_as_str
+
+TF_IMPORT_ERROR = """
+ERROR: This is a TensorFlow model and requires tensorflow to be installed on this device. 
+Please install lobe-python with lobe[tf] or lobe[all] options. 
+If that doesn't work, please go to https://www.tensorflow.org/install for instructions.
+"""
 
 try:
     import tensorflow as tf
+    from tensorflow.python.training.tracking.tracking import AutoTrackable
 except ImportError:
-    raise ImportError("ERROR: This is a TensorFlow model and requires tensorflow to be installed on this device. Please install lobe-python with lobe[tf] or lobe[all] options. If that doesn't work, please go to https://www.tensorflow.org/install for instructions.")
+    raise ImportError(TF_IMPORT_ERROR)
 
 
-class TFModel(object):
+class TFModel(Backend):
     """
-    Generic wrapper for running a tensorflow model
+    Generic wrapper for running a TensorFlow model from Lobe.
     """
     def __init__(self, signature: Signature):
+        super(TFModel, self).__init__(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.model: AutoTrackable = 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):

src/lobe/backends/tf/image_backend.py

@@ -0,0 +1,162 @@
+from .backend import TFModel, TF_IMPORT_ERROR
+from ..backend import ImageBackend
+from ...signature import ImageClassificationSignature
+
+import numpy as np
+
+from ...signature_constants import SUPPORTED_EXPORT_VERSIONS, LABEL_CONFIDENCES, LABEL_CONFIDENCES_COMPAT, IMAGE_INPUT, TENSOR_NAME
+from ...utils import dict_get_compat
+
+try:
+    import tensorflow as tf
+except ImportError:
+    raise ImportError(TF_IMPORT_ERROR)
+
+
+class TFImageModel(TFModel, ImageBackend):
+    signature: ImageClassificationSignature
+
+    def __init__(self, signature: ImageClassificationSignature):
+        super(TFImageModel, self).__init__(signature=signature)
+
+    def gradcam_plusplus(self, image: np.ndarray, label=None) -> np.ndarray:
+        """
+        Implementation of Grad-CAM++,
+        adapted for TF 2.x from the original source: https://github.com/adityac94/Grad_CAM_plus_plus
+
+         @article{chattopadhyay2017grad,
+          title={Grad-CAM++: Generalized Gradient-based Visual Explanations for Deep Convolutional Networks},
+          author={Chattopadhyay, Aditya and Sarkar, Anirban and Howlader, Prantik and Balasubramanian, Vineeth N},
+          journal={arXiv preprint arXiv:1710.11063},
+          year={2017}
+        }
+        """
+        labels = self.signature.classes
+        # Get the output index of the desired label for visualizing
+        # If no desired label is given, find the predicted label by running the model on the image
+        if label is None:
+            label_idx = self._get_predicted_label_argmax(image=image)
+        else:
+            # if we are batched, get the indices by looping, otherwise just get the index
+            if isinstance(label, list):
+                label_idx = [labels.index(_label) for _label in label]
+            else:
+                label_idx = [labels.index(label)]
+        # create a one-hot vector of our label indices to use as a mask for the output cost
+        label_idx = tf.one_hot(label_idx, depth=len(labels))
+        if len(label_idx) != len(image):
+            raise ValueError(
+                f"Supplied label (or list of labels) does not match the the number of input images. Images : {len(image)}, labels: {len(label_idx)}"
+            )
+
+        with self.lock:
+            # now we want to get the derivatives of the output with respect to the last conv layer
+            # get the layer name of the confidences logits output and the last convolutional layer
+            last_fc_tensor, last_conv_tensor = self._get_last_fc_and_conv_tensors()
+
+            # now get the function that returns the fc and conv tensors from the image
+            input_image_name = self.signature.inputs[IMAGE_INPUT][TENSOR_NAME]
+            last_conv_fn = self.model.prune(input_image_name, last_conv_tensor.name)
+            last_fc_fn = self.model.prune(last_conv_tensor.name, last_fc_tensor.name)
+
+            # get the last conv out
+            last_conv_out = last_conv_fn(tf.constant(image))
+
+            # take the 3 derivatives of the cost wrt the conv layer
+            with tf.GradientTape() as t3:
+                t3.watch(last_conv_out)
+                with tf.GradientTape() as t2:
+                    t2.watch(last_conv_out)
+                    with tf.GradientTape() as t1:
+                        t1.watch(last_conv_out)
+                        last_fc_out = last_fc_fn(last_conv_out)
+                        # get the output neuron corresponding to the class of interest
+                        cost = last_fc_out * label_idx
+                        # first derivative
+                        conv_first_grad = t1.gradient(cost, last_conv_out)
+                    # second derivative
+                    conv_second_grad = t2.gradient(conv_first_grad, last_conv_out)
+                # triple derivative
+                conv_third_grad = t3.gradient(conv_second_grad, last_conv_out)
+
+            batch, _, _, filters = last_conv_out.shape
+
+            global_sum = tf.math.reduce_sum(last_conv_out, axis=[1, 2])
+
+            alpha_num = conv_second_grad
+            broadcasted_global_sum = tf.reshape(global_sum, (batch, 1, 1, filters))
+            alpha_denom = conv_second_grad * 2.0 + conv_third_grad * broadcasted_global_sum
+            alpha_denom = tf.where(alpha_denom != 0.0, alpha_denom, tf.ones(alpha_denom.shape))
+            alphas = alpha_num / alpha_denom
+
+            weights = tf.maximum(conv_first_grad, 0.0)
+
+            alphas_thresholding = tf.where(weights != 0.0, alphas, 0.0)
+
+            alpha_normalization_constant = tf.math.reduce_sum(alphas_thresholding, axis=[1, 2])
+            alpha_normalization_constant_processed = tf.where(alpha_normalization_constant != 0.0,
+                                                              alpha_normalization_constant,
+                                                              tf.ones(alpha_normalization_constant.shape))
+
+            alphas /= tf.reshape(alpha_normalization_constant_processed, (batch, 1, 1, filters))
+
+            deep_linearization_weights = tf.math.reduce_sum((weights * alphas), axis=[1, 2])
+            broadcasted_deep_lin_weights = tf.reshape(deep_linearization_weights, (batch, 1, 1, filters))
+            grad_CAM_map = tf.math.reduce_sum(broadcasted_deep_lin_weights * last_conv_out, axis=3)
+
+            # Passing through ReLU
+            cam = tf.maximum(grad_CAM_map, 0)
+            cam_max = tf.math.reduce_max(cam, axis=[1, 2])
+            cam /= tf.reshape(cam_max, (batch, 1, 1))  # scale 0 to 1.0
+            return cam.numpy()
+
+    def _get_predicted_label_argmax(self, image: np.ndarray):
+        """
+        Given an image, run our model and return the array of predicted argmax indices.
+        """
+        result = self.predict(data=image)
+        if self.signature.export_version not in SUPPORTED_EXPORT_VERSIONS:
+            raise ValueError(
+                f"Lobe model export version {self.signature.export_version} not supported. Need one of: {SUPPORTED_EXPORT_VERSIONS}")
+
+        # grab the list of confidences
+        confidences, _ = dict_get_compat(
+            in_dict=result, current_key=LABEL_CONFIDENCES, compat_keys=LABEL_CONFIDENCES_COMPAT, default=[]
+        )
+        return tf.argmax(confidences, axis=1)
+
+    def _get_last_fc_and_conv_tensors(self):
+        """
+        Gets the tensor that represents the last fully-connected layer outputs (logits).
+        Since the 'Confidences' output is the softmax tensor, its input is the last FC layer.
+
+        Also find the tensor that is the last convolution layer's output before any pooling (this will be the
+        RELU output before the global max or avg pooling)
+        """
+        confidences_out, _ = dict_get_compat(
+            in_dict=self.signature.outputs, current_key=LABEL_CONFIDENCES, compat_keys=LABEL_CONFIDENCES_COMPAT
+        )
+        softmax_tensor = self.model.graph.get_tensor_by_name(confidences_out.get(TENSOR_NAME))
+        # get the op (softmax)'s inputs -- the last fc layer tensor will be the only (first) input to this op
+        last_fc_tensor = softmax_tensor.op.inputs[0]
+
+        # now from the last fc layer, bfs search for the closest max pooling op and find its input -- that is the
+        # last conv layer output (the RELU tensor)
+        last_conv_tensor = None
+        visited, queue = [], []
+        queue.append(last_fc_tensor)
+        while queue:
+            tensor = queue.pop()
+            visited.append(tensor.name)
+            op = tensor.op
+            # if this was from the max/avg pool op, get the input tensor
+            # (which is the output of the last conv layer's relu)
+            if op.type in ["Max", "Mean"]:
+                last_conv_tensor = op.inputs[0]
+                break
+            else:
+                for input_tensor in op.inputs:
+                    if input_tensor.name not in visited:
+                        queue.append(input_tensor)
+
+        return last_fc_tensor, last_conv_tensor