fix: reformat code for imagenette

.gitattributes

@@ -25,3 +25,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+
+# for macOS
+.DS_Store

app.py

@@ -0,0 +1,34 @@
+# import the necessary packages
+from utilities import config
+from utilities import model
+from utilities import visualization
+from tensorflow import keras
+import gradio as gr
+
+# load the models from disk
+conv_stem = keras.models.load_model(
+	config.IMAGENETTE_STEM_PATH,
+	compile=False
+)
+conv_trunk = keras.models.load_model(
+	config.IMAGENETTE_TRUNK_PATH,
+	compile=False
+)
+conv_attn = keras.models.load_model(
+	config.IMAGENETTE_ATTN_PATH,
+	compile=False
+)
+
+# create the patch conv net
+patch_conv_net = model.PatchConvNet(
+	stem=conv_stem,
+	trunk=conv_trunk,
+	attention_pooling=conv_attn,
+)
+
+# get the plot attention function
+plot_attention = visualization.PlotAttention(model=patch_conv_net)
+iface = gr.Interface(
+	fn=plot_attention,
+	inputs=[gr.inputs.Image(label="Input Image")],
+	outputs="image").launch()

model/att_pool/keras_metadata.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fefb606f0aeb214dcfa9cf9786955f6b7ecb7bdd116e007c44264af75936bfca
+size 15848

model/att_pool/saved_model.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0057c7816b0d297c8cdb02958c1a0044220eea870f923e38ba2c1a7fb01c9448
+size 324550

model/stem/keras_metadata.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e41ea234e1a8f95c77c1308242971d818b50ec3212d3b1f1f34d9042d77f1270
+size 11998

model/stem/saved_model.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee90a781360388d5b7e2035be35b349f4209a6e58881bbc02e2c01cea8130877
+size 96815

model/trunk/keras_metadata.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c090252e37268009bf4e4fc8866b2984bfdfbe88341ddb54dea464c9eb365bc4
+size 23883

model/trunk/saved_model.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:255bf760cd4df017b1d4ff670f889ac9c5985d9b41d1d2d1d401c3030797164b
+size 359160

utilities/config.py

@@ -0,0 +1,8 @@
+# import the necessary packages
+import os
+
+# define the path to the model
+MODEL_PATH = "model"
+IMAGENETTE_ATTN_PATH = os.path.join(MODEL_PATH, "att_pool")
+IMAGENETTE_STEM_PATH = os.path.join(MODEL_PATH, "stem")
+IMAGENETTE_TRUNK_PATH = os.path.join(MODEL_PATH, "trunk")

utilities/model.py

@@ -0,0 +1,30 @@
+# import the necessary packages
+from tensorflow import keras
+import tensorflow as tf
+
+# Patch conv
+class PatchConvNet(keras.Model):
+	def __init__(
+		self,
+		stem,
+		trunk,
+		attention_pooling,
+		**kwargs,
+	):
+		super().__init__(**kwargs)
+		self.stem = stem
+		self.trunk = trunk
+		self.attention_pooling = attention_pooling
+
+	@tf.function(
+	input_signature=[
+		tf.TensorSpec(shape=[None, None, None, 3], dtype=tf.uint8)
+	])
+	def call(self, images):
+		# pass through the stem
+		x = self.stem(images)
+		# pass through the trunk
+		x = self.trunk(x)
+		# pass through the attention pooling block
+		predictions, viz_weights = self.attention_pooling(x)
+		return predictions, viz_weights

utilities/visualization.py

@@ -0,0 +1,45 @@
+# import the necessary packages
+from tensorflow.keras import layers
+import tensorflow as tf
+import matplotlib.pyplot as plt
+import math
+
+class PlotAttention:
+	def __init__(self, model):
+		self.model = model
+	
+	def  __call__(self, image):
+		# resize the image to a 224, 224 dim
+		image = tf.image.convert_image_dtype(image, tf.float32)
+		image = tf.image.resize(image, (224, 224))
+		image = image[tf.newaxis, ...]
+
+		# pass through the stem
+		test_x = self.model.stem(image)
+		# pass through the trunk
+		test_x = self.model.trunk(test_x)
+		# pass through the attention pooling block
+		_, test_viz_weights = self.model.attention_pooling(test_x)
+		test_viz_weights = test_viz_weights[tf.newaxis, ...]
+		
+		# reshape the vizualization weights
+		num_patches = tf.shape(test_viz_weights)[-1]
+		height = width = int(math.sqrt(num_patches))
+		test_viz_weights = layers.Reshape((height, width))(test_viz_weights)
+		
+		index = 0
+		selected_image = image[index]
+		selected_weight = test_viz_weights[index]
+
+		fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(10, 5))
+		ax[0].imshow(selected_image)
+		ax[0].set_title(f"Original")
+		ax[0].axis("off")
+		
+		img = ax[1].imshow(selected_image)
+		ax[1].imshow(selected_weight, cmap='inferno', alpha=0.6, extent=img.get_extent())
+		ax[1].set_title(f"Attended")
+		ax[1].axis("off")
+
+		plt.axis("off")
+		return plt