chore: final

app.py

@@ -43,8 +43,9 @@ with st.spinner("Generating the heat maps... HOLD ON!"):
 		image=preprocessed_img_orig
 	)
 
-	plt = utils.plot(attentions=attentions, image=preprocessed_img_orig)
+	utils.plot(attentions=attentions, image=preprocessed_img_orig)
 
 # Show the attention maps
 st.title("Attention 🔥 Maps")
-st.image(plt, caption="Attention Heat Maps")
+image = Image.open("heat_map.png")
+st.image(image, caption="Attention Heat Maps")

utils.py

@@ -97,4 +97,4 @@ def plot(attentions, image):
                 img_count += 1
     
     plt.tight_layout()
-    return plt
+    plt.savefig("heat_map.png")