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

AI summary

Overview

Research area: Computer vision, specifically semantic segmentation of structural elements in 2D architectural floor plans as a prerequisite for automated 3D reconstruction.

Technical level: Advanced. The paper assumes familiarity with encoder-decoder segmentation networks, Transformer backbones, skip connections, and region-based loss functions.

Scope: The paper proposes and evaluates MitUNet, a hybrid Mix-Transformer encoder plus U-Net decoder architecture trained with an asymmetric Tversky loss, for precise wall segmentation on the CubiCasa5k dataset and a newly released 500-image regional floor plan dataset.

What This Paper Is About

Automatically converting 2D raster floor plans into 3D building models is expensive to do by hand, and existing segmentation models struggle with thin wall structures and geometric precision. This paper's goal is to segment walls accurately enough to serve as a reliable structural basis for vectorization and Scan-to-BIM pipelines, particularly on regional floor plans whose drawing conventions differ from standard open datasets.

Key Contributions

1. MitUNet architecture: A hybrid network that pairs a hierarchical Mix-Transformer encoder (the MiT-b4 backbone from SegFormer, pre-trained on ImageNet, with channel dimensions {64, 128, 320, 512}) with a full U-Net decoder path (decoder channels {256, 128, 64, 32} from deepest to shallowest) enhanced by scSE attention blocks after each convolutional layer. The authors replace SegFormer's lightweight MLP decoder, which they argue loses high-frequency edge information through bilinear upsampling.

2. Loss function analysis for wall segmentation: A comparative study of Dice, Focal, Lovasz-Softmax, and Tversky losses, including a hyperparameter ablation on the Tversky parameters, concluding that asymmetric penalties (α = 0.6, β = 0.4) suppress false-positive noise and produce cleaner, vectorization-ready masks.

3. Annotation refinement procedure: A preprocessing pipeline that generates separate binary masks for doors and windows, structurally dilates them by a fixed margin of approximately 30 pixels, subtracts them from the wall mask to carve out precise openings, and applies morphological closing with a 5 × 5 kernel to ensure wall continuity.

4. Public release of code and data: Source code at https://github.com/aliasstudio/mitunet and the regional dataset (Floor Plan CIS) at https://doi.org/10.5281/zenodo.17871079.

Main Findings

• MitUNet with Tversky loss leads Table 1: On the Regional Dataset, MitUNet (mit_b4, Tversky) achieved the highest mIoU (87.91%), Precision (94.82%), Accuracy (98.86%), and Boundary IoU (85.01%) among the models compared, with Recall of 92.35% and VRAM usage of 1751 MiB. Of the other three MitUNet loss variants, Dice reached 87.61% mIoU / 84.89% B-IoU, Lovasz 87.31% / 84.48%, and Focal 86.94% / 84.15%.

• Comparable accuracy at higher memory cost for CNN baselines: UNet++ (ResNet50, Lovasz) achieved 87.25% mIoU and 84.68% B-IoU but used 3311 MiB of VRAM, versus 1751 MiB for MitUNet. DeepLabV3+ (ResNet50, Lovasz) reached 85.16% mIoU and 81.59% B-IoU at 947 MiB, and UPerNet (mit_b4, Dice) reached 86.67% mIoU and 82.91% B-IoU at 2211 MiB.

• The U-Net decoder matters under asymmetric loss: SegFormer with the same Tversky configuration (α = 0.6, β = 0.4) achieved the lowest B-IoU among its variants at 81.31%, with Recall 90.85% and Precision 93.44%. The authors attribute this to the MLP decoder's bilinear upsampling being unable to realize thin boundaries when false positives are penalized strictly.

• Tversky parameter trade-off is monotonic: Increasing α in the ablation raised Precision while lowering Recall, mIoU and B-IoU: (0.6/0.4) gave 92.35% Recall, 94.82% Precision, 87.91% mIoU, 85.01% B-IoU; (0.7/0.3) gave 89.20/96.25/86.20/82.75; (0.8/0.2) gave 87.13/97.14/84.95/81.38; (0.9/0.1) gave 81.41/98.23/80.23/75.40.

• Fine-tuning on CubiCasa5k pre-trained weights improved all metrics: In Table 3, Lovasz reached 88.75% mIoU and 85.93% B-IoU, Dice reached 88.67% mIoU and 85.98% B-IoU, and Tversky (0.6/0.4) reached 88.53% mIoU with the highest Precision of 94.99% (compared to approximately 93.6% for Dice and Lovasz). The authors state this Precision advantage is deliberate because symmetric losses can dilate boundaries to maximize overlap.

• Qualitative claims: The authors report that Tversky loss leaves ventilation shafts correctly hollow where Dice and Focal hallucinate solid blocks and internal noise, and that MitUNet minimizes "staircase" edge artifacts relative to standard losses, reducing the need for smoothing that distorts corner geometry.

• Regional dataset differences are documented: Compared to CubiCasa5k, the authors describe three domain shifts: texture-based material encoding (solid black load-bearing walls versus hatching for partitions), non-Manhattan geometry such as curved or angled outer walls, and high-density semantic clutter including furniture outlines, dimension lines, and text overlays.

Methodology in Plain English

The authors built a segmentation model that tries to get both the "big picture" and the fine details right. The encoder is a Mix-Transformer, which processes an image in stages and produces feature maps at four resolutions, letting the model reason about global layout without losing locality (the encoder uses overlapping patch merging to avoid grid artifacts). The decoder is a full U-Net reconstruction path that progressively upsamples those features and fuses them with high-resolution encoder features through skip connections, with scSE blocks that recalibrate features spatially and by channel after each convolution.

Training is done in two stages. First, models are pre-trained on CubiCasa5k (5,000 images) to learn general geometry. Then they are fine-tuned on the authors' regional dataset of 500 floor plans collected from real estate listings in the Russian Federation and CIS region, using a reduced learning rate of 1e-5. Inputs are standardized to 512 × 512 pixels with a batch size of 4, normalized with ImageNet mean [0.485, 0.456, 0.406] and standard deviation [0.229, 0.224, 0.225]. Augmentation includes random affine scaling between [0.9, 1.1], rotation between [-15°, 15°], translation at p = 0.7, elastic and grid distortions at p = 0.2 to mimic paper warping, brightness/contrast at p = 0.5, CLAHE at p = 0.2, and Gaussian and ISO noise.

The optimization uses Adam, 30 epochs per stage, batch size 4, and a ReduceLROnPlateau scheduler with factor 0.5 and patience 3 monitoring validation IoU. Each experiment was repeated three times, and the reported metrics are averages for the best validation IoU. The data were aggregated into a single pool and split 80/20 for training and validation with a fixed random seed of 42. All experiments ran on a single NVIDIA RTX 4060 Ti with 16 GB VRAM using PyTorch and Albumentations.

The loss comparison shows why asymmetric penalties were chosen: Dice optimizes the F1 score but can encourage dilated masks, Focal down-weights easy negatives but lacks boundary control, Lovasz-Softmax directly optimizes IoU but gives no explicit control over precision versus recall, and Tversky adds separate α and β weights for false positives and false negatives.

Why This Matters

Impact on research: The paper argues that wall segmentation errors propagate into downstream 3D generation as topological inconsistencies or distorted geometry, so improving boundary fidelity at this stage matters more than aggregate pixel overlap. By reporting Boundary IoU alongside mIoU and by releasing both code and a regional dataset, the work also supports domain-adaptation research on floor plans whose drawing conventions differ from Western open datasets. The reported comparison shows that a Transformer-CNN hybrid can match or exceed CNN baselines while using substantially less memory than UNet++ (1751 MiB versus 3311 MiB).

Real-world applications:

• Automated 3D reconstruction and visualization of indoor spaces from 2D raster floor plans.
• Real estate marketing, where fast 3D visualization of listed properties has direct commercial value.
• Project planning and space visualization for private individuals.
• Vectorization and Scan-to-BIM pipelines, where sharp, low-noise wall masks reduce the post-processing smoothing that distorts corners.

Industry relevance: The authors note that conventional architectural tools such as Autodesk Revit, AutoCAD, SketchUp, and ArchiCAD target manual design and BIM workflows rather than automatic reconstruction from raw 2D raster plans, requiring professional expertise and user intervention. This leaves a gap for high-throughput automation, and the paper positions MitUNet as a foundation for that.

Future Directions

• Extending the segmentation module into an end-to-end vectorization framework that directly generates topological graphs from raster images, which the authors name as their future work.
• Studying domain adaptation more systematically now that the regional dataset is public, particularly for texture-based material encoding, non-Manhattan geometries, and dense semantic clutter.
• Revisiting the loss trade-off: on the fine-tuned model, Lovasz and Dice posted marginally higher mIoU and B-IoU than Tversky, so the case for asymmetric penalties rests on a Precision advantage and qualitative crispness rather than on every reported metric.
• Clarifying whether MitUNet generalizes to drawing conventions beyond the two datasets tested; the paper does not report cross-domain evaluation of the CubiCasa5k-pretrained model without fine-tuning, nor inference latency or parameter counts.

Target Audience

Researchers and practitioners working on floor plan analysis, indoor scene understanding, semantic segmentation architectures, and Scan-to-BIM or automated 3D reconstruction pipelines. It is most useful to readers already comfortable with encoder-decoder segmentation networks and loss-function design, and to engineers who need a reproducible baseline with released code and data for regional floor plan styles. Readers looking for deployment-oriented metrics such as inference speed, model size, or training time will not find them reported here.

Authors’ 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.