Skip to content

Latest commit

 

History

History
246 lines (185 loc) · 9.11 KB

File metadata and controls

246 lines (185 loc) · 9.11 KB

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