chore: house cleaning

app.py

@@ -1,8 +1,11 @@
 # import the necessary packages
 from utilities import config
 from utilities import model
-from utilities import visualization
 from tensorflow import keras
+from tensorflow.keras import layers
+import tensorflow as tf
+import matplotlib.pyplot as plt
+import math
 import gradio as gr
 
 # load the models from disk
@@ -19,15 +22,42 @@ conv_attn = keras.models.load_model(
 	compile=False
 )
 
-# create the patch conv net
-patch_conv_net = model.PatchConvNet(
-	stem=conv_stem,
-	trunk=conv_trunk,
-	attention_pooling=conv_attn,
-)
+def plot_attention(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 = conv_stem(image)
+	# pass through the trunk
+	test_x = conv_trunk(test_x)
+	# pass through the attention pooling block
+	_, test_viz_weights = conv_attn(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
 
-# 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")],

utilities/model.py

@@ -1,30 +0,0 @@
-# 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

@@ -1,45 +0,0 @@
-# import the necessary packages
-from tensorflow.keras import layers
-import tensorflow as tf
-import matplotlib.pyplot as plt
-import math
-
-class PlotAttention(object):
-	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