giswqs/whu-building-unetplusplus-efficientnet-b4

WHU Building Detection — EfficientNet-B4 + UNet++

A semantic segmentation model for building detection in high-resolution aerial imagery, trained on the WHU Building Dataset.

Model Description

Property	Value
Architecture	UNet++
Encoder	EfficientNet-B4 (ImageNet pretrained)
Framework	segmentation-models-pytorch (SMP)
Training Framework	PyTorch Lightning
Input	3-channel RGB, 512x512 tiles
Output	2-class mask (Background=0, Building=1)
Parameters	~20.8M
Model Size	~84 MB

Performance

Evaluated on the WHU Building Dataset test split (1,228 tiles):

Metric	Score
IoU	0.9054
Dice	0.9503
Best Val IoU	0.9434

Training Details

• Dataset: WHU Building Dataset — 5,732 training tiles (512x512 RGB at 0.3m resolution)
• Validation split: 20% of training data
• Optimizer: AdamW (lr=1e-4, weight_decay=1e-4)
• Loss: CrossEntropyLoss
• Epochs: 36 (early stopping, patience=10)
• Batch size: 16
• GPU: NVIDIA RTX 6000 Ada (48GB)
• Encoder weights: ImageNet pretrained

Quick Start

Installation

pip install geoai-py timm segmentation-models-pytorch

Inference with GeoAI

import geoai

# Run building detection on a GeoTIFF
geoai.timm_segmentation_from_hub(
    input_path="input_image.tif",
    output_path="building_prediction.tif",
    repo_id="giswqs/whu-building-unetplusplus-efficientnet-b4",
    window_size=512,
    overlap=256,
    batch_size=4,
)

# Vectorize to building footprints
gdf = geoai.orthogonalize(
    input_path="building_prediction.tif",
    output_path="building_footprints.geojson",
    epsilon=2.0,
)

Manual Loading

import json
import torch
import segmentation_models_pytorch as smp

# Load config
with open("config.json") as f:
    config = json.load(f)

# Create model
model = smp.UnetPlusPlus(
    encoder_name="efficientnet-b4",
    encoder_weights=None,
    in_channels=3,
    classes=2,
)

# Load weights
state_dict = torch.load("model.pth", map_location="cpu")
model.load_state_dict(state_dict)
model.eval()

Example Notebook

See the full inference notebook with visualization and analysis:

Dataset

The WHU Building Dataset consists of aerial imagery at 0.3m resolution with binary building masks:

• Train: 5,732 tiles (512x512 RGB)
• Val: 1,228 tiles
• Test: 1,228 tiles

Reference

Ji, S., Wei, S., & Lu, M. (2019). Fully Convolutional Networks for Multisource Building Identification. IEEE Transactions on Geoscience and Remote Sensing, 57(1), 108-120.

License

This model is released under the Apache 2.0 License.

Links

• GeoAI package: https://github.com/opengeos/geoai
• Documentation: https://geoai.gishub.org