BASeg: Boundary-Aware Segmentation

📋 Table of Contents

• Overview
• Features
• Installation
• Quick Start
• Dataset Preparation
• Model Training
• Model Testing & Evaluation
• Project Structure

Overview

BASeg: Boundary-Aware Segmentation Model with Multi-Scale Fusion This repository contains the complete implementation including:

• 🏗️ State-of-the-art model architecture combining Global State Module and Cross Fusion module
• 📊 Support for multi-class semantic segmentation
• 🔬 Angle Loss and Mahalanobis-Distance-Based Boundary Loss
• 🚀 Distributed training support for multi-GPU setups
• 📈 Comprehensive evaluation metrics (Accuracy, F1-Score, mIoU, Kappa coefficient)

Features

• Boundary-Aware Segmentation: Special attention to object boundaries through architecture and loss function
• Multi-Scale Fusion: Leverages features at multiple scales for better context understanding
• Multi-GPU Training: Distributed Data Parallel (DDP) support for efficient training

Training Objective

BASeg combines two complementary boundary-aware objectives:

• Angle Loss penalizes disagreement between predicted and ground-truth contour orientations. It uses the absolute cosine of the orientation difference, so directions separated by $\pi$ represent the same unoriented boundary.
• Mahalanobis-Distance-Based Boundary Loss weights cross-entropy using a region-wise Mahalanobis-distance prior, emphasizing structurally important boundary pixels.

Installation

Prerequisites

Ensure you have the following installed:

• Python: 3.7 or higher (recommended: 3.8+)
• CUDA: 11.8 or higher (recommended: 12.4)
• cuDNN: Compatible with your CUDA version

Step 1: Create Virtual Environment (Recommended)

# Using venv
python -m venv venv
source venv/bin/activate 

# Or using conda
conda create -n baseg python=3.8
conda activate baseg

Step 2: Install Dependencies

pip install -r requirements.txt

Key Dependencies:

• PyTorch >= 1.9.0 (with CUDA support)
• torchvision
• scikit-image
• scikit-learn
• scipy
• numpy
• Pillow
• timm (for pretrained models)
• einops
• transformers
• tqdm

To install PyTorch with CUDA support (example for CUDA 12.1):

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

Step 3: Add the Pretrained VMamba-Tiny Checkpoint

Download the ImageNet-pretrained VMamba-Tiny checkpoint and place it at:

BASeg/pretrain/vmamba_tiny_e292.pth

The expected project structure is:

BASeg/
├── pretrain/
│   └── vmamba_tiny_e292.pth
├── train.py
└── utils.py

Run the training command from inside the BASeg directory so that the relative checkpoint path resolves correctly. The pretrained ResNet-18 and DINOv3 weights are downloaded automatically by timm and Hugging Face on the first run; only the VMamba-Tiny checkpoint must be placed manually.

Quick Start

1. Prepare Your Dataset

Update the following parameters in utils.py:

MAIN_FOLDER = "/path/to/your/dataset/"
DATA_FOLDER = MAIN_FOLDER + "images_png/{}.png"
LABEL_FOLDER = MAIN_FOLDER + "masks_png/{}.png"

# Update train_ids and test_ids with your dataset IDs
train_ids = ["id1", "id2", ...]
test_ids = ["test_id1", "test_id2", ...]

2. Configure Training Settings

Edit utils.py to adjust model hyperparameters:

# Model configuration
WINDOW_SIZE = (256, 256)    # Patch size for training
STRIDE = 32                 # Stride for testing
BATCH_SIZE = 32             # Batch size
IN_CHANNELS = 3             # RGB input
N_CLASSES = 7               # Number of segmentation classes

# Mode selection
MODE = "Train"              # Change to "Test" for inference
DATASET = "Urban"           # Dataset name

3. Train the Model

python train.py

The model auto detects and uses all available GPUs

4. Test/Evaluate the Model

Change the mode in utils.py:

MODE = "Test"

Then run:

python train.py

Dataset Preparation

Included Sample Data

A small sample from GCD-25k is included in data/GCD-25k/. It contains 20 RGB image tiles (two tiles from each of ten cities) and 20 corresponding single-channel masks. All tiles are 512 x 512 pixels. This sample is intended for checking the data pipeline only; it is too small for a meaningful training or evaluation experiment.

The sample uses the following correspondence:

data/GCD-25k/images_png/Abuja_0.png
data/GCD-25k/masks_png/Abuja_0_label.png

In general, an image named <id>.png corresponds to a mask named <id>_label.png. The mask values are 0, 1, and 2, matching the configured classes background, building, and road, respectively.

To use the included sample, configure utils.py as follows:

MAIN_FOLDER = "./data/GCD-25k/"
DATA_FOLDER = MAIN_FOLDER + "images_png/{}.png"
LABEL_FOLDER = MAIN_FOLDER + "masks_png/{}_label.png"

LABELS = ["background", "building", "road"]
N_CLASSES = len(LABELS)
WEIGHTS = torch.tensor([1, 1, 1], dtype=torch.float)

Populate train_ids and test_ids with the image names without .png, for example "Abuja_0". Keep the training and test lists non-empty.

The included masks are already zero-based (0 to 2). Therefore, when using this sample, the mask-loading code in dataset.__getitem__ must not subtract one. Use:

label = np.asarray(io.imread(self.label_files[random_idx]), dtype="int64")

instead of loading the mask with - 1. The subtraction is only appropriate for datasets whose stored class indices begin at 1.

Supported Format

Your dataset should be organized as follows:

dataset/
├── images_png/
│   ├── 1366.png
│   ├── 1367.png
│   └── ...
└── masks_png/
    ├── 1366.png
    ├── 1367.png
    └── ...

Image Requirements

• Format: PNG (supports multi-channel images)
• Input Images: RGB (3 channels)
• Labels: Single-channel grayscale with class indices

Model Training

Basic Training

python train.py

Training Configuration

Edit train.py to customize:

# Learning rate, optimizer, scheduler
# Loss function weights
# Number of epochs
# Checkpoint saving frequency

Multi-GPU Distributed Training

The training script supports automatic distributed training detection:

# Automatic setup (recommended)
python train.py

# Or explicit distributed launch
torchrun --nproc_per_node=4 train.py

Monitoring Training

The script prints:

• GPU availability and device information
• Per-epoch training loss
• Validation metrics
• Checkpoint paths

Checkpoints

Trained model checkpoints are saved during training. The model supports:

• Loading pretrained VMamba encoder weights
• Resuming training from checkpoints
• Fine-tuning on new datasets

Model Testing & Evaluation

Inference

Switch to test mode and run:

# In utils.py
MODE = "Test"

python train.py

Advanced Usage

Custom Loss Function Weights

Adjust class weights in utils.py:

WEIGHTS = torch.tensor([1, 2, 2, 1, 1, 1, 1], dtype=torch.float)