CubiCasa5K: A Dataset and an Improved Multi-Task Model for Floorplan Image Analysis

Paper: CubiCasa5K: A Dataset and an Improved Multi-Task Model for Floorplan Image Analysis

Multi-Task Model

The model uses the neural network architecture presented in Raster-to-Vector: Revisiting Floorplan Transformation [1]. The pre- and post-processing parts are modified to suit our dataset, but otherwise the pipeline follows the torch implementation of [1] as much as possible. Our model utilizes the multi-task uncertainty loss function presented in Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics. An example of our trained model's prediction can be found in the samples.ipynb file.

Dataset

CubiCasa5K is a large-scale floorplan image dataset containing 5000 samples annotated into over 80 floorplan object categories. The dataset annotations are performed in a dense and versatile manner by using polygons for separating the different objects. You can download the CubiCasa5K dataset from here and extract the zip file to data/ folder.

Requirements

The model is written for Python 3.11 and Pytorch 2.4 with CUDA enabled GPU. See the pyproject.toml for dependencies (NB: support is limited to linux-64 and osx-arm64). If you want to use the Dockerfile you need to have docker and nvidia-docker2 installed. We use pre-built image anibali/pytorch:cuda-9.0 as a starting point and install other required libraries using pip. To create the container run in the:

docker build -t cubi -f Dockerfile .

To start JupyterLab in the container:

docker run --rm -it --init \
  --runtime=nvidia \
  --ipc=host \
  --publish 1111:1111 \
  --user="$(id -u):$(id -g)" \
  --volume=$PWD:/app \
  -e NVIDIA_VISIBLE_DEVICES=0 \
  cubi jupyter-lab --port 1111 --ip 0.0.0.0 --no-browser

You can now open a terminal in JupyterLab web interface to execute more commands in the container.

Installation steps

1. Get Pixi installed via (you might need a shell restart afterwards):

curl -fsSL https://pixi.sh/install.sh | bash

2. Install the dev environment for model evaluation, or the cuda environment for both training and evaluation:

pixi install -e default

Database creation

We create a LMDB database of the dataset, where we store the floorplan image, segmentation tensors and heatmap coordinates. This way we can access the data faster during training and evaluation. The downside however is that the database takes about 105G of hard drive space. There is an option to parse the SVG file on the go but it is slow for training. Commands to create the database:

python create_lmdb.py --txt val.txt
python create_lmdb.py --txt test.txt
python create_lmdb.py --txt train.txt

Train

python train.py

Different training options can be found in the script file. Tensorboard is not included in the docker container. You need to run it outside and point it to cubi_runs/ folder. For each run a new folder is created with a timestamp as the folder name.

tensorboard --logdir runs_cubi/

Evaluation

Our model weights file can be downloaded here. Once the weights file is in the project folder evaluation can be done. Also you can run the jupyter notebook file to see how the model is performing for different floorplans.

python eval.py --weights model_best_val_loss_var.pkl

Additional option for evaluation can be found in the script file. The results can be found in runs_cubi/ folder.

Todo

• Modify create_lmdb.py to save files as uint8 (now using float32 which is the main reason why the lmdb file gets as big as over 100 gbytes).
• Modify augmentations.py to operate with numpy arrays (the reason why it currently utilizes torch tensors is the fact that in our earlier version we applied augmentations to heatmap tensors and not to heatmap dicts which is the correct way to do it)