code cleanup

Dockerfile

@@ -12,7 +12,7 @@ RUN apt-get update && apt-get install -y \
     ffmpeg \
     && rm -rf /var/lib/apt/lists/*
 
-# Set Matplotlib temp folder
+# Set Matplotlib temp folder. This is needed to avoid warnings for missing fonts
 ENV MPLCONFIGDIR=/tmp/matplotlib
 
 # Set working directory
@@ -33,4 +33,4 @@ RUN pip install 'git+https://github.com/facebookresearch/detectron2.git'
 COPY . .
 
 # Run
-CMD ["python3", "app.py"]
+CMD ["python3", "app.py"]

app.py

@@ -5,15 +5,9 @@ import sys
 import json
 import shutil
 import gdown
-import time
 from PIL import Image
-from io import BytesIO
 
-# ==================================
-# SETUP
-# ==================================
-
-print("🚀 Gradio App Starting...")
+print("Gradio App Starting...")
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
@@ -34,22 +28,14 @@ os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(JSON_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# Copy logo.png to /tmp
-logo_src = os.path.join(BASE_DIR, "public", "logo.png")
-logo_dst = "/tmp/logo.png"
-
-# Only copy if not already there
-if not os.path.exists(logo_dst):
-    shutil.copy(logo_src, logo_dst)
-
 # Download model if missing
 if not os.path.exists(model_path):
-    print("🚀 Model file not found! Downloading...")
+    print("Model file not found! Downloading...")
     try:
         gdown.download(GDRIVE_URL, model_path, quiet=False, use_cookies=False)
-        print("✅ Model downloaded successfully.")
+        print("Model downloaded successfully.")
     except Exception as e:
-        print(f"❌ Failed to download model: {e}")
+        print(f"Failed to download model: {e}")
 
 # Import model
 sys.path.append(MODEL_DIR)
@@ -57,16 +43,11 @@ from rcnn_model.scripts.rcnn_run import main, write_config
 
 cfg = write_config()
 
-# ==================================
-# MAIN PREDICTION FUNCTION
-# ==================================
-
 def predict(uploaded_file_path):
-    print("Inside Predict:" + uploaded_file_path)
     if uploaded_file_path is None:
         return None, None, "No file uploaded.", None
 
-    # Save uploaded file to temp
+    # Save uploaded file to tmp folder
     uploaded_path = os.path.join(UPLOAD_DIR, "input_image.png")
     shutil.copy(uploaded_file_path, uploaded_path)
 
@@ -95,17 +76,11 @@ def predict(uploaded_file_path):
 
     return result_img, json.dumps(result_json, indent=2), None, download_json_path, uploaded_path
 
-# ==================================
-# GRADIO UI
-# ==================================
-
 with gr.Blocks() as demo:
-    # Header
     with gr.Row():
         gr.Markdown(
             f"""
     <div style='display: flex; align-items: center; justify-content: center;'>
-        <img src='file/tmp/logo.png' style='height: 50px; margin-right: 10px;'/>
         <h1>Inovonics 2D Floorplan Vectorizer</h1>
     </div>
     """
@@ -118,19 +93,12 @@ with gr.Blocks() as demo:
             run_button = gr.Button("Run Vectorizer 🔥")
 
         with gr.Column():
-            output_image = gr.Image(label="🖼 Output Vectorized Image")
-            download_button = gr.File(label="⬇️ Download JSON", visible=True)
-            output_json = gr.JSON(label="🧾 Output JSON")
+            output_image = gr.Image(label="Output Vectorized Image")
+            download_button = gr.File(label="Download JSON", visible=True)
+            output_json = gr.JSON(label="Output JSON")
 
     error_output = gr.Textbox(label="Error Message", visible=False)
 
-    # Logic binding
-    # run_button.click(
-    #     predict,
-    #     inputs=[uploaded_file],
-    #     outputs=[output_image, output_json, error_output, download_button, uploaded_image_display]
-    # )
-
     run_button.click(
     lambda x: (x, gr.update(interactive=False)),
     inputs=[uploaded_file],

app_streamlit.py

@@ -1,200 +0,0 @@
-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

rcnn_model/extraction/floorplan_sampler.py

@@ -13,13 +13,6 @@ import pylab
 pylab.rcParams['figure.figsize'] = (128.0, 160.0)
 from rcnn_model.utils.floorplan_vectorizer_utils import get_image_size, draw_from_coco
 
-
-# sys.path.append(str(from_root("utils")))
-# from floorplan_vectorizer_utils import get_image_size, draw_from_coco
-
-
-### Main functionality ###
-
 data_directory_root = str(from_root("dataset/"))+"/"
 category_filter = [2]
 image_sample_room_count_threshold = 4

rcnn_model/extraction/svg_to_json.py

@@ -13,13 +13,6 @@ from from_root import from_root
 from rcnn_model.preprocessing.cleaning_single_image import preprocess_image
 from rcnn_model.utils.floorplan_vectorizer_utils import get_image_size, draw_from_coco
 
-
-
-# sys.path.append(str(from_root("preprocessing")))
-# from cleaning_images import preprocessing
-# sys.path.append(str(from_root("utils")))
-# from floorplan_vectorizer_utils import get_image_size, draw_from_coco
-
 ### After running, its split with https://github.com/akarazniewicz/cocosplit
 ### This may or may not be temporary
 

rcnn_model/preprocessing/cleaning_single_image.py

@@ -1,37 +1,20 @@
-# single_image_cleaning.py
-
+#This file is a version of the original cleaning_images.py fiele, but it has been modified to only process a single image.
 import cv2
 import numpy as np
-import os
 
 def preprocess_image(image_path):
-    """
-    Preprocess a single floorplan image: denoising, CLAHE, edge enhancement.
-    """
-    print(f"🧹 Preprocessing image: {image_path}")
+    print(f" Preprocessing image: {image_path}")
     
     image = cv2.imread(image_path)
     if image is None:
-        print(f"❌ Error: Could not read image from {image_path}")
+        print(f"Error: Could not read image from {image_path}")
         return None
-
-    # Convert to grayscale
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-
-    # Apply Gaussian blur
     denoisy_img = cv2.GaussianBlur(gray, (5, 5), 0)
-
-    # Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)
     clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
     enhanced = clahe.apply(denoisy_img)
-
-    # Apply threshold
     _, thresholded = cv2.threshold(enhanced, 150, 255, cv2.THRESH_BINARY)
-
-    # Detect edges
     edges = cv2.Canny(thresholded, 100, 220, apertureSize=3)
-
-    # Detect lines and draw them
     output_img = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
     lines = cv2.HoughLinesP(edges, rho=1, theta=np.pi / 180, threshold=50,
                             minLineLength=35, maxLineGap=5)
@@ -41,8 +24,7 @@ def preprocess_image(image_path):
             x1, y1, x2, y2 = line[0]
             cv2.line(output_img, (x1, y1), (x2, y2), (210, 210, 210), 1)
 
-    # Blend the original image with line-enhanced version
     blended_image = cv2.addWeighted(image, 0.7, output_img, 0.3, 0)
 
-    print(f"✅ Preprocessing complete for: {image_path}")
+    print(f"Preprocessing complete for: {image_path}")
     return blended_image

rcnn_model/preprocessing/svg_to_yolo.py

@@ -19,7 +19,6 @@ def extract_svg_elements(svg_file):
     svg_width = float(root.get("width", "1"))
     svg_height = float(root.get("height", "1"))
 
-    # floorplans = {}
     floorplans = {"Door": [], "Window": [], "Space": []}
     for floorplan in root.findall(".//svg:g[@class]", namespaces=namespace):
         class_attr = floorplan.get("class", "").strip()
@@ -54,11 +53,8 @@ def extract_svg_elements(svg_file):
     return floorplans, svg_width, svg_height
 
 def get_bounding_box(polygons, svg_width, svg_height):
-    """Compute YOLO bounding box from polygon points."""
     all_x, all_y = [], []
     print(polygons)
-    # for polygon in polygons:
-        # print(polygon)
     points = polygons.strip().split(" ")
     for point in points:
         x, y = map(float, point.split(","))
@@ -78,7 +74,6 @@ def get_bounding_box(polygons, svg_width, svg_height):
     return (x_center, y_center, width, height)
 
 def save_yolo_annotations(floorplans, output_dir, filename):
-    """Save extracted bounding boxes in YOLO format."""
     os.makedirs("dataset", exist_ok=True)
     os.makedirs(output_dir, exist_ok=True)
 

rcnn_model/scripts/rcnn_config.py

@@ -26,8 +26,6 @@ def write_config():
     cfg.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION = 14
     cfg.MODEL.ROI_MASK_HEAD.NUM_CONV = 3
     cfg.TEST.DETECTIONS_PER_IMAGE = 120
-# Added this extra line
-    # cfg.OUTPUT_DIR = str(from_root("rcnn_model/output"))
 
     return cfg
 

rcnn_model/scripts/rcnn_eval.py

@@ -12,14 +12,10 @@ from detectron2.utils.visualizer import Visualizer
 import random
 import matplotlib.pyplot as plt
 import time
-from rcnn_config import write_config
 from from_root import from_root
 from rcnn_model.utils.floorplan_vectorizer_utils import check_image_size_thresh
 
 
-# sys.path.append(str(from_root("utils")))
-# from floorplan_vectorizer_utils import check_image_size_thresh
-
 results_directory = str(from_root("results"))+"/"
 max_image_size = 700*500
 

rcnn_model/scripts/rcnn_run.py

@@ -6,24 +6,11 @@ import matplotlib.pyplot as plt
 import torch
 from datetime import datetime
 from rcnn_config import write_config
-import sys
 from from_root import from_root
 from rcnn_model.preprocessing.cleaning_single_image import preprocess_image
 from rcnn_model.utils.floorplan_vectorizer_utils import draw_from_coco, bitmask_to_polygon
 from rcnn_model.extraction.annotation_builder import AnnotationBuilder as AnnBuild
 
-
-
-# sys.path.append(str(from_root("preprocessing")))
-# from cleaning_images import preprocessing
-# sys.path.append(str(from_root("utils")))
-# from floorplan_vectorizer_utils import draw_from_coco, bitmask_to_polygon
-# sys.path.append(str(from_root("dataset/extraction_scripts")))
-# from annotation_builder import AnnotationBuilder as AnnBuild
-
-# results_directory = str(from_root("results"))+"/"
-# sample_data_directory = str(from_root("models/rcnn/sample_data"))+"/"
-
 results_directory = "/tmp/results/"
 sample_data_directory = "rcnn_model/sample/"
 model_directory = "/tmp/output/"
@@ -34,15 +21,13 @@ def main(cfg,img_source_path, coco_dest_filename, val_img_dest_filename):
     print("Inside main: " + img_source_path)
     #configure model
     cfg.DATALOADER.NUM_WORKERS = 1
-    cfg.SOLVER.IMS_PER_BATCH = 1 
-    # Changed this file here                                          
+    cfg.SOLVER.IMS_PER_BATCH = 1                                        
     cfg.MODEL.WEIGHTS = os.path.join(model_directory, "model_final.pth")
     cfg.MODEL.DEVICE = "cpu"
     predictor = DefaultPredictor(cfg)
 
     #run
     prediction_runner(img_source_path, results_directory+coco_dest_filename, results_directory+val_img_dest_filename, predictor)
-    #prediction_runner(img_source_path, results_directory+coco_dest_filename, results_directory+val_img_dest_filename, predictor, segmented_prediction=True, scale_factor=.5)
     print("SAVED to "+results_directory+coco_dest_filename+" and "+results_directory+val_img_dest_filename)                                                  
 
 
@@ -174,4 +159,5 @@ if __name__ == "__main__":
     print("Inside Main Function Call")
     cfg = write_config()
     run_index=-1
+    #Sample run for a single image
     main(cfg,sample_data_directory+"F1_original.png","cubicasa_result_run_"+str(run_index)+".json","cubicasa_result_run_"+str(run_index)+".png")

rcnn_model/utils/coco_to_inovonics_json.py

@@ -2,7 +2,6 @@
 from pycocotools.coco import COCO
 import json
 from inovonics_ann_builder import InovonicsAnnotationBuilder as InovAnnBuild
-from from_root import from_root
 
 def main(coco_source_path, inovonics_anns_dest_path, img_ids=[1]):
     coco = COCO(coco_source_path)

rcnn_model/utils/floorplan_vectorizer_utils.py

@@ -39,7 +39,6 @@ def draw_from_coco(id,coco,annotated_img_dest_path,category_filter = [0,1,2,3],
 
 def get_blank_image(width,height):
     blank = io.imread(str(from_root("models/rcnn/westmoor_check/white_bg.png")))
-    #return cv2.resize(blank, (0,0), fx=width, fy=height)
     return blank[:width,:height,:]
 
 

rcnn_model/utils/inovonics_ann_builder.py

@@ -1,6 +1,4 @@
 import json
-import numpy as np
-from datetime import datetime
 import random
 
 version_number = "0.0.1"