updated app.py

app.py

@@ -28,7 +28,7 @@ model_path = os.path.join(OUTPUT_DIR, "model_final.pth")
 GOOGLE_DRIVE_FILE_ID = "1yr64AOgaYZPTcQzG6cxG6lWBENHR9qjW"
 GDRIVE_URL = f"https://drive.google.com/uc?id={GOOGLE_DRIVE_FILE_ID}"
 
-# Make folders if they don't exist
+# Create necessary folders
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(JSON_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
@@ -53,7 +53,7 @@ sys.path.append(MODEL_DIR)
 from rcnn_run import main, write_config
 
 # ==================================
-# CSS Styling
+# PAGE CONFIG
 # ==================================
 
 st.set_page_config(
@@ -62,24 +62,6 @@ st.set_page_config(
     initial_sidebar_state="collapsed"
 )
 
-st.markdown(
-    """
-    <style>
-    .stApp { background-color: #FAFAFA; }
-    .header-title { font-size: 2.5rem; font-weight: bold; text-align: center;
-                    background: linear-gradient(to right, #D4ECDD, #EAF4F4);
-                    color: #2C3E50; padding: 20px; border-radius: 12px; }
-    .upload-container { display: flex; flex-direction: column; align-items: center;
-                        justify-content: center; background: white; padding: 20px;
-                        border-radius: 10px; box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1); }
-    .json-container { background: #F5F5F5; padding: 15px; border-radius: 10px;
-                      font-family: monospace; overflow-y: auto; max-height: 400px;
-                      white-space: pre-wrap; }
-    </style>
-    """,
-    unsafe_allow_html=True
-)
-
 # ==================================
 # HEADER
 # ==================================
@@ -130,11 +112,14 @@ if uploaded_file is not None:
             progress_bar = st.progress(0)
             status_text = st.empty()
 
-            # === 🔥 Run Model Here ===
+            # === 🔥 Model Run Here ===
             input_image = uploaded_path
             output_json_name = uploaded_file.name.replace(".png", "_result.json").replace(".jpg", "_result.json").replace(".jpeg", "_result.json")
             output_image_name = uploaded_file.name.replace(".png", "_result.png").replace(".jpg", "_result.png").replace(".jpeg", "_result.png")
 
+            output_json_path = os.path.join(JSON_DIR, output_json_name)
+            output_image_path = os.path.join(JSON_DIR, output_image_name)
+
             cfg = write_config()
             print("⚙️ Model config created. Running model...")
 
@@ -144,15 +129,9 @@ if uploaded_file is not None:
                 progress_bar.progress(i)
                 status_text.text(f"Preprocessing: {i}%")
 
-
-            output_json_path = os.path.join(JSON_DIR, output_json_name)
-            output_image_path = os.path.join(JSON_DIR, output_image_name)
+            # Run model
             main(cfg, input_image, output_json_path, output_image_path)
             print("✅ Model run complete.")
-            # main(cfg, input_image, output_json_name, output_image_name)
-            # # Prepare output paths
-            # output_json_path = os.path.join(JSON_DIR, output_json_name)
-            # output_image_path = os.path.join(JSON_DIR, output_image_name)
 
             while not os.path.exists(output_json_path):
                 print("Waiting for JSON output...")
@@ -179,26 +158,26 @@ if uploaded_file is not None:
             st.session_state.processing_complete = True
 
         # ==================================
-        # DISPLAY Output Image + Download Buttons + JSON side-by-side
+        # DISPLAY OUTPUTS
         # ==================================
 
         out_col1, out_col2 = st.columns(2)
 
         with out_col1:
             if os.path.exists(output_image_path):
-                # st.image(output_image_path, caption="🖼 Output Vectorized Image", use_container_width=True)
                 with open(output_image_path, "rb") as img_file:
                     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"
-                )
+                    img_file.seek(0)
+                    st.download_button(
+                        label="Download Output Image",
+                        data=img_file,
+                        file_name="floorplan_output.png",
+                        mime="image/png"
+                    )
 
+            if os.path.exists(output_json_path):
                 json_str = json.dumps(st.session_state.json_output, indent=4)
                 st.download_button(
                     label="Download JSON",
@@ -206,8 +185,6 @@ if uploaded_file is not None:
                     file_name="floorplan_output.json",
                     mime="application/json"
                 )
-            else:
-                st.warning("⚠️ Output image not found.")
 
         with out_col2:
             st.markdown("<div class='json-container'>", unsafe_allow_html=True)