Update app.py

app.py

@@ -1,6 +1,14 @@
 import streamlit as st
 import matplotlib.pyplot as plt
 import numpy as np
+import io
+
+# Convert the plot to bytes for download
+def fig_to_bytes(fig):
+    img_bytes = io.BytesIO()
+    fig.savefig(img_bytes, format="png")
+    img_bytes.seek(0)
+    return img_bytes.read()
 
 # Function to draw the plot
 def draw_plot(length, width, unit):
@@ -50,11 +58,3 @@ if length and width:
         file_name="plot_image.png",
         mime="image/png"
     )
-
-# Convert the plot to bytes for download
-def fig_to_bytes(fig):
-    import io
-    img_bytes = io.BytesIO()
-    fig.savefig(img_bytes, format="png")
-    img_bytes.seek(0)
-    return img_bytes.read()