trying gradio

README.md

@@ -1,14 +1,14 @@
 ---
 title: 2D Floorplan Vectorizer
 emoji: 🏠
-colorFrom: blue
-colorTo: green
-sdk: docker
+colorFrom: green
+colorTo: blue
+sdk: gradio
+sdk_version: "4.15.0"
 app_file: app.py
 pinned: false
 ---
 
-
 # 2D Floorplan Vectorizer
 
 A Streamlit web app that allows you to upload 2D floorplan images and automatically vectorize them into COCO-style annotations using a trained Mask R-CNN model.

app.py

@@ -1,200 +1,110 @@
-import streamlit as st
-import json
-import time
-from PIL import Image
+import gradio as gr
 import os
 import sys
+import json
 import shutil
 import gdown
+import time
+from PIL import Image
 from io import BytesIO
 
 # ==================================
 # SETUP
 # ==================================
 
-print("🚀 Streamlit App Starting...")
+print("🚀 Gradio App Starting...")
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
-# Setup Paths
+# Paths
 UPLOAD_DIR = "/tmp/uploads/"
-MODEL_DIR = os.path.join(BASE_DIR, "rcnn_model", "scripts")
 JSON_DIR = "/tmp/results/"
 OUTPUT_DIR = "/tmp/output/"
-SAMPLE_DIR = os.path.join(BASE_DIR, "rcnn_model", "sample")
+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")
 
-# Google Drive file download link
+# Google Drive model
 GOOGLE_DRIVE_FILE_ID = "1yr64AOgaYZPTcQzG6cxG6lWBENHR9qjW"
 GDRIVE_URL = f"https://drive.google.com/uc?id={GOOGLE_DRIVE_FILE_ID}"
 
-# Create necessary folders
+# Create folders
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(JSON_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# ==================================
-# DOWNLOAD MODEL IF MISSING
-# ==================================
-
+# Download model if missing
 if not os.path.exists(model_path):
-    print("🚀 Model file not found! Downloading from Google Drive...")
+    print("🚀 Model file not found! Downloading...")
     try:
         gdown.download(GDRIVE_URL, model_path, quiet=False)
         print("✅ Model downloaded successfully.")
     except Exception as e:
         print(f"❌ Failed to download model: {e}")
 
-# ==================================
-# IMPORT MODEL RUNNER
-# ==================================
-
+# Import model
 sys.path.append(MODEL_DIR)
-from rcnn_model.scripts.rcnn_run import main, write_config
+from rcnn_run import main, write_config
+
+cfg = write_config()
 
 # ==================================
-# PAGE CONFIG
+# MAIN PREDICTION FUNCTION
 # ==================================
 
-st.set_page_config(
-    page_title="2D Floorplan Vectorizer",
-    layout="wide",
-    initial_sidebar_state="collapsed"
-)
+def predict(uploaded_file):
+    if uploaded_file is None:
+        return None, None, "No file uploaded."
 
-# ==================================
-# HEADER
-# ==================================
+    # Save uploaded image
+    input_bytes = uploaded_file.read()
+    img = Image.open(BytesIO(input_bytes)).convert("RGB")
+    input_filename = uploaded_file.name
+    uploaded_path = os.path.join(UPLOAD_DIR, input_filename)
+    img.save(uploaded_path)
+    print(f"✅ Image saved to {uploaded_path}")
 
-st.image(logo_path, width=250)
-st.markdown("<div class='header-title'>2D Floorplan Vectorizer</div>", unsafe_allow_html=True)
+    # 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")
 
-# ==================================
-# FILE UPLOAD SECTION
-# ==================================
+    output_json_path = os.path.join(JSON_DIR, output_json_name)
+    output_image_path = os.path.join(JSON_DIR, output_image_name)
 
-st.subheader("Upload your Floorplan Image")
-uploaded_file = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg"])
+    # Run model
+    main(cfg, uploaded_path, output_json_path, output_image_path)
 
-# Initialize session state
-if "processing_complete" not in st.session_state:
-    st.session_state.processing_complete = False
-if "json_output" not in st.session_state:
-    st.session_state.json_output = None
+    # 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
 
 # ==================================
-# IMAGE + JSON Layout
+# GRADIO UI
 # ==================================
 
-col1, col2 = st.columns([1, 2])
+with gr.Blocks() as demo:
+    gr.Markdown("<h1 style='text-align: center;'>🏠 Inovonics 2D Floorplan Vectorizer</h1>")
 
-# ==================================
-# MAIN LOGIC
-# ==================================
+    with gr.Row():
+        with gr.Column():
+            uploaded_file = gr.File(label="Upload your Floorplan Image", type="file")
+            run_button = gr.Button("Run Vectorizer 🔥")
+
+        with gr.Column():
+            output_image = gr.Image(label="🖼 Output Vectorized Image")
+            output_json = gr.JSON(label="🧾 Output JSON")
+
+    error_output = gr.Textbox(label="Error Message", visible=False)
+
+    run_button.click(
+        predict,
+        inputs=[uploaded_file],
+        outputs=[output_image, output_json, error_output]
+    )
 
-if uploaded_file is not None:
-    print("📤 File Uploaded:", uploaded_file.name)
-
-    image_bytes = uploaded_file.read()
-    img = Image.open(BytesIO(image_bytes)).convert("RGB")
-
-    uploaded_path = os.path.join(UPLOAD_DIR, uploaded_file.name)
-    with open(uploaded_path, "wb") as f:
-        f.write(uploaded_file.getbuffer())
-    print("✅ Uploaded file saved at:", uploaded_path)
-
-    with col1:
-        st.markdown("<div class='upload-container'>", unsafe_allow_html=True)
-        st.image(Image.open(uploaded_path), caption="Uploaded Image", use_container_width=True)
-        st.markdown("</div>", unsafe_allow_html=True)
-
-    with col2:
-        if not st.session_state.processing_complete:
-            status_placeholder = st.empty()
-            status_placeholder.info("⏳ Model is processing the uploaded image...")
-            progress_bar = st.progress(0)
-            status_text = st.empty()
-
-            # === 🔥 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...")
-
-            # Simulate progress
-            for i in range(1, 30):
-                time.sleep(0.01)
-                progress_bar.progress(i)
-                status_text.text(f"Preprocessing: {i}%")
-
-            # Run model
-            main(cfg, input_image, output_json_path, output_image_path)
-            print("✅ Model run complete.")
-
-            while not os.path.exists(output_json_path):
-                print("Waiting for JSON output...")
-                time.sleep(0.5)
-
-            for i in range(30, 100):
-                time.sleep(0.01)
-                progress_bar.progress(i)
-                status_text.text(f"Postprocessing: {i}%")
-
-            progress_bar.empty()
-            status_text.text("✅ Processing Complete!")
-            status_placeholder.success("✅ Model finished and JSON is ready!")
-
-            # Read generated JSON
-            if os.path.exists(output_json_path):
-                with open(output_json_path, "r") as jf:
-                    st.session_state.json_output = json.load(jf)
-                    print("📄 JSON Output Loaded Successfully.")
-            else:
-                st.session_state.json_output = {"error": "JSON output not generated."}
-                print("❌ JSON output missing.")
-
-            st.session_state.processing_complete = True
-
-        # ==================================
-        # DISPLAY OUTPUTS
-        # ==================================
-
-        out_col1, out_col2 = st.columns(2)
-
-        with out_col1:
-            if os.path.exists(output_image_path):
-                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)
-                    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",
-                    data=json_str,
-                    file_name="floorplan_output.json",
-                    mime="application/json"
-                )
-
-        with out_col2:
-            st.markdown("<div class='json-container'>", unsafe_allow_html=True)
-            st.json(st.session_state.json_output)
-            st.markdown("</div>", unsafe_allow_html=True)
-
-else:
-    st.warning("⚠️ No image uploaded yet.")
-    st.session_state.processing_complete = False
+demo.launch()

app_streamlit.py

@@ -0,0 +1,200 @@
+import streamlit as st
+import json
+import time
+from PIL import Image
+import os
+import sys
+import shutil
+import gdown
+from io import BytesIO
+
+# ==================================
+# SETUP
+# ==================================
+
+print("🚀 Streamlit App Starting...")
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# Setup Paths
+UPLOAD_DIR = "/tmp/uploads/"
+MODEL_DIR = os.path.join(BASE_DIR, "rcnn_model", "scripts")
+JSON_DIR = "/tmp/results/"
+OUTPUT_DIR = "/tmp/output/"
+SAMPLE_DIR = os.path.join(BASE_DIR, "rcnn_model", "sample")
+logo_path = os.path.join(BASE_DIR, "public", "logo.png")
+model_path = os.path.join(OUTPUT_DIR, "model_final.pth")
+
+# Google Drive file download link
+GOOGLE_DRIVE_FILE_ID = "1yr64AOgaYZPTcQzG6cxG6lWBENHR9qjW"
+GDRIVE_URL = f"https://drive.google.com/uc?id={GOOGLE_DRIVE_FILE_ID}"
+
+# Create necessary folders
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+os.makedirs(JSON_DIR, exist_ok=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+# ==================================
+# DOWNLOAD MODEL IF MISSING
+# ==================================
+
+if not os.path.exists(model_path):
+    print("🚀 Model file not found! Downloading from Google Drive...")
+    try:
+        gdown.download(GDRIVE_URL, model_path, quiet=False)
+        print("✅ Model downloaded successfully.")
+    except Exception as e:
+        print(f"❌ Failed to download model: {e}")
+
+# ==================================
+# IMPORT MODEL RUNNER
+# ==================================
+
+sys.path.append(MODEL_DIR)
+from rcnn_model.scripts.rcnn_run import main, write_config
+
+# ==================================
+# PAGE CONFIG
+# ==================================
+
+st.set_page_config(
+    page_title="2D Floorplan Vectorizer",
+    layout="wide",
+    initial_sidebar_state="collapsed"
+)
+
+# ==================================
+# HEADER
+# ==================================
+
+st.image(logo_path, width=250)
+st.markdown("<div class='header-title'>2D Floorplan Vectorizer</div>", unsafe_allow_html=True)
+
+# ==================================
+# FILE UPLOAD SECTION
+# ==================================
+
+st.subheader("Upload your Floorplan Image")
+uploaded_file = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg"])
+
+# Initialize session state
+if "processing_complete" not in st.session_state:
+    st.session_state.processing_complete = False
+if "json_output" not in st.session_state:
+    st.session_state.json_output = None
+
+# ==================================
+# IMAGE + JSON Layout
+# ==================================
+
+col1, col2 = st.columns([1, 2])
+
+# ==================================
+# MAIN LOGIC
+# ==================================
+
+if uploaded_file is not None:
+    print("📤 File Uploaded:", uploaded_file.name)
+
+    image_bytes = uploaded_file.read()
+    img = Image.open(BytesIO(image_bytes)).convert("RGB")
+
+    uploaded_path = os.path.join(UPLOAD_DIR, uploaded_file.name)
+    with open(uploaded_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
+    print("✅ Uploaded file saved at:", uploaded_path)
+
+    with col1:
+        st.markdown("<div class='upload-container'>", unsafe_allow_html=True)
+        st.image(Image.open(uploaded_path), caption="Uploaded Image", use_container_width=True)
+        st.markdown("</div>", unsafe_allow_html=True)
+
+    with col2:
+        if not st.session_state.processing_complete:
+            status_placeholder = st.empty()
+            status_placeholder.info("⏳ Model is processing the uploaded image...")
+            progress_bar = st.progress(0)
+            status_text = st.empty()
+
+            # === 🔥 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...")
+
+            # Simulate progress
+            for i in range(1, 30):
+                time.sleep(0.01)
+                progress_bar.progress(i)
+                status_text.text(f"Preprocessing: {i}%")
+
+            # Run model
+            main(cfg, input_image, output_json_path, output_image_path)
+            print("✅ Model run complete.")
+
+            while not os.path.exists(output_json_path):
+                print("Waiting for JSON output...")
+                time.sleep(0.5)
+
+            for i in range(30, 100):
+                time.sleep(0.01)
+                progress_bar.progress(i)
+                status_text.text(f"Postprocessing: {i}%")
+
+            progress_bar.empty()
+            status_text.text("✅ Processing Complete!")
+            status_placeholder.success("✅ Model finished and JSON is ready!")
+
+            # Read generated JSON
+            if os.path.exists(output_json_path):
+                with open(output_json_path, "r") as jf:
+                    st.session_state.json_output = json.load(jf)
+                    print("📄 JSON Output Loaded Successfully.")
+            else:
+                st.session_state.json_output = {"error": "JSON output not generated."}
+                print("❌ JSON output missing.")
+
+            st.session_state.processing_complete = True
+
+        # ==================================
+        # DISPLAY OUTPUTS
+        # ==================================
+
+        out_col1, out_col2 = st.columns(2)
+
+        with out_col1:
+            if os.path.exists(output_image_path):
+                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)
+                    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",
+                    data=json_str,
+                    file_name="floorplan_output.json",
+                    mime="application/json"
+                )
+
+        with out_col2:
+            st.markdown("<div class='json-container'>", unsafe_allow_html=True)
+            st.json(st.session_state.json_output)
+            st.markdown("</div>", unsafe_allow_html=True)
+
+else:
+    st.warning("⚠️ No image uploaded yet.")
+    st.session_state.processing_complete = False

requirements.txt

@@ -8,4 +8,5 @@ shapely
 matplotlib
 labelme2coco
 numpy
-from_root
+from_root
+gradio