Swin-UNet — Satellite Imagery Segmentation with Transformers

Revolutionize geospatial analysis with Swin-UNet — a cutting-edge solution for satellite imagery segmentation using Swin Transformers and UNet. Achieve state-of-the-art precision in road extraction, land cover classification, and infrastructure mapping. Ideal for GIS professionals, urban planners, and AI researchers. Train, test, and infer with just three lines of code.

from keras_swin_unet import swin_train, swin_infer

swin_train(data="./dataset", epochs=50)           # train
swin_infer(image="input.jpg", output="out.png")   # predict

Swin-UNet combines the Swin Transformer (hierarchical shifted-window self-attention) with the U-Net encoder-decoder design. It captures global context that CNNs miss while remaining computationally efficient. This package wraps the full architecture into a production-ready Keras pipeline — training, validation, inference, and visualization in a single call.

Installation

pip install keras-swin-unet

Requires Python 3.10+ and TensorFlow 2.19+.

Quick Start

1. Prepare Your Data

dataset/
  images/    # RGB images (.png, .jpg, .tif, .npy)
  masks/     # Pixel-level labels (same filenames, any image format)

For binary segmentation, masks should have pixel values 0 (background) and 255 (foreground). For multi-class, use values 0, 1, 2, ... and set num_classes accordingly.

2. Train

from keras_swin_unet import swin_train

swin_train(
    data="./dataset",
    model_dir="./checkpoint",
    num_classes=2,
    epochs=50,
    bs=8,
    loss_type="focal_dice",
    input_shape=[256, 256, 3],
)

That's it. The training loop handles data splitting (80/10/10), normalization, augmentation-ready loading, focal loss with class balancing, early stopping, model checkpointing, and saves evaluation metrics as JSON.

3. Predict

from keras_swin_unet import swin_infer

# Batch inference on a test set
swin_infer(data="./dataset", model_dir="./checkpoint")

# Single image
swin_infer(image="photo.jpg", output="segmentation.png", model_dir="./checkpoint")

CLI (Alternative)

swin-unet train --data ./dataset --epochs 50 --loss-type focal_dice
swin-unet infer --image photo.jpg --output result.png --model-dir ./checkpoint

Loss Functions

Choose the loss that fits your data. Change it with loss_type="...".

All losses are proper Keras Loss subclasses — models save and load with model.save() / load_model() without custom objects for the loss.

API

swin_train(**kwargs)

Only data is required. Everything else has sensible defaults.

swin_infer(**kwargs)

Architecture

Input [H, W, 3]
  → Patch Embedding (4×4) → [H/4, W/4, C]
  → Encoder Stage 0: Swin Blocks → skip₀
  → Patch Merging → [H/8, W/8, 2C]
  → Encoder Stage 1: Swin Blocks → skip₁
  → Patch Merging → [H/16, W/16, 4C]
  → Encoder Stage 2: Swin Blocks → skip₂
  → Patch Merging → [H/32, W/32, 8C]
  → Encoder Stage 3: Swin Blocks → skip₃
  → Bottleneck: Swin Blocks
  → Patch Expanding → concat skip₂ → Decoder Stage 0
  → Patch Expanding → concat skip₁ → Decoder Stage 1
  → Patch Expanding → concat skip₀ → Decoder Stage 2
  → Patch Expanding (4×) → pixel resolution
  → Conv2D Softmax → Segmentation Map

Each Swin Transformer block alternates regular and shifted-window multi-head self-attention. Shifted windows allow cross-window communication without the quadratic cost of global attention.

Advanced Usage

For full control, build the model directly:

from swin_transformer.model_loader import get_model
from swin_transformer.AUC_LOSS import FocalDiceLoss

model = get_model(
    input_size=(512, 512, 3),
    filter_num_begin=96,
    depth=4,
    stack_num_down=2,
    stack_num_up=2,
    patch_size=[4, 4],
    num_heads=[3, 6, 12, 24],
    window_size=[7, 7, 7, 7],
    num_mlp=1024,
    num_classes=5,           # multi-class segmentation
)
model.compile(
    optimizer="adam",
    loss=FocalDiceLoss(alpha=0.25, gamma=2.0, focal_weight=0.5),
    metrics=["accuracy"],
)
model.fit(train_loader, validation_data=val_loader, epochs=100)

Performance

PennFudan Pedestrian Detection

Binary segmentation, 256×256, 170 images, trained from scratch on CPU. No pre-training, no fine-tuning. 63 epochs, 11.4M params.

Best prediction from the test set (IoU 0.813, Dice 0.897):

Left to right: input image, ground truth mask, model prediction, error map (green = true positive, red = false positive, blue = false negative).

Multi-Sample Comparison

Three test samples with per-sample IoU/Dice scores. The model segments pedestrians accurately despite heavy class imbalance (~10% foreground pixels) and only 118 training images.

Applications

• Remote Sensing & GIS — road extraction, land cover classification, building detection
• Medical Imaging — organ segmentation, tumor delineation, cell counting
• Autonomous Driving — lane marking, pedestrian detection, drivable area
• Agriculture — crop field boundaries, vegetation indices
• Industrial — defect detection, part segmentation, quality control

Citation

If you use this package in your research, please cite the original works:

@inproceedings{cao2022swinunet,
  title={Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation},
  author={Cao, Hu and Wang, Yueyue and Chen, Joy and Jiang, Dongsheng and
          Zhang, Xiaopeng and Tian, Qi and Wang, Manning},
  booktitle={ECCV Workshops},
  year={2022}
}

@inproceedings{liu2021swin,
  title={Swin Transformer: Hierarchical Vision Transformer using Shifted Windows},
  author={Liu, Ze and Lin, Yutong and Cao, Yue and Hu, Han and Wei, Yixuan and
          Zhang, Zheng and Lin, Stephen and Guo, Baining},
  booktitle={ICCV},
  year={2021}
}

License

MIT.

Maintainer: Laeeq Aslam — laeeq.aslam.100@gmail.com