Update app.py

app.py

@@ -162,13 +162,14 @@ class Visualization:
         
         plt.tight_layout()
         
-        # Convert matplotlib figure to PIL Image
-        fig.canvas.draw()
-        img_array = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
-        img_array = img_array.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+        # Save to BytesIO and convert to PIL Image
+        buf = BytesIO()
+        fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
+        buf.seek(0)
+        img = Image.open(buf)
         plt.close(fig)
         
-        return Image.fromarray(img_array)
+        return img
 
 def download_dataset():
     """Download the dataset using kagglehub"""