updated app.py

app.py

@@ -183,15 +183,18 @@ if uploaded_file is not None:
 
         with out_col1:
             if os.path.exists(output_image_path):
-                st.image(output_image_path, caption="🖼 Output Vectorized Image", use_container_width=True)
-
+                # st.image(output_image_path, caption="🖼 Output Vectorized Image", use_container_width=True)
                 with open(output_image_path, "rb") as img_file:
-                    st.download_button(
-                        label="Download Output Image",
-                        data=img_file,
-                        file_name="floorplan_output.png",
-                        mime="image/png"
-                    )
+                    image = Image.open(img_file)
+                    st.image(image, caption="🖼 Output Vectorized Image", use_container_width=True)
+
+                img_file.seek(0)  # Reset file pointer
+                st.download_button(
+                label="Download Output Image",
+                data=img_file,
+                file_name="floorplan_output.png",
+                mime="image/png"
+                )
 
                 json_str = json.dumps(st.session_state.json_output, indent=4)
                 st.download_button(