initial commit

.gitignore

@@ -0,0 +1,27 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Caches
+.cache/
+.ipynb_checkpoints/
+
+# Virtual environment
+venv/
+.env/
+*.env
+
+# MacOS system files
+.DS_Store
+
+# Model files
+*.pth
+*.ckpt
+
+# Streamlit temp files
+.streamlit/
+
+# Uploads and results
+rcnn_model/uploads/
+rcnn_model/results/

Dockerfile

@@ -0,0 +1,36 @@
+# Use Python base image
+FROM python:3.10-slim
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    build-essential \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Clone and install Detectron2
+RUN pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu
+RUN git clone https://github.com/facebookresearch/detectron2.git && \
+    pip install -e detectron2
+
+# Copy your app code
+COPY . .
+
+# Expose Streamlit default port
+EXPOSE 7860
+
+# Streamlit Environment Variables
+ENV STREAMLIT_SERVER_PORT=7860
+ENV STREAMLIT_SERVER_HEADLESS=true
+ENV STREAMLIT_SERVER_ENABLECORS=false
+
+# Command to run the app
+CMD ["streamlit", "run", "app.py"]

README.md

@@ -1,10 +1,75 @@
+# 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.
+
 ---
-title: Floorplan Vectorizer
-emoji: 📉
-colorFrom: pink
-colorTo: gray
-sdk: docker
-pinned: false
+
+## How to Run the App
+
+1. **Clone the repository:**
+
+    ```bash
+    git clone <this-repo-link>
+    cd inovonics-ui-vectorizer
+    ```
+
+2. **Install the required Python packages:**
+
+    ```bash
+    pip install -r requirements.txt
+    ```
+
+3. **Download the pretrained model:**
+
+    - Download `model_final.pth` from [Google Drive here](https://drive.google.com/file/d/1yr64AOgaYZPTcQzG6cxG6lWBENHR9qjW/view?usp=sharing).
+    - Place it inside:
+
+      ```plaintext
+      inovonics-ui-vectorizer/rcnn_model/output/model_final.pth
+      ```
+
+4. **Run the app:**
+
+    ```bash
+    streamlit run app.py
+    ```
+
+5. Open your browser at [http://localhost:8501](http://localhost:8501) to start using the app!
+
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## Project Structure
+
+```plaintext
+inovonics-ui-vectorizer/
+├── app.py                     # Streamlit frontend app
+├── public/
+│   └── logo.png                # App logo
+├── rcnn_model/
+│   ├── extraction/             # Extract information from uploaded png image
+│   │   └── annotation_builder.py      
+│   │   └── floorplan_sampler.py
+│   │   └── from_labelme_runner.py
+│   │   └── svg_to_json.py   
+│   ├── output/                 # Empty folder while cloning. Place the pth file here
+│   ├── preprocessing/          # Preprocess the image before sending to model
+│   │   └── cleaning_images.py  
+│   │   └── cleaning_single_image.py 
+│   │   └── splitting_dataset.py
+│   │   └── svg_to_yolo.py    
+│   ├── results/                # Empty folder while cloning. The resulting image and JSON will be stored here
+│   ├── sample/                 # Sample images for the model       
+│   ├── scripts/                # Model training, evaluation and inference. Streamlit runs the rcnn_run.py file from the frontend
+│   │   └── rcnn_config.py    
+│   │   └── rcnn_eval.py  
+│   │   └── rcnn_full_tuner.py 
+│   │   └── rcnn_run.py  
+│   │   └── rcnn_train.py     
+│   ├── uploads/                # Temporary folder for streamlit to store the user uploaded image
+│   ├── utils/                  # Utility functions during model train and preprocessing
+│   │   └── coco_to_inovonics_json.py
+│   │   └── floorplan_vectorizer_utils.py
+│   │   └── inovonics_ann_builder.py
+├── README.md                   # (this file)
+├── requirements.txt            # Python dependencies
+└── .gitignore                  # Files to ignore during Git commits

app.py

@@ -0,0 +1,177 @@
+import streamlit as st
+import json
+import time
+from PIL import Image
+import os
+import sys
+import gdown
+
+st.set_page_config(
+    page_title="2D Floorplan Vectorizer",
+    layout="wide",
+    initial_sidebar_state="collapsed"
+)
+
+print("Streamlit App Starting...")
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# Setup Paths
+UPLOAD_DIR = os.path.join(BASE_DIR, "rcnn_model", "uploads")
+MODEL_DIR = os.path.join(BASE_DIR, "rcnn_model", "scripts")
+JSON_DIR = os.path.join(BASE_DIR, "rcnn_model", "results")
+OUTPUT_DIR = os.path.join(BASE_DIR, "rcnn_model", "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}"
+
+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}")
+
+sys.path.append(MODEL_DIR)
+from rcnn_model.scripts.rcnn_run import main, write_config
+
+
+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
+)
+
+st.image(logo_path, width=250)
+st.markdown("<div class='header-title'>2D Floorplan Vectorizer</div>", unsafe_allow_html=True)
+st.subheader("Upload your Floorplan Image")
+uploaded_file = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg"])
+
+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
+
+col1, col2 = st.columns([1, 2])
+
+if uploaded_file is not None:
+    print("File Uploaded:", uploaded_file.name)
+
+    # Save uploaded file
+    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)
+
+    # Display uploaded image
+    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()
+
+            # Run Model
+            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")
+
+            cfg = write_config()
+            print("Model config created. Running model...")
+
+            # Simulate progress bar
+            for i in range(1, 30):
+                time.sleep(0.01)
+                progress_bar.progress(i)
+                status_text.text(f"Preprocessing: {i}%")
+
+            main(cfg, input_image, output_json_name, output_image_name)
+            print("Model run complete.")
+
+            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...")
+                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
+
+        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:
+                    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(
+                    label="Download JSON",
+                    data=json_str,
+                    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)
+            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

packages.txt

@@ -0,0 +1,2 @@
+libgl1-mesa-glx
+libglib2.0-0

rcnn_model/extraction/annotation_builder.py

@@ -0,0 +1,69 @@
+import json
+import numpy as np
+from datetime import datetime
+
+version_number = "0.0.1"
+
+class AnnotationBuilder:
+    #creates the base structure of the coco format
+    def __init__(self):
+        self.licenses = []
+        self.categories = [{"id": 0,
+                             "name": "Wall",
+                             "supercategory": "none"},
+                           { "id": 1,
+                             "name": "Door",
+                             "supercategory": "none"},
+                           {"id": 2,
+                             "name": "Room",
+                             "supercategory": "none"},
+                           {"id": 3,
+                            "name": "Window",
+                            "supercategory": "none"}
+                          ]
+        self.images = []
+        self.annotations = []
+
+    def set_info(self, description, data_source_name, data_source_url, data_source_creation_date):
+        self.info = [{"year":2025,
+                        "version":version_number,
+                        "description":description,
+                        "contributor":data_source_name,
+                        "url":data_source_url,
+                        "date_created":data_source_creation_date.strftime("%Y-%m-%dT%H:%M:%S")}]
+
+    def add_license(self, license_name, license_url):
+        self.licenses.append({"id":len(self.licenses),
+                          "url":license_url,
+                          "name":license_name})
+
+    def add_image(self, filename, width, height):
+        id = len(self.images)
+        self.images.append({"id":id,
+                            "width":width,
+                            "height":height,
+                            "file_name":filename, #filename should be the image path relative to the cocofile's path
+                            "license":0,
+                            "date_captured":datetime.now().strftime("%Y-%m-%dT%H:%M:%S")})
+        return id
+
+    def add_annotation(self, image_id, category_id, poly):
+        id = len(self.annotations)
+        segmentation = np.array(poly.exterior.coords).astype(int).ravel().tolist()[:-2]
+        x,y,x2,y2 = tuple(map(int, poly.bounds))
+        self.annotations.append({"id":id,
+                                 "image_id":image_id,
+                                 "category_id":category_id,
+                                 "segmentation":[segmentation], 
+                                 "area":poly.area,
+                                 "bbox":[x,y,x2-x,y2-y],
+                                 "iscrowd":0})
+        return id, poly
+    
+    def final_output(self):
+        return {"info":self.info, "licenses":self.licenses, "categories":self.categories, "images":self.images, "annotations":self.annotations}
+    
+    def save_file(self, filepath):
+        coco_file = open(filepath,'w')
+        json.dump(self.final_output(),coco_file,indent=4)
+        coco_file.close()

rcnn_model/extraction/floorplan_sampler.py

@@ -0,0 +1,257 @@
+import cv2
+import numpy as np
+import copy
+from pycocotools.coco import COCO
+import shapely
+from shapely import geometry
+import sys
+import random
+from datetime import datetime
+from from_root import from_root
+from annotation_builder import AnnotationBuilder as AnnBuild
+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
+min_sample_size = 400
+max_sample_size = 800
+samples_per_image = 30
+
+def main():
+    sample_from_labelme2coco_dataset("train",data_directory_root+"annotations/","sample_data/","validation_images/")
+    sample_from_labelme2coco_dataset("val",data_directory_root+"annotations/","sample_data/","validation_images/")
+
+
+### Core Sampling Logic ###
+
+#sample from dataset cocofile created by labelme2coco
+#dataset_name should only be "train", "val", or "dataset" based on labelme2coco's output naming conventions
+def sample_from_labelme2coco_dataset(dataset_name,annotation_source_dir,sample_img_dest_dir,validation_img_dest_dir=""):
+    #initialize annbuilder
+    ann_builder = AnnBuild()
+    ann_builder.set_info("manual annotations of Inovonics and university provided data","inovonData","NA",datetime(2019,5,24))
+    ann_builder.add_license("TODO", "TODO")
+    coco = COCO(annotation_source_dir+dataset_name+".json")
+    print("Coco Loaded")
+
+    #reading
+    for img_id in coco.getImgIds():
+        take_samples_from_image(ann_builder, img_id, coco, sample_img_dest_dir)
+
+    #save
+    ann_builder.save_file(annotation_source_dir+dataset_name+"_sampled_data.json")
+
+    #validation images
+    if(validation_img_dest_dir != ""):
+        validation_coco = COCO(annotation_source_dir+dataset_name+"_sampled_data.json")
+        validation_images(dataset_name,validation_coco,validation_img_dest_dir)
+
+
+def take_samples_from_image(ann_builder, img_id, coco, img_dest):
+    #set up image name
+    source_img_filename = coco.imgs[img_id]['file_name']
+    source_img = cv2.imread(data_directory_root+source_img_filename,cv2.IMREAD_COLOR)
+    img_name = source_img_filename[1:-4]
+    img_name = img_name[img_name.index("/"):]
+    img_name = img_name[1:]
+
+    #set up mirroring
+    mirrored_imgs = [source_img, np.fliplr(source_img), np.flipud(source_img), np.flipud(np.fliplr(source_img))]
+    mirror_tags = ["","_h","_v","_hv"]
+    print("Processing image "+str(img_id)+": "+img_name)
+
+    #run sampler
+    for m in range(0,len(mirrored_imgs)):
+        #load mirror of image
+        img = mirrored_imgs[m]
+        tag = mirror_tags[m]
+        mirrored_anns = mirror_coco_coordinates(coco,img_id,m)
+
+        if len(mirrored_anns) > 10:
+            #collect samples
+            for i in range(0,samples_per_image):
+                img_dest_path = data_directory_root+img_dest+img_name+tag+"_"+str(i)+".png"
+                take_sample(ann_builder, mirrored_anns, img, img_dest_path)
+        else:
+            img_dest_path = data_directory_root+img_dest+img_name+tag+".png"
+            take_full_image(ann_builder, mirrored_anns, img, img_dest_path)
+
+
+def take_sample(ann_builder, annotations, img, img_dest_path):
+    #Take a random sample with at least a certain number of room bounding boxes overlapping
+    sample_annotations, cropped, room_count = random_sample_selection(annotations,img)
+    while(room_count < image_sample_room_count_threshold):
+        sample_annotations, cropped, room_count = random_sample_selection(annotations,img)
+    
+    #sav the cropped image portion of the final sample
+    cv2.imwrite(img_dest_path, cropped)
+    print("  sample saved to "+img_dest_path)
+    sample_width, sample_height = get_image_size(img_dest_path)
+    sampled_img_id = ann_builder.add_image(img_dest_path, sample_width, sample_height)
+
+    #crop annotations
+    crop_area = define_crop_area(0,0,sample_width, sample_height)
+    add_cropped_annotations(ann_builder, sampled_img_id, sample_annotations, crop_area)
+
+
+def take_full_image(ann_builder, annotations, img, img_dest_path):
+    #sav the cropped image portion of the final sample
+    cv2.imwrite(img_dest_path, img)
+    print("  whole image saved to "+img_dest_path)
+    width, height = get_image_size(img_dest_path)
+    img_id = ann_builder.add_image(img_dest_path, width, height)
+    for ann in annotations:
+        ann_builder.add_annotation(img_id, ann["category_id"], segmentation_to_polygon(ann["segmentation"]))
+
+
+def random_sample_selection(annotations, img):
+    #get bounds of original image
+    init_width = len(img)
+    init_height = len(img[0])
+
+    #randomly select a rectangle
+    sample_x = random.randrange(0,init_width-min_sample_size-1)
+    sample_y = random.randrange(0,init_height-min_sample_size-1)
+    sample_width = random.randrange(min_sample_size,min(max_sample_size, init_width-sample_x))
+    sample_height = random.randrange(min_sample_size,min(max_sample_size, init_height-sample_y))
+
+    #create cropped image and offset annotation coordinates immediately for easier data transfer
+    cropped = img[sample_y:sample_y+sample_height,sample_x:sample_x+sample_width]
+    sampled_annotations, room_count = offset_annotation_coordinates(annotations,sample_x,sample_y,sample_width,sample_height)
+
+    #return values
+    return sampled_annotations, cropped, room_count
+
+
+### Annotation Cropping ###
+
+def define_crop_area(x, y, width, height):
+    cropped_area = geometry.Polygon([(x,y),
+                                    (x+width,y),
+                                    (x+width,y+height),
+                                    (x,y+height)])
+    return cropped_area
+
+
+def add_cropped_annotations(ann_builder, img_id, annotations, cropped_area):
+    for ann in annotations:
+        #get intersecting area
+        poly = crop_polygon(ann["segmentation"],cropped_area,ann)
+
+        #handle convex rooms that weren't split
+        if(isinstance(poly,geometry.Polygon)):
+            ann_builder.add_annotation(img_id, ann["category_id"], poly)
+        
+        #handle concave rooms that were split
+        elif(isinstance(poly,geometry.GeometryCollection) or isinstance(poly,geometry.MultiPolygon)):
+            for subpoly in poly.geoms:
+                if(isinstance(subpoly,geometry.Polygon)):
+                    ann_builder.add_annotation(img_id, ann["category_id"], subpoly)
+
+
+def crop_polygon(segmentation,crop_area,id):
+    #reformat into shapely geometry Polygon
+    poly = segmentation_to_polygon(segmentation)
+
+    #check shape validity (most common error is self overlapping)
+    if(not shapely.is_valid(poly)):
+        print(id)
+        print(shapely.is_valid_reason(poly))
+        return None
+    
+    #check size and return
+    cropped_poly = shapely.intersection(poly,crop_area)
+    if(cropped_poly.area > 0):
+        return cropped_poly
+    else:
+        return None
+
+def segmentation_to_polygon(segmentation):
+    points = np.array(segmentation[0])
+    points = points.reshape(int(len(segmentation[0])/2),2)
+    return geometry.Polygon([[p[0], p[1]] for p in points])
+
+
+### Applying Geometry Modifications
+
+def mirror_coco_coordinates(coco, img_id, mirroring_index):
+    #instantiate copy
+    original_annotations = coco.imgToAnns[img_id]
+    annotations = copy.deepcopy(original_annotations)
+    new_annotations = []
+
+    #get width
+    width = coco.imgs[img_id]['width']
+    height = coco.imgs[img_id]['height']
+
+    #apply mirroring
+    for ann in annotations:
+        if(ann["category_id"] in category_filter):
+            for i in range(0,len(ann['bbox'])):
+                apply_mirroring_to_coord(ann['bbox'], i, mirroring_index, width, height)
+            for i in range(0,len(ann['segmentation'][0])):
+                apply_mirroring_to_coord(ann['segmentation'][0], i, mirroring_index, width, height)
+            new_annotations.append(ann)
+    
+    return new_annotations
+
+
+def apply_mirroring_to_coord(array, index, mirroring_index, width, height):
+    if(index%2 == 0):
+        if(mirroring_index%2==1):
+            array[index] = width-array[index]
+    else:
+        if(mirroring_index>1):
+            array[index] = height-array[index]
+            
+
+def offset_annotation_coordinates(original_annotations,x_offset,y_offset,width,height):
+    #instantiate copy
+    annotations = copy.deepcopy(original_annotations)
+    new_annotations = []
+    room_count = 0
+
+    #apply offfset
+    for ann in annotations:
+        if(check_bounding_box_overlap(ann['bbox'], x_offset, y_offset, width, height)):
+            room_count += 1
+        for i in range(0,len(ann['bbox'])):
+            apply_offset_to_coord(ann['bbox'], i, x_offset, y_offset)
+        for i in range(0,len(ann['segmentation'][0])):
+            apply_offset_to_coord(ann['segmentation'][0], i, x_offset, y_offset)
+        new_annotations.append(ann)
+
+    return new_annotations, room_count
+
+
+def apply_offset_to_coord(array, index, x_offset, y_offset):
+    if(index%2 == 0):
+        array[index] -= x_offset
+    else:
+        array[index] -= y_offset
+
+
+def check_bounding_box_overlap(bbox, x_offset, y_offset, width, height):
+    boundary_threshold = 25 #so a row of rooms with a one pixel sliver within the area don't get counted
+    within_horizontal_bounds = max(bbox[0],bbox[2]) >= x_offset+boundary_threshold and min(bbox[0],bbox[2]) <= x_offset+width-boundary_threshold
+    within_vertical_bounds = max(bbox[1],bbox[3]) >= y_offset+boundary_threshold and min(bbox[1],bbox[3]) <= y_offset+height-boundary_threshold
+    return within_horizontal_bounds and within_vertical_bounds
+
+
+### Validation Display ###
+
+def validation_images(dataset_name,coco,validation_img_target_dir):
+        count = 1
+        for i in np.random.choice(coco.getImgIds(),8):
+            draw_from_coco(i,coco,data_directory_root+validation_img_target_dir+dataset_name+"_sampling_validation_"+str(count)+".png")
+            count+=1

rcnn_model/extraction/from_labelme_runner.py

@@ -0,0 +1,61 @@
+import os
+import labelme2coco
+from PIL import Image
+from from_root import from_root
+import floorplan_sampler
+import json
+
+
+def main():
+    convert_from_labelme()
+
+    handle_jpeg_files("annotations/train.json")
+    handle_jpeg_files("annotations/val.json")
+
+    print("NOW FOR SAMPLING")
+    floorplan_sampler.main()
+
+def convert_from_labelme():
+    os.chdir(str(from_root("dataset")))
+    labelme_source_dir = "labelme_data"
+    annotation_dest_dir = "annotations"
+    training_split_percentage = .8
+    labelme2coco.convert(labelme_source_dir, annotation_dest_dir, training_split_percentage, category_id_start=0)
+
+
+def handle_jpeg_files(coco_path):
+    #open file
+    file = open(coco_path,"r+")
+    coco = json.load(file)
+
+    #find and edit jpeg images
+    for image in coco["images"]:
+        img_name = image["file_name"]
+        if(".jpg" in img_name or ".jpeg" in img_name):
+            new_img_name = convert_to_png(img_name)
+            image["file_name"]=new_img_name
+
+    #save
+    file.seek(0)
+    json.dump(coco, file, indent="  ")
+    file.close()
+
+
+def convert_to_png(img_path):
+    #load image
+    img = Image.open(img_path)
+
+    #remove .jpg or .jpeg from path
+    if(".jpeg" in img_path):
+        img_path = img_path[0:-5]
+    else:
+        img_path = img_path[0:-4]
+    
+    #add .png and save
+    img_path += ".png"
+    img.save(img_path)
+    return img_path
+
+
+
+main()

rcnn_model/extraction/svg_to_json.py

@@ -0,0 +1,194 @@
+from xml.dom import minidom
+import cv2
+import numpy as np
+import math
+from annotation_builder import AnnotationBuilder as AnnBuild
+from pycocotools.coco import COCO
+from shapely import geometry
+from datetime import datetime
+import os
+import sys
+import random
+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
+
+### Main functionality ###
+
+scale_factor = .5
+dataset_root = str(from_root("dataset"))+"/"
+
+def main():
+    extract_all_cubicasa_anns(True)
+
+
+def extract_all_cubicasa_anns(export_image=False):
+    #initialize annotation builder
+    ann_builder = AnnBuild()
+    ann_builder.set_info("converted from cubicasa 5k SVG file","cubicasa 5k","https://github.com/cubicasa/cubicasa5k",datetime(2019,5,24))
+    ann_builder.add_license("Creative Commons Attribution-NonCommercial 4.0 International License", "http://creativecommons.org/licenses/by-nc/4.0/")
+    #iterate through cubicasa files
+    for name in os.listdir(str(from_root(dataset_root+"cubicasa_data/"))):
+        process_cubicasa_image(ann_builder, name)
+    #save data
+    print("SAVING TO annotations/cubicasa_coco.json")
+    ann_builder.save_file(str(from_root(dataset_root+"annotations/cubicasa_coco.json")))
+    if(export_image):
+        save_validation_images(str(from_root(dataset_root+"annotations/cubicasa_coco.json")))
+
+
+def process_cubicasa_image(ann_builder, name):
+    #load and preprocess image
+    print("\nprocessing "+name)
+    source_img_path = str(from_root(dataset_root+"cubicasa_data/"+name+"/F1_scaled.png"))
+    processed_img_path = str(from_root(dataset_root+"preprocessed/casa"+name+".png"))
+    apply_preprocessing(source_img_path, processed_img_path)
+
+    #load svg
+    source_svg_path = str(from_root(dataset_root+"cubicasa_data/"+name+"/model.svg"))
+    print("from "+source_svg_path)
+    print("image in "+processed_img_path)
+
+    #extract data from svg
+    try:
+        width, height = get_image_size(processed_img_path)
+        ann_builder = process_cubicasa(ann_builder, source_svg_path, processed_img_path, width, height)
+    except:
+        print("ERROR while extracting "+name)
+        print(sys.exc_info())
+
+
+def find_svg(path, name):
+    for file in os.listdir(path):
+        found_name = file.startswith(name+"_gt_")
+        if(found_name):
+            found_svg = file.endswith(".svg")
+            if(found_svg):
+                return path+file
+
+
+def process_cubicasa(ann_builder, sourve_svg_path, source_img_path, width, height):
+    #Get points
+    doc = minidom.parse(sourve_svg_path)
+    walls = extract_casa_elements_with_id("Wall",doc)
+    windows = extract_casa_elements_with_id("Window",doc)
+    doors = extract_casa_elements_with_id("Door",doc)
+    doc.unlink()
+    #export to JSON and potentially imges for visual confirmation that the process works
+    ann_builder = export_to_builder_casa(ann_builder,source_img_path,width,height,walls,doors,windows)
+    return ann_builder
+
+
+### Coco Formatting/Export ###
+
+def export_to_builder_casa(ann_builder,source_img,width,height,walls,doors,windows):
+    #initialization
+    id = ann_builder.add_image(source_img, width, height)
+    #walls
+    wall_polygons = get_features_from_ann_set(walls)
+    door_polygons = get_features_from_ann_set(doors)
+    window_polygons = get_features_from_ann_set(windows)
+    features = wall_polygons + door_polygons + window_polygons
+    rooms = create_rooms_from_features(features, width, height)
+    for poly in rooms.geoms:
+        ann_builder.add_annotation(id, 2, poly)
+    return ann_builder
+
+
+def get_features_from_ann_set(set, coco = None, image_id = 0, category_id = 0):
+    polygons = []
+    for points in set:
+        poly = geometry.Polygon([[p[0], p[1]] for p in points])
+        if(coco is not None):
+            coco.add_annotation(image_id, category_id, poly)
+        polygons.append(poly)
+    return polygons
+
+
+def create_rooms_from_features(features, width, height):
+    room_polygons = geometry.Polygon([(0,0),
+                                    (width,0),
+                                    (width,height),
+                                    (0,height)
+                                    ])
+    for poly in features:
+        room_polygons = room_polygons.difference(poly,3)
+    return geometry.MultiPolygon(room_polygons.geoms[1:]) #this eliminates the exterior from the rooms
+
+
+def apply_preprocessing(source_path, processed_path):
+    img = cv2.imread(source_path)
+    small_img = cv2.resize(img, (0,0), fx=scale_factor, fy=scale_factor)
+    cv2.imwrite(processed_path,small_img)
+    processed_img = preprocess_image(processed_path)
+    #small = cv2.resize(processed, (0,0), fx=scale_factor, fy=scale_factor)
+    cv2.imwrite(processed_path,processed_img)
+    print(get_image_size(source_path))
+    print(get_image_size(processed_path))
+
+
+### SVG element extraction ###
+
+def get_casa_size(doc):
+    path = doc.getElementsByTagName('svg')[0]
+    return int(float(path.getAttribute('width'))), int(float(path.getAttribute('height')))
+
+
+def extract_casa_elements_with_id(id, doc):
+    elements = []
+    for path in doc.getElementsByTagName('g'):
+        #iterates through everything and finds items labelled as walls
+        if(id in path.getAttribute('id')):
+            #luckily, the first attribute after all of these is a polygon containing a list of coordinate points
+            string = path.firstChild.getAttribute('points')
+            points = points_string_to_int_points(string)
+            elements.append(points)
+    return elements
+
+
+### Helper Functions ###
+
+def quadrilateral_to_line(points):
+    base_point = [0,0]
+    points.sort(key=lambda p: check_distance(base_point,p))
+    base_point = points[0]
+    points.sort(key=lambda p: check_distance(base_point,p))
+    return np.array([get_midpoint(points[0], points[1]), get_midpoint(points[2], points[3])])
+
+
+def check_distance(point_A, point_B):
+    return math.sqrt(((point_A[0]-point_B[0]) ** 2) + ((point_A[1]-point_B[1]) ** 2))
+
+
+def get_midpoint(point_A, point_B):
+    return np.array([round((point_A[0]+point_B[0])/2), round((point_A[1]+point_B[1])/2)])
+
+
+def points_string_to_int_points(string):
+    return [[int(round(float(pi)*scale_factor)) for pi in p.split(",")] for p in string.split()]
+
+
+### Validation Images ###
+
+def save_validation_images(filepath):
+    count = 0
+    result = COCO(filepath)
+    for id in random.sample(result.getImgIds(), 15):
+        print("IMAGE "+str(result.imgs[id]))
+        validation_path = str(from_root(dataset_root+"validation_images/casa_"+str(count)))
+        draw_from_coco(id, result, validation_path)
+        count+=1
+
+
+
+main()

rcnn_model/preprocessing/cleaning_images.py

@@ -0,0 +1,67 @@
+import cv2
+import numpy as np
+import os
+
+# Function to preprocess the image
+def preprocessing(image_path):
+    image = cv2.imread(image_path)
+
+    if image is None:
+        print(f"Warning: Could not read {image_path}")
+        return None
+
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    denoisy_img = cv2.GaussianBlur(gray, (5, 5), 0)
+
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    clahe_img = clahe.apply(denoisy_img)
+
+    _, thresholded_img = cv2.threshold(clahe_img, 150, 255, cv2.THRESH_BINARY)
+    
+    edges = cv2.Canny(thresholded_img, 100, 220, apertureSize=3)
+    lines = cv2.HoughLinesP(edges, rho=1, theta=np.pi / 180, threshold=50,
+                            minLineLength=35, maxLineGap=5)
+
+    output_img = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
+
+    if lines is not None:
+        for line in lines:
+            x1, y1, x2, y2 = line[0]
+            cv2.line(output_img, (x1, y1), (x2, y2), (210, 210, 210), 1)
+
+    blended_image = cv2.addWeighted(image, 0.7, output_img, 0.3, 0)
+
+    return blended_image
+
+# Define paths
+source_root = "../cubicasa5k"
+output_dir = "dataset/images"
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+os.chdir(script_dir)
+
+print(f"Fixed Working Directory: {os.getcwd()}")
+
+# Create output directories if they don't exist
+os.makedirs("dataset", exist_ok=True)
+os.makedirs(output_dir, exist_ok=True)
+
+# Iterate over subfolders 1, 2, 3, ..., n
+for subfolder in os.listdir(source_root):
+    subfolder_path = os.path.join(source_root, subfolder)
+
+    if os.path.isdir(subfolder_path):  # Ensure it's a directory
+        image_path = os.path.join(subfolder_path, "F1_original.png")
+
+        if os.path.exists(image_path):
+            processed_img = preprocessing(image_path)
+
+            if processed_img is not None:
+                output_filename = f"{subfolder}.png"  # Save with subfolder name
+                output_path = os.path.join(output_dir, output_filename)
+                cv2.imwrite(output_path, processed_img)
+                print(f"Processed: {image_path} -> {output_path}")
+        else:
+            print(f"Skipping {subfolder}: F1_original.png not found")
+
+print("Processing completed.")

rcnn_model/preprocessing/cleaning_single_image.py

@@ -0,0 +1,48 @@
+# single_image_cleaning.py
+
+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}")
+    
+    image = cv2.imread(image_path)
+    if image is None:
+        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)
+    
+    if lines is not None:
+        for line in lines:
+            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}")
+    return blended_image

rcnn_model/preprocessing/splitting_dataset.py

@@ -0,0 +1,45 @@
+import os
+import random
+import shutil
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+os.chdir(script_dir)
+
+print(f"Fixed Working Directory: {os.getcwd()}")
+
+dataset_path = "dataset"
+image_dir = os.path.join(dataset_path, "images")
+label_dir = os.path.join(dataset_path, "yolo_annotations")
+
+train_dir = "dataset/train"
+val_dir = "dataset/val"
+test_dir = "dataset/test"
+
+for d in [train_dir, val_dir, test_dir]:
+    os.makedirs(os.path.join(d, "images"), exist_ok=True)
+    os.makedirs(os.path.join(d, "yolo_annotations"), exist_ok=True)
+
+images = [f for f in os.listdir(image_dir) if f.endswith(".png")]
+random.shuffle(images)
+
+split_ratio = [0.8, 0.1, 0.1]
+train_split = int(split_ratio[0] * len(images))
+val_split = int(split_ratio[1] * len(images)) + train_split
+
+train_images = images[:train_split]
+val_images = images[train_split:val_split]
+test_images = images[val_split:]
+
+for img in train_images:
+    shutil.move(os.path.join(image_dir, img), os.path.join(train_dir, "images", img))
+    shutil.move(os.path.join(label_dir, img.replace(".png", ".txt")), os.path.join(train_dir, "yolo_annotations", img.replace(".png", ".txt")))
+
+for img in val_images:
+    shutil.move(os.path.join(image_dir, img), os.path.join(val_dir, "images", img))
+    shutil.move(os.path.join(label_dir, img.replace(".png", ".txt")), os.path.join(val_dir, "yolo_annotations", img.replace(".png", ".txt")))
+
+for img in test_images:
+    shutil.move(os.path.join(image_dir, img), os.path.join(test_dir, "images", img))
+    shutil.move(os.path.join(label_dir, img.replace(".png", ".txt")), os.path.join(test_dir, "yolo_annotations", img.replace(".png", ".txt")))
+
+print("Dataset split completed!")

rcnn_model/preprocessing/svg_to_yolo.py

@@ -0,0 +1,120 @@
+import xml.etree.ElementTree as ET
+import re
+import os
+import glob
+
+namespace = {"svg": "http://www.w3.org/2000/svg"}
+
+
+YOLO_CLASSES = {
+    "Door": 0,
+    "Window": 1,
+    "Space": 2
+}
+
+def extract_svg_elements(svg_file):
+    tree = ET.parse(svg_file)
+    root = tree.getroot()
+    
+    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()
+
+        
+        if class_attr in ["Floorplan Floor-1", "Floorplan Floor-2"]:
+
+            for element in floorplan.iter():
+                class_attr = element.get("class")
+                if class_attr:
+                    if any(cat in class_attr for cat in YOLO_CLASSES.keys()):
+                        polygons = []
+                        for poly in element.findall(".//svg:polygon", namespaces=namespace):
+                            points = poly.get("points")
+                            if points:
+                                polygons.append(points)
+
+                        if polygons:
+                            category = next((cat for cat in YOLO_CLASSES if cat in class_attr), None)
+                            print(type(polygons[0]))
+                            if category:
+                                bbox = get_bounding_box(polygons[0], svg_width, svg_height)
+
+                                if "Space" in class_attr:
+                                    name_label = re.sub(r'\b[Ss]pace\b', '', class_attr).strip()
+                                    floorplans["Space"].append({
+                                        "name": name_label, "bbox": bbox
+                                    })
+                                else:
+                                    floorplans[category].append({"bbox": bbox})
+
+    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(","))
+        all_x.append(x)
+        all_y.append(y)
+    print(all_x, all_y)
+    # Bounding Box Calculation
+    x_min, x_max = min(all_x), max(all_x)
+    y_min, y_max = min(all_y), max(all_y)
+
+    # Convert to YOLO format (normalized)
+    x_center = (x_min + x_max) / 2 / svg_width
+    y_center = (y_min + y_max) / 2 / svg_height
+    width = (x_max - x_min) / svg_width
+    height = (y_max - y_min) / 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)
+
+   
+    output_file = f"{output_dir}/{filename}.txt"
+
+    with open(output_file, "w") as f:
+        for category, elements in floorplans.items():
+            class_id = YOLO_CLASSES[category]
+
+            for element in elements:
+                bbox = element["bbox"]
+                yolo_line = f"{class_id} {bbox[0]:.6f} {bbox[1]:.6f} {bbox[2]:.6f} {bbox[3]:.6f}\n"
+                f.write(yolo_line)
+
+    print(f"YOLO annotations saved in '{output_dir}'")
+
+
+input_folder = "../cubicasa5k/high_quality/"  
+output_folder = "dataset/yolo_annotations" 
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+os.chdir(script_dir)
+
+print(f"Fixed Working Directory: {os.getcwd()}")
+
+subfolders = glob.glob(os.path.join(input_folder, "*"))  
+
+for subfolder in subfolders:
+    svg_file = os.path.join(subfolder, "model.svg")  
+    
+    if os.path.exists(svg_file):  
+        filename = os.path.basename(subfolder)  
+        print(f"Processing: {svg_file} ...")
+        
+        floorplans, svg_width, svg_height = extract_svg_elements(svg_file)
+        save_yolo_annotations(floorplans, output_folder, filename)
+
+print(" All SVG files have been processed!")

rcnn_model/scripts/rcnn_config.py

@@ -0,0 +1,38 @@
+import os
+from detectron2 import model_zoo
+from detectron2.config import get_cfg
+from detectron2.data.datasets import register_coco_instances
+from from_root import from_root
+
+
+def write_config():
+    cfg = get_cfg()
+    cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml"))
+    cfg.DATASETS.TRAIN = ("inodata_train","cubicasa_train")
+    cfg.DATASETS.PROPOSAL_FILES_TRAIN = ("inodata_train")
+    cfg.DATASETS.TEST = ()
+    cfg.SOLVER.BASE_LR = .0005
+    cfg.SOLVER.MAX_ITER = 100
+    cfg.SOLVER.CHECKPOINT_PERIOD = 1000
+    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256
+    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4
+    cfg.INPUT.MASK_FORMAT = "polygon"
+    cfg.MODEL.RPN.NMS_THRESH = 0.8
+    cfg.SOLVER.NUM_DECAYS = 2
+    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml")
+    cfg.SOLVER.STEPS = (50,75)
+    cfg.MODEL.ROI_HEADS.POSITIVE_FRACTION = .7
+    cfg.SOLVER.GAMMA = 0.4
+    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
+
+os.chdir(str(from_root()))
+register_coco_instances("cubicasa_train",{},"dataset/annotations/cubicasa_train.json","dataset/")
+register_coco_instances("inodata_train",{},"dataset/annotations/train_sampled_data.json","dataset/")
+register_coco_instances("inodata_val",{},"dataset/annotations/val_sampled_data.json","dataset/")
+register_coco_instances("cubicasa_val",{},"dataset/annotations/cubicasa_test.json","dataset/")

rcnn_model/scripts/rcnn_eval.py

@@ -0,0 +1,108 @@
+
+import sys
+sys.path.append("/workspaces/tensorflow-gpu/cocoapi/PythonAPI/pycocotools")
+import os
+import cv2
+from detectron2.data import DatasetCatalog
+import detectron2.data as ddata
+from detectron2.engine import DefaultPredictor
+from detectron2.evaluation import COCOEvaluator
+from detectron2.utils.visualizer import ColorMode
+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
+
+def main(cfg,results_filename = "eval_results.txt"):
+    #update config file
+    cfg.DATALOADER.NUM_WORKERS = 1
+    cfg.SOLVER.IMS_PER_BATCH = 1                                           
+    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
+
+    #run evaluation
+    results = standard_evaluation(cfg)
+
+    #save results
+    file = open(results_directory+results_filename,"w")
+    file.write(str(results))
+    file.close()
+
+
+### Evaluation ###
+
+def standard_evaluation(cfg):
+    #load predictor
+    predictor = DefaultPredictor(cfg)
+    test_data_loader = ddata.build_detection_test_loader(cfg, "inodata_val")
+
+    #save some validation images
+    save_validation_images(predictor)
+
+    #create evaluator
+    evaluator = COCOEvaluator("inodata_val",tasks={"segm","bbox"},output_dir="./eval_output",distributed=False,max_dets_per_image=50,allow_cached_coco=False)
+    print("EVALUATING")
+    evaluator.reset()
+
+    #load results into evaluator
+    for inputs, outputs in block_prediction(test_data_loader, predictor):
+        evaluator.process(inputs,outputs)
+        del inputs
+        del outputs
+        print("|",end="")
+        time.sleep(.5)
+    print("")
+
+    #run evaluator
+    results = evaluator.evaluate()
+    print(results)
+    print("EVALUATED")
+    return results
+
+
+def block_prediction(loader, predictor):
+    for data in loader:
+        if(check_image_size_thresh(data[0]["file_name"],max_image_size)):
+            image = cv2.imread(data[0]["file_name"])
+            result = predictor(image)
+            yield data, [result]
+            del image
+            del result
+
+
+### Validation Images ###
+
+def save_validation_images(predictor):
+    val_img_id = 1
+    for d in random.sample(DatasetCatalog.get("inodata_val"), 16):
+        save_image(d,predictor,val_img_id)
+        val_img_id += 1
+
+
+def save_image(d, predictor, val_img_id):
+    try:
+        if(check_image_size_thresh(d["file_name"],max_image_size)):
+            #set load image
+            val_img_dest_path = "models/rcnn/validation_images/RCNN_val_image_"+str(val_img_id)+".png"
+            im = cv2.imread(d["file_name"])
+            outputs = predictor(im)
+
+            #save image
+            v = Visualizer(im[:,:,::-1],scale=0.5,instance_mode=ColorMode.IMAGE_BW)
+            out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+            plt.imshow(out.get_image()[:,:,::-1])
+            plt.axis('off')
+            plt.savefig(val_img_dest_path,bbox_inches='tight',pad_inches=0)
+            print("Saved validation image to "+val_img_dest_path)
+            plt.clf()
+    except:
+        print("ERROR SAVING IMAGE")

rcnn_model/scripts/rcnn_full_tuner.py

@@ -0,0 +1,19 @@
+from rcnn_config import write_config
+import rcnn_train
+import rcnn_eval
+import rcnn_run
+
+run_iter_counts = [3000,4000,5000,6000,7000,8000,9000,10000]
+
+def main():
+    cfg = write_config()
+    for i in range(0,len(run_iter_counts)):
+        iters = run_iter_counts[i]
+        if(iters > 0):
+            cfg.SOLVER.MAX_ITER = iters
+            cfg.SOLVER.STEPS = (int(iters*.5),int(iters*.75))
+            rcnn_train.main(cfg)
+            rcnn_run.default_sample(cfg, i)
+            rcnn_eval.main(cfg,"eval_results_run_"+str(i)+".txt")
+
+main()

rcnn_model/scripts/rcnn_run.py

@@ -0,0 +1,173 @@
+import os
+import cv2
+from pycocotools.coco import COCO
+from detectron2.engine import DefaultPredictor
+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 = "rcnn_model/results/"
+sample_data_directory = "rcnn_model/sample/"
+
+def main(cfg,img_source_path, coco_dest_filename, val_img_dest_filename):
+    os.chdir(str(from_root()))
+
+    #configure model
+    cfg.DATALOADER.NUM_WORKERS = 1
+    cfg.SOLVER.IMS_PER_BATCH = 1                                           
+    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "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)                                                  
+
+
+### Main Runner ###
+
+def prediction_runner(filename, coco_dest_path, val_img_dest_path, predictor, segmented_prediction = False, scale_factor = 0):
+    #set up annotation builder
+    ann_builder = instantiate_ann_build()
+
+    #set up image
+    initImg = preprocess_image(filename)
+    init_width = initImg.shape[1]
+    init_height = initImg.shape[0]
+    img_id = ann_builder.add_image(str(from_root(filename)), init_width, init_height)
+    print("cleaned")
+
+    #resize image
+    scaled_width = 800
+    scaled_height = 800
+    img = cv2.resize(initImg, (scaled_width,scaled_height))
+    if(scale_factor > 0):
+        img = cv2.resize(initImg, (0,0), fx=scale_factor, fy=scale_factor)
+        scaled_width = img.shape[1]
+        scaled_height = img.shape[0]
+    print("resized")
+
+    #run prediction
+    if(segmented_prediction):
+        run_segmented_prediction(ann_builder, predictor, img, img_id, scaled_width, scaled_height, scale_factor)
+    else:
+        run_prediction(ann_builder, predictor, img, img_id, init_width/scaled_width, init_height/scaled_height)
+
+    #save the file
+    ann_builder.save_file(str(from_root(coco_dest_path)))
+
+    #visualize
+    coco = COCO(str(from_root(coco_dest_path)))
+    draw_from_coco(0, coco, val_img_dest_path)
+
+
+def instantiate_ann_build():
+    ann_builder = AnnBuild()
+    ann_builder.set_info("generated annotations of Inovonics and university provided data","inovonData","NA",datetime.now())
+    ann_builder.add_license("TODO", "TODO")
+    return ann_builder
+
+
+def prediction_outputs_to_annotations(annotations, outputs, img_id, base_ann_id, x_offset=0, y_offset=0, h_scale_factor=1, v_scale_factor=1):
+    ann_id = base_ann_id
+    for i in range(0,len(outputs["instances"].to(torch.device("cpu")).pred_masks)):
+        mask = outputs["instances"].to(torch.device("cpu")).pred_masks[i]
+        class_id = outputs["instances"].to(torch.device("cpu")).pred_classes[i].item()
+        score = outputs["instances"].to(torch.device("cpu")).scores[i].item()
+        annotations.append(bitmask_to_polygon(ann_id, img_id, class_id, score, mask, x_offset=x_offset, y_offset=y_offset,scale_factor_width = h_scale_factor,scale_factor_height = v_scale_factor))
+        ann_id += 1
+    return annotations, ann_id
+
+
+### Standard Prediction ###
+
+def run_prediction(ann_builder, predictor, img, img_id, h_scale_factor, v_scale_factor):
+    outputs = predictor(img)
+    print("predicted")
+    annotations = []
+    annotations, annId = prediction_outputs_to_annotations(annotations, outputs, img_id, 0, 0, 0, h_scale_factor, v_scale_factor)
+    ann_builder.annotations = annotations
+    print("annotations converted")
+
+
+### Segmented Prediction ###
+
+def run_segmented_prediction(ann_builder, predictor, img, img_id, width, height, scale_factor, segment_size = 800):
+    #initialize
+    count = 1
+    subimg_dest_path = "models/rcnn/westmoor_check/subimgs/subimg_"
+    annotations = []
+    ann_id = 0
+
+    #iterate through segments
+    for xi in range(0, int(width/segment_size)+1):
+        for yi in range(0, int(height/segment_size)+1):
+            #calculate subimg area
+            h_base, v_base, h_boundary, v_boundary = get_subimg_area(xi,yi,width,height,segment_size)
+
+            #save subimgs
+            if(h_boundary > h_base and v_boundary > v_base):
+                subimg = img[v_base:v_boundary,h_base:h_boundary,:]
+                save_subimg(subimg, subimg_dest_path, count, h_boundary-h_base, v_boundary-v_base)
+
+                #get annotations
+                outputs = predictor(subimg)
+                annotations, ann_id = prediction_outputs_to_annotations(annotations, outputs, img_id, ann_id, h_base*1/scale_factor, h_base*1/scale_factor, 1/scale_factor, 1/scale_factor)
+            count += 1
+    ann_builder.annotations = annotations
+
+
+def get_subimg_area(xi,yi,img_width,img_height,segment_size):
+    h_base = xi*segment_size
+    v_base = yi*segment_size
+    h_boundary = (xi+1)*segment_size
+    if(h_boundary >= img_width):
+            h_boundary = img_width-1
+    v_boundary = (yi+1)*segment_size
+    if(v_boundary >= img_height):
+            v_boundary = img_height-1
+    return h_base, v_base, h_boundary, v_boundary
+
+
+def save_subimg(subimg, subimg_dest_path, count, width, height):
+    plt.figure(figsize=(width, height),dpi=1)
+    plt.imshow(subimg)
+    plt.axis('off')
+    plt.savefig(str(from_root(subimg_dest_path+str(count)+".png")),bbox_inches='tight',pad_inches=0)
+    plt.clf()
+
+
+def default_sample(cfg, run_index = -1):
+    if(run_index >= 0):
+        main(cfg,sample_data_directory+"westmoor_floor_2_floorplan.png","westmoor_result_run_"+str(run_index)+".json","westmoor_result_run_"+str(run_index)+".png")
+        main(cfg,sample_data_directory+"REV2 ENSURE Layout - Springs at the Waterfront - PR-2023-3661 - CN0012630 - 11.02.2023 (1)-16.png","rev2_result_run_"+str(run_index)+".json","rev2_result_run_"+str(run_index)+".png")
+        main(cfg,sample_data_directory+"F1_original.png","cubicasa_result_run_"+str(run_index)+".json","cubicasa_result_run_"+str(run_index)+".png")
+    else:
+        main(cfg,sample_data_directory+"westmoor_floor_2_floorplan.png","westmoor_result.json","westmoor_result.png")
+        main(cfg,sample_data_directory+"REV2 ENSURE Layout - Springs at the Waterfront - PR-2023-3661 - CN0012630 - 11.02.2023 (1)-16.png","rev2_result.json","rev2_result.png")
+        main(cfg,sample_data_directory+"F1_original.png","cubicasa_result.json","cubicasa_result.png")
+
+if __name__ == "__main__":
+    print("Inside Main Function Call")
+    cfg = write_config()
+    run_index=-1
+    main(cfg,sample_data_directory+"F1_original.png","cubicasa_result_run_"+str(run_index)+".json","cubicasa_result_run_"+str(run_index)+".png")

rcnn_model/scripts/rcnn_train.py

@@ -0,0 +1,21 @@
+import os
+from detectron2.modeling import build_model
+from detectron2.data.datasets import register_coco_instances
+from detectron2.engine import DefaultTrainer
+from from_root import from_root
+
+def main(cfg):
+    os.chdir(str(from_root()))
+    register_coco_instances("cubicasa_train",{},"dataset/annotations/cubicasa_train.json","dataset/")
+    register_coco_instances("cubicasa_val",{},"dataset/annotations/cubicasa_test.json","dataset/")
+
+    
+    cfg.DATALOADER.NUM_WORKERS = 4
+    cfg.SOLVER.IMS_PER_BATCH = 2
+
+    model = build_model(cfg)
+
+    model.train()
+    trainer = DefaultTrainer(cfg=cfg)
+    trainer.resume_or_load(resume=False)
+    trainer.train()

rcnn_model/utils/coco_to_inovonics_json.py

@@ -0,0 +1,34 @@
+
+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)
+    if(len(img_ids) == 1):
+        for img_id in img_ids:
+            coco_img_to_inovonics_json(coco, inovonics_anns_dest_path, img_id)
+    else:
+        for img_id in img_ids:
+            coco_img_to_inovonics_json(coco, inovonics_anns_dest_path[0:-5]+"_"+str(img_id)+".json", img_id)
+
+    
+def coco_img_to_inovonics_json(coco, inovonics_anns_dest_path, img_id=0):
+    #iterate
+    annotation_full_file = []
+    count = 0
+    for ann in coco.imgToAnns[img_id]:
+        print(ann)
+        inov_ann_build = InovAnnBuild()
+        inov_ann_build.set_id(str(count))
+        inov_ann_build.set_body("Room "+str(count))
+        inov_ann_build.set_type("Selection")
+        inov_ann_build.set_target("FragmentSelector","http://www.w3.org/TR/media-frags/",ann["bbox"])
+        annotation_full_file.append(inov_ann_build.final_output())
+        count+=1
+
+    #save file
+    coco_file = open(inovonics_anns_dest_path,'w')
+    json.dump(annotation_full_file,coco_file,indent=4)
+    coco_file.close()

rcnn_model/utils/floorplan_vectorizer_utils.py

@@ -0,0 +1,104 @@
+from PIL import Image
+import matplotlib.pyplot as plt
+import skimage.io as io
+import numpy as np
+from pycocotools import mask
+from skimage import measure
+from shapely import geometry
+from from_root import from_root
+
+### Image Size Checking ###
+
+def check_image_size_thresh(png_path, areathreshold):
+        width, height = get_image_size(png_path)
+        return width*height <= areathreshold
+
+def get_image_size(png_path):
+    img = Image.open(png_path)
+    return img.width, img.height
+
+
+### Visualization ###
+
+def draw_from_coco(id,coco,annotated_img_dest_path,category_filter = [0,1,2,3], blank_bg = False):
+    filename = coco.imgs[id]["file_name"]
+    image = io.imread(str(from_root(filename)))
+    if(image is not None):
+        plt.figure(figsize=(image.shape[1],image.shape[0]),dpi=1)
+        if(blank_bg):
+            image = get_blank_image(image.shape[0],image.shape[1])
+        plt.imshow(image)
+        plt.axis('off')
+        annotation_ids = coco.getAnnIds(imgIds=[id], catIds=category_filter, iscrowd=None)
+        annotations = coco.loadAnns(annotation_ids)
+        coco.showAnns(annotations)
+        plt.savefig(annotated_img_dest_path,bbox_inches='tight',pad_inches=0)
+        plt.clf()
+        print("---")
+        print("Saved Validation Image "+annotated_img_dest_path)
+
+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,:]
+
+
+### Converting Bitmask to Polygon ###
+
+#modifeid version of code from Waspinator on https://github.com/cocodataset/cocoapi/issues/131
+def bitmask_to_polygon(id, im_id, cat_id, score, ground_truth_binary_mask, x_offset = 0, y_offset = 0, scale_factor_width = 1, scale_factor_height = 1):
+    fortran_ground_truth_binary_mask = np.asfortranarray(ground_truth_binary_mask)
+    encoded_ground_truth = mask.encode(fortran_ground_truth_binary_mask)
+    ground_truth_area = mask.area(encoded_ground_truth)
+    ground_truth_bounding_box = mask.toBbox(encoded_ground_truth)
+    contours = measure.find_contours(ground_truth_binary_mask.numpy(), 0.5)
+
+    bbox = []
+    for i in range(0,len(ground_truth_bounding_box.tolist())):
+        if(i%2 == 0):
+            bbox.append(scale_factor_width*ground_truth_bounding_box.tolist()[i]+x_offset)
+        else:
+            bbox.append(scale_factor_height*ground_truth_bounding_box.tolist()[i]+y_offset)
+
+    polygon = []
+    for contour in contours:
+        contour = np.flip(contour, axis=1)
+        segmentation = contour.ravel().tolist()
+        for i in range(0,len(segmentation)):
+            if(i%2 == 0):
+                segmentation[i] = int(segmentation[i]*scale_factor_width)
+            else:
+                segmentation[i] = int(segmentation[i]*scale_factor_height)
+        polygon.append(segmentation)
+
+    segmentations = []
+    if(len(polygon) >  0):
+        smoothed_polygons = polygon_smoothing_and_offset(polygon[0], x_offset, y_offset)#toPolygon[0]#
+        for segmentation in smoothed_polygons:
+            segmentations.append(segmentation)
+    
+    annotation = {
+            "segmentation": segmentations,
+            "area": ground_truth_area.tolist(),
+            "iscrowd": 0,
+            "image_id": im_id,
+            "bbox": bbox,
+            "category_id": cat_id,
+            "id": id,
+            "score": score
+        }
+    
+    return annotation
+
+def polygon_smoothing_and_offset(polygon, x_offset, y_offset):
+    points = []
+    for i in range(0,int(len(polygon)/2)):
+        points.append([polygon[(2*i)]+x_offset,polygon[(2*i)+1]+y_offset])
+    if(len(points) < 4):
+        return []
+    poly = geometry.Polygon(points)
+    for i in [1,2,3,5,8,12,15,18]:
+        poly = poly.simplify(i)
+    return [np.array(poly.exterior.coords).astype(int).ravel().tolist()[:-2]]
+
+

rcnn_model/utils/inovonics_ann_builder.py

@@ -0,0 +1,52 @@
+import json
+import numpy as np
+from datetime import datetime
+import random
+
+version_number = "0.0.1"
+
+class InovonicsAnnotationBuilder:
+    #creates the base structure of the coco format
+    def __init__(self):
+        self.annotation_id = ""
+        self.body = []
+        self.type = ""
+        self.target = {}
+
+    def set_id(self, id):
+        self.annotation_id = id
+
+    def set_body(self, room_name, locator_value=None):
+        color = self.generate_color()
+        self.body = [{"type":"TextualBody",
+                     "value":room_name},
+                     {"type":"highlighting",
+                     "value":color},
+                     {"type":"locators",
+                     "value":locator_value}]
+        
+    def generate_color(self):
+        red = str(random.randint(0,255))
+        green = str(random.randint(0,255))
+        blue = str(random.randint(0,255))
+        return "rgb("+red+","+green+","+blue+")"
+
+    def set_type(self, type):
+        self.type = type
+
+    def set_target(self, type, url, bbox):
+        rect = self.bbox_to_rect(bbox)
+        self.target = {"selector":{"type":type,
+                                 "conformsTo":url,
+                                 "value":rect}}
+    
+    def bbox_to_rect(self,bbox):
+        return "xywh=pixel"+str(bbox[0])+","+str(bbox[1])+","+str(bbox[2])+","+str(bbox[3])
+    
+    def final_output(self):
+        return {"annotation_id":self.annotation_id, "body":self.body, "type":self.type, "target":self.target}
+    
+    def save_file(self, filepath):
+        coco_file = open(filepath,'w')
+        json.dump(self.final_output(),coco_file,indent=4)
+        coco_file.close()

requirements.txt

@@ -0,0 +1,5 @@
+streamlit>=1.30
+opencv-python
+pycocotools
+Pillow
+gdown