Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import matplotlib.pyplot as plt
+import io
 
 def convert_to_meters(value, unit):
     conversion_factors = {
@@ -10,16 +11,16 @@ def convert_to_meters(value, unit):
     }
     return value * conversion_factors[unit]
 
-def draw_floor_plan(length, width):
+def draw_floor_plan(length, width, unit):
     fig, ax = plt.subplots(figsize=(8, 8))
     ax.plot([0, length, length, 0, 0], [0, 0, width, width, 0], marker='o', color='blue')
     ax.set_xlim(-1, length + 1)
     ax.set_ylim(-1, width + 1)
     ax.set_aspect('equal', adjustable='box')
-    ax.set_title("Floor Plan")
-    ax.set_xlabel("Length (m)")
-    ax.set_ylabel("Width (m)")
-    st.pyplot(fig)
+    ax.set_title(f"Floor Plan ({unit})")
+    ax.set_xlabel(f"Length ({unit})")
+    ax.set_ylabel(f"Width ({unit})")
+    return fig
 
 st.title("Floor Plan Generator")
 
@@ -28,13 +29,27 @@ st.write("Enter the dimensions of the floor plan and select the unit of measurem
 # Input fields
 length = st.number_input("Enter the length:", min_value=0.0, format="%.2f")
 width = st.number_input("Enter the width:", min_value=0.0, format="%.2f")
-unit = st.selectbox("Select the unit:", ["m", "ft", "cm", "in"], index=0)
+unit = st.selectbox("Select the unit:", ["m", "ft", "cm", "in"], index=1)
 
 if st.button("Generate Floor Plan"):
     if length > 0 and width > 0:
-        length_m = convert_to_meters(length, unit)
-        width_m = convert_to_meters(width, unit)
-        st.success(f"Dimensions in meters: Length = {length_m:.2f} m, Width = {width_m:.2f} m")
-        draw_floor_plan(length_m, width_m)
+        length_converted = convert_to_meters(length, unit) if unit != "m" else length
+        width_converted = convert_to_meters(width, unit) if unit != "m" else width
+        st.success(f"Dimensions: Length = {length:.2f} {unit}, Width = {width:.2f} {unit}")
+        fig = draw_floor_plan(length, width, unit)
+
+        # Show the plot
+        st.pyplot(fig)
+
+        # Add a download button
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png")
+        buf.seek(0)
+        st.download_button(
+            label="Download Floor Plan",
+            data=buf,
+            file_name="floor_plan.png",
+            mime="image/png"
+        )
     else:
         st.error("Please enter positive values for length and width.")