chore: attention only plot

app.py

@@ -43,21 +43,18 @@ def plot_attention(image):
 	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")
-
+	img = plt.imshow(selected_image)
+	plt.imshow(
+		selected_weight,
+		cmap="inferno",
+		alpha=0.6,
+		extent=img.get_extent()
+	)
 	plt.axis("off")
 	return plt
 
 iface = gr.Interface(
 	fn=plot_attention,
 	inputs=[gr.inputs.Image(label="Input Image")],
-	outputs="image").launch()
+	outputs="Attention Map").launch()