add: plot and bar

app.py

@@ -1,27 +1,34 @@
 # import the necessary packages
 from utilities import config
+from utilities import load_model
 from tensorflow.keras import layers
-from tensorflow import keras
 import tensorflow as tf
 import matplotlib.pyplot as plt
 import math
 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
+(conv_stem, conv_trunk, conv_attn) = load_model.loader(
+	stem=config.IMAGENETTE_STEM_PATH,
+	trunk=config.IMAGENETTE_TRUNK_PATH,
+	attn=config.IMAGENETTE_ATTN_PATH,
 )
 
-def plot_attention(image):
+# load labels
+labels = [
+	'tench',
+	'english springer',
+	'cassette player',
+	'chain saw',
+	'church',
+	'french horn',
+	'garbage truck',
+	'gas pump',
+	'golf ball',
+	'parachute'
+]
+
+def get_results(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))
@@ -32,7 +39,7 @@ def plot_attention(image):
 	# pass through the trunk
 	test_x = conv_trunk(test_x)
 	# pass through the attention pooling block
-	_, test_viz_weights = conv_attn(test_x)
+	logits, test_viz_weights = conv_attn(test_x)
 	test_viz_weights = test_viz_weights[tf.newaxis, ...]
 	
 	# reshape the vizualization weights
@@ -52,9 +59,16 @@ def plot_attention(image):
 		extent=img.get_extent()
 	)
 	plt.axis("off")
-	return plt
+
+	prediction = tf.nn.softmax(logits, axis=-1)
+
+	return plt, {labels[i]: float(prediction[i]) for i in range(10)}
 
 iface = gr.Interface(
-	fn=plot_attention,
-	inputs=[gr.inputs.Image(label="Input Image")],
-	outputs="image").launch()
+	fn=get_results,
+	inputs=gr.inputs.Image(label="Input Image"),
+	outputs=[
+		gr.outputs.Image(label="Attention Map"),
+		gr.outputs.Label(num_top_classes=10)
+	]
+).launch()

utilities/load_model.py

@@ -0,0 +1,19 @@
+# import the necessary packages
+from tensorflow import keras
+
+def loader(stem, trunk, attn):
+	# load the models from disk
+	conv_stem = keras.models.load_model(
+		stem,
+		compile=False
+	)
+	conv_trunk = keras.models.load_model(
+		trunk,
+		compile=False
+	)
+	conv_attn = keras.models.load_model(
+		attn,
+		compile=False
+	)
+
+	return (conv_stem, conv_trunk, conv_attn)