Enhancing Floor Plan Recognition: A Hybrid Mix-Transformer and U-Net Approach for Precise Wall Segmentation

TL;DR

This paper introduces MitUNet, a hybrid neural network combining Mix-Transformer and U-Net with attention mechanisms, achieving high-precision wall segmentation for 3D indoor space reconstruction.

Contribution

The novel MitUNet architecture effectively combines global context and fine details, improving wall segmentation accuracy over existing methods.

Findings

MitUNet outperforms standard models on CubiCasa5k and regional datasets.

It achieves high boundary accuracy and structural correctness in wall masks.

The approach balances precision and recall using Tversky loss.

Abstract

Automatic 3D reconstruction of indoor spaces from 2D floor plans necessitates high-precision semantic segmentation of structural elements, particularly walls. However, existing methods often struggle with detecting thin structures and maintaining geometric precision. To address this, we introduce MitUNet, a hybrid neural network designed to bridge the gap between global semantic context and fine-grained structural details. Our architecture combines a Mix-Transformer encoder with a U-Net decoder enhanced with spatial and channel attention blocks. Optimized with the Tversky loss function, this approach achieves a balance between precision and recall, ensuring accurate boundary recovery. Experiments on the CubiCasa5k dataset and the regional dataset demonstrate MitUNet's superiority in generating structurally correct masks with high boundary accuracy, outperforming standard models. This tool provides a robust foundation for automated 3D reconstruction pipelines. To ensure reproducibility and facilitate future research, the source code and the regional dataset are publicly available at https://github.com/aliasstudio/mitunet and https://doi.org/10.5281/zenodo.17871079, respectively.