Skip to content

Latest commit

Β 

History

1 Commit

Folders and files

NameName
Last commit message
Last commit date
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

CABNet: Content-Adaptive Building Segmentation Network

CABNet Architecture

PyTorch Python License arXiv

πŸ“‹ Overview

CABNet is a novel deep learning architecture for building extraction from remote sensing imagery. It features adaptive multi-scale convolution, corner-guided enhancement, bidirectional feature pyramid, and dual-stream boundary refinement for accurate building segmentation.

πŸ”‘ Key Features

  • Content-Adaptive Scale Convolution (CASC): Dynamically predicts optimal receptive field scale based on content
  • Corner-Guided Feature Enhancement (CGFE): Leverages building corner detection for enhanced feature representation
  • Bidirectional Feature Pyramid (BiFPN): Efficient top-down and bottom-up feature fusion
  • Local Window Cross-Attention (LWCA): Linear complexity attention mechanism with relative position encoding
  • Dual-Stream Boundary Refinement (DSBR): Independent semantic and boundary modeling with interactive enhancement

πŸ—οΈ Architecture

Input Image
    ↓
β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
β”‚                    Encoder (ResNet-50)                        β”‚
β”‚   β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”   β”‚
β”‚   β”‚  C1     β”‚ β†’  β”‚  C2     β”‚ β†’  β”‚  C3     β”‚ β†’  β”‚  C4     β”‚   β”‚
β”‚   β”‚ H/4     β”‚    β”‚ H/8     β”‚    β”‚ H/16    β”‚    β”‚ H/32    β”‚   β”‚
β”‚   β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜    β””β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”˜    β””β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”˜   β”‚
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”˜
                                         ↓              ↓
                                    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                                    β”‚  CASC + CGFE + LWCA   β”‚
                                    β”‚  (Feature Enhancement) β”‚
                                    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                              ↓
                                    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                                    β”‚       BiFPN            β”‚
                                    β”‚ (Bidirectional Fusion) β”‚
                                    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                              ↓
                                    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                                    β”‚        DSBR            β”‚
                                    β”‚ (Boundary Refinement)  β”‚
                                    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                              ↓
                              β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                              β”‚  Segmentation  β”‚   Boundary   β”‚
                              β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜

πŸ“¦ Installation

Requirements

  • Python >= 3.8
  • PyTorch >= 2.0
  • CUDA >= 11.8 (for GPU support)

Install from source

git clone https://github.com/yourusername/CABNet.git
cd CABNet
pip install -r requirements.txt

πŸš€ Quick Start

Inference

import torch
from models import CABNet

# Create model
model = CABNet(num_classes=2, backbone='resnet50', pretrained=True)
model.eval()

# Inference
image = torch.randn(1, 3, 512, 512)
with torch.no_grad():
    outputs = model(image)
    segmentation = outputs['seg'].argmax(dim=1)
    boundary = outputs['boundary']

Training

import torch
from models import CABNet, CABNetLoss

# Create model and loss
model = CABNet(num_classes=2, backbone='resnet50', pretrained=True).cuda()
criterion = CABNetLoss(num_classes=2, lambda_boundary=2.0, lambda_corner=0.5)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)

# Training loop
model.train()
for epoch in range(200):
    for images, masks, boundaries in train_loader:
        images = images.cuda()
        targets = {
            'seg': masks.cuda(),
            'boundary': boundaries.cuda()
        }
        
        outputs = model(images)
        loss_dict = criterion(outputs, targets)
        
        optimizer.zero_grad()
        loss_dict['total'].backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

See train.py for a complete training script.

πŸ“ Project Structure

CABNet/
β”œβ”€β”€ configs/
β”‚   └── default.yaml         # Default configuration
β”œβ”€β”€ datasets/
β”‚   └── building_dataset.py  # Dataset implementation
β”œβ”€β”€ models/
β”‚   β”œβ”€β”€ __init__.py
β”‚   └── cabnet.py            # CABNet model
β”œβ”€β”€ utils/
β”‚   β”œβ”€β”€ __init__.py
β”‚   β”œβ”€β”€ metrics.py           # Evaluation metrics
β”‚   └── visualization.py     # Visualization tools
β”œβ”€β”€ docs/
β”‚   └── assets/              # Documentation assets
β”œβ”€β”€ train.py                 # Training script
β”œβ”€β”€ test.py                  # Testing script
β”œβ”€β”€ requirements.txt         # Dependencies
β”œβ”€β”€ LICENSE                  # MIT License
└── README.md                # This file

πŸ“Š Results

Performance on Building Extraction Datasets

Dataset IoU (%) F1 (%) Boundary IoU (%)
WHU Building - - -
Inria Aerial - - -
Massachusetts Buildings - - -

Results will be updated after training.

πŸ“– Model Components

Content-Adaptive Scale Convolution (CASC)

CASC dynamically predicts optimal receptive field scale for each spatial location:

$$w = \text{Softmax}(\text{GlobalPredict}(F) + \text{LocalRefine}(F))$$

$$F' = \sum_{k}(w_k \odot \text{Conv}_k(F)) + \text{Residual}(F)$$

Corner-Guided Feature Enhancement (CGFE)

CGFE leverages Harris corner detection to enhance structural features:

$$M = \begin{bmatrix} \sum I_x^2 & \sum I_x I_y \ \sum I_x I_y & \sum I_y^2 \end{bmatrix}$$

$$R = \det(M) - k \cdot \text{trace}(M)^2$$

Loss Function

The total loss combines multiple components:

$$\mathcal{L}_{total} = \mathcal{L}_{semantic} + \lambda_b \mathcal{L}_{boundary} + \lambda_c \mathcal{L}_{corner}$$

Where:

  • $\mathcal{L}_{semantic}$: Cross-Entropy + LovΓ‘sz-Softmax Loss
  • $\mathcal{L}_{boundary}$: Boundary IoU Loss
  • $\mathcal{L}_{corner}$: Corner Supervision Loss

βš™οΈ Configuration

Key hyperparameters in configs/default.yaml:

Parameter Default Description
num_classes 2 Number of segmentation classes
backbone resnet50 Encoder backbone
lr 1e-4 Learning rate
weight_decay 1e-4 Weight decay
lambda_boundary 2.0 Boundary loss weight
lambda_corner 0.5 Corner loss weight
window_size 8 LWCA window size

πŸ”§ Training Tips

  1. Learning Rate: Start with 1e-4, use cosine annealing
  2. Data Augmentation: Random flip, rotation, color jitter
  3. Mixed Precision: Enable AMP for faster training
  4. Gradient Clipping: Clip gradients to max_norm=1.0
  5. Warm-up: Use 5-10 epochs of warm-up

πŸ“ Citation

If you find this work useful, please cite:

@article{cabnet2024,
  title={CABNet: Content-Adaptive Building Segmentation Network},
  author={Your Name},
  journal={arXiv preprint arXiv:XXXX.XXXXX},
  year={2024}
}

πŸ“„ License

This project is licensed under the MIT License - see the LICENSE file for details.

πŸ™ Acknowledgements

πŸ“§ Contact

For questions or collaboration, please open an issue or contact your.email@example.com.


Made with ❀️ for the remote sensing community

About

Content-Adaptive Building Segmentation Network for Remote Sensing

Resources

Contributing

Stars

3 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages