chore: checking with plt obj

app.py

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

utils.py

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