polish UI

app.py

@@ -22,12 +22,9 @@ UPLOAD_DIR = "/tmp/uploads/"
 JSON_DIR = "/tmp/results/"
 OUTPUT_DIR = "/tmp/output/"
 MODEL_DIR = os.path.join(BASE_DIR, "rcnn_model", "scripts")
-logo_path = os.path.join(BASE_DIR, "public", "logo.png")
-# model_path = os.path.join(OUTPUT_DIR, "model_final.pth")
-#  changine the model directory to the tmp directory
+logo_path = os.path.join(BASE_DIR, "public", "logo.png")  # Your Inovonics image
 model_path = os.path.join(OUTPUT_DIR, "model_final.pth")
 
-# Google Drive model
 GOOGLE_DRIVE_FILE_ID = "1yr64AOgaYZPTcQzG6cxG6lWBENHR9qjW"
 GDRIVE_URL = f"https://drive.google.com/uc?id={GOOGLE_DRIVE_FILE_ID}"
 
@@ -40,7 +37,6 @@ os.makedirs(OUTPUT_DIR, exist_ok=True)
 if not os.path.exists(model_path):
     print("🚀 Model file not found! Downloading...")
     try:
-        # gdown.download(GDRIVE_URL, model_path, quiet=False)
         gdown.download(GDRIVE_URL, model_path, quiet=False, use_cookies=False)
         print("✅ Model downloaded successfully.")
     except Exception as e:
@@ -59,55 +55,55 @@ cfg = write_config()
 def predict(uploaded_file_path):
     print("Inside Predict:" + uploaded_file_path)
     if uploaded_file_path is None:
-        return None, None, "No file uploaded."
+        return None, None, None, "No file uploaded."
 
     uploaded_path = os.path.join(UPLOAD_DIR, "input_image.png")
-    print("Saved uploaded image to:", uploaded_path)
+    shutil.copy(uploaded_file_path, uploaded_path)
+
     input_filename = "input_image.png"
 
-    # Prepare output paths
-    output_json_name = input_filename.replace(".png", "_result.json").replace(".jpg", "_result.json").replace(".jpeg", "_result.json")
-    output_image_name = input_filename.replace(".png", "_result.png").replace(".jpg", "_result.png").replace(".jpeg", "_result.png")
+    output_json_name = input_filename.replace(".png", "_result.json")
+    output_image_name = input_filename.replace(".png", "_result.png")
 
     output_json_path = os.path.join(JSON_DIR, output_json_name)
     output_image_path = os.path.join(JSON_DIR, output_image_name)
 
-    # print(f"Before calling main in app.py: {uploaded_file.name}")
-    # Run model
     main(cfg, uploaded_file_path, output_json_name, output_image_name)
 
-    # Read outputs
     result_img = Image.open(output_image_path) if os.path.exists(output_image_path) else None
     result_json = {}
     if os.path.exists(output_json_path):
         with open(output_json_path, "r") as jf:
             result_json = json.load(jf)
 
-    return result_img, json.dumps(result_json, indent=2), None
+    return uploaded_path, result_img, json.dumps(result_json, indent=2), output_json_path
 
 # ==================================
 # GRADIO UI
 # ==================================
 
 with gr.Blocks() as demo:
-    gr.Markdown("<h1 style='text-align: center;'>🏠 Inovonics 2D Floorplan Vectorizer</h1>")
+    with gr.Row():
+        gr.Image(logo_path, height=80, width=80)
+        gr.Markdown("<h1 style='text-align: center;'> Inovonics 2D Floorplan Vectorizer</h1>")
 
     with gr.Row():
         with gr.Column():
             uploaded_file = gr.File(label="Upload your Floorplan Image", type="filepath")
-            # uploaded_file = gr.File(label="Upload your Floorplan Image", type="file")
+            uploaded_img_preview = gr.Image(label="📤 Uploaded Image", interactive=False)
             run_button = gr.Button("Run Vectorizer 🔥")
 
         with gr.Column():
             output_image = gr.Image(label="🖼 Output Vectorized Image")
             output_json = gr.JSON(label="🧾 Output JSON")
+            download_button = gr.File(label="⬇️ Download JSON")
 
     error_output = gr.Textbox(label="Error Message", visible=False)
 
     run_button.click(
         predict,
         inputs=[uploaded_file],
-        outputs=[output_image, output_json, error_output]
+        outputs=[uploaded_img_preview, output_image, output_json, download_button]
     )
 
-demo.launch(server_name="0.0.0.0", server_port=7860,share=True)
+demo.launch(server_name="0.0.0.0", server_port=7860, share=True)