Possible to do heatmaps using tflite models?

[Medicmind]

Is it possible to do heatmaps on keras_cv_attention_models once they are converted to tflite format? To save server space I convert trained models to tflite format but I still also want to do heatmaps with new images. Is this possible? From what I've read you cannot convert a tflite model back to a keras model.

[leondgarse]

I think it's possible but not easy refering Applying Grad-Cam to a model that was converted to Tensorflow Lite. Just we need the original keras model anyway, for exporting a TFLite model with train signature. Not sure if anyone has done that...

[leondgarse]

OK, I think it's workng now. Two functions grads_to_heatmap and apply_heatmap are separated out from previous make_and_apply_gradcam_heatmap, and added a converting class ModelWithGradForTFLite with gradcam signature.

from keras_cv_attention_models import efficientnet, test_images
mm = efficientnet.EfficientNetV2B0()

from keras_cv_attention_models.visualizing.gradcam_heatmap import ModelWithGradForTFLite, grads_to_heatmap, apply_heatmap
saved_model = mm.name
bb = ModelWithGradForTFLite(mm)
signatures = {'serving_default': bb.serving_default.get_concrete_function(), 'gradcam': bb.gradcam.get_concrete_function()}
tf.saved_model.save(bb, saved_model, signatures=signatures)

# Convert the model
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model)
# enable TensorFlow and TensorFlow ops.
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
tflite_model = converter.convert()
open(mm.name + ".tflite", "wb").write(tflite_model)  # Will save a much larger model than without these signatures

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
gradcam = interpreter.get_signature_runner("gradcam")

image = test_images.dog_cat()
gg = gradcam(inputs=tf.expand_dims(tf.image.resize(image , [224, 224]) / 255, 0))
grads, last_conv_layer_output, class_channel = gg['grads'], gg['last_conv_layer_output'], gg['class_channel']
# Show
heatmap = grads_to_heatmap(grads, last_conv_layer_output, class_channel=class_channel, use_v2=True)
superimposed_img = apply_heatmap(image, heatmap, alpha=0.8, plot=True)

You may find the raw code test in kecam_test.ipynb TFLite Heatmap. But still it's not recommended, as the converted TFLite model will be much larger, almost ~3 times. Just a good practice with TFLite signature.
But from an already converted TFLite model, don't know how...

[Medicmind]

Awesome thanks. And it gives the same heatmap each time. I made few small changes that seem to make it work for my images.

from keras_cv_attention_models import efficientnet, test_images
import tensorflow as tf
from PIL import Image
import numpy as np
from keras_cv_attention_models.visualizing.gradcam_heatmap import ModelWithGradForTFLite, grads_to_heatmap, apply_heatmap
mm = efficientnet.EfficientNetV2S(num_classes=2,pretrained=None)
mm.load_weights('EfficientNetV2S_latest.h5', by_name=True, skip_mismatch=True) #EfficientNetV2S_latest.h5')

saved_model = mm.name
bb = ModelWithGradForTFLite(mm)
signatures = {'serving_default': bb.serving_default.get_concrete_function(), 'gradcam': bb.gradcam.get_concrete_function()}
tf.saved_model.save(bb, saved_model, signatures=signatures)

# Convert the model
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model)
# enable TensorFlow and TensorFlow ops.
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
tflite_model = converter.convert()

open(mm.name + ".tflite", "wb").write(tflite_model)  # Will save a much larger model than without these signatures

interpreter = tf.lite.Interpreter(model_content=tflite_model)

#interpreter = tf.lite.Interpreter(model_path='efficientnet_v2-s.tflite')

interpreter.allocate_tensors()
gradcam = interpreter.get_signature_runner("gradcam")

#image = test_images.dog_cat()
fl="20051020_43882_0100_PP.jpg" 
image=Image.open(fl) 
imageorig=np.array(image)
image_size=384 
input_data_type=np.float32
image = np.array(image.resize((image_size,image_size)), dtype=input_data_type)
image= np.expand_dims(image, 0) #axis=0)  
#gg = gradcam(tf.expand_dims(image.resize((image_size, image_size)) / 255, 0))
gg = gradcam(inputs=image)
grads, last_conv_layer_output, class_channel = gg['grads'], gg['last_conv_layer_output'], gg['class_channel']
# Show
heatmap = grads_to_heatmap(grads, last_conv_layer_output, class_channel=class_channel, use_v2=True)
superimposed_img = apply_heatmap(imageorig, heatmap, alpha=0.8, plot=True)
im = Image.fromarray((255*superimposed_img).astype(np.uint8))
im.save("heattflite"+fl+".png")