SatlasPretrain Models

To fine-tune a SatlasPretrain model on a custom dataset (like UCMerced) using Torchgeo, you should extend the ClassificationTask class.

In the configure_models method, you must:

1. Load the model using swin_v2_b with the Swin_V2_B_Weights.SENTINEL2_RGB_SI_SATLAS weights.
2. Adjust the first layer's input channels to match your task's input requirements.
3. Replace the model's head with a new torch.nn.Linear layer that matches your number of classes.

Example implementation:

from torchgeo.trainers import ClassificationTask
from torchgeo.models import Swin_V2_B_Weights, swin_v2_b
import torch

class SatlasClassificationTask(ClassificationTask):
    def configure_models(self):
        # Load pretrained Satlas weights
        weights = Swin_V2_B_Weights.SENTINEL2_RGB_SI_SATLAS
        self.model = swin_v2_b(weights)

        # Replace first layer's input channels
        first_layer = self.model.features[0][0]
        self.model.features[0][0] = torch.nn.Conv2d(
            3, # Number of input channels
            first_layer.out_channels,
            kernel_size=first_layer.kernel_size,
            stride=first_layer.stride,
            padding=first_layer.padding,
            bias=(first_layer.bias is not None)
        )

        # Replace last layer's output features with the number of classes
        self.model.head = torch.nn.Linear(
            in_features=1024, 
            out_features=self.hparams["num_classes"],
            bias=True
        )

    def on_validation_epoch_end(self):
        # Accessing metrics logged during the current validation epoch
        val_loss = self.trainer.callback_metrics.get('val_loss', 'N/A')
        val_acc = self.trainer.callback_metrics.get('val_OverallAccuracy', 'N/A')
        print(f"Epoch {self.current_epoch} Validation - Loss: {val_loss}, Accuracy: {val_acc}")

# Usage
task = SatlasClassificationTask(num_classes=21)

class SatlasClassificationTask(ClassificationTask):
    def configure_models(self):
        weights = Swin_V2_B_Weights.SENTINEL2_RGB_SI_SATLAS
        self.model = swin_v2_b(weights)

        first_layer = self.model.features[0][0]
        self.model.features[0][0] = torch.nn.Conv2d(3,
                                    first_layer.out_channels,
                                    kernel_size=first_layer.kernel_size,
                                    stride=first_layer.stride,
                                    padding=first_layer.padding,
                                    bias=(first_layer.bias is not None))

        self.model.head = torch.nn.Linear(in_features=1024, out_features=self.hparams["num_classes"], bias=True)

This guide demonstrates how to fine-tune a Sentinel2_SwinB_SI_RGB model using the EuroSAT classification task. The process involves downloading the dataset, preparing a PyTorch Dataset and DataLoader, initializing the pretrained model with a classification head, and running a standard training loop.

Model Initialization

To use the model for classification, use the Weights manager to load the backbone with a specific head configuration. For EuroSAT, you should set num_categories=10 and use satlaspretrain_models.Head.CLASSIFY.

Data Preprocessing Requirements

When using Sentinel-2 TCI images, ensure the input is normalized: pixel values (0-255) must be divided by 255 to scale them to the 0-1 range.

import satlaspretrain_models
import torch

weights_manager = satlaspretrain_models.Weights()
model = weights_manager.get_pretrained_model(
    "Sentinel2_SwinB_SI_RGB", 
    fpn=True, 
    head=satlaspretrain_models.Head.CLASSIFY, 
    num_categories=10, 
    device='cpu'
)
model = model.to('cpu')

While the repository does not support loading pretrained prediction heads (as they are task-specific), you can attach randomly initialized heads to the backbone or FPN to facilitate fine-tuning. Use the head and num_categories arguments in get_pretrained_model().

Available Head Architectures (satlaspretrain_models.Head):

• Segmentation: U-Net Decoder w/ Cross Entropy loss
• Detection: Faster R-CNN Decoder
• Instance Segmentation: Mask R-CNN Decoder
• Regression: U-Net Decoder w/ L1 loss
• Classification: Pooling + Linear layers
• Multi-label Classification: Pooling + Linear layers

Example: Attaching a classification head:

model = weights_manager.get_pretrained_model(
    "Aerial_SwinB_MI", 
    fpn=True, 
    head=satlaspretrain_models.Head.CLASSIFY, 
    num_categories=2
)

# Backbone and FPN parameters initialized from checkpoint, head parameters initialized randomly.
model = weights_manager.get_pretrained_model(MODEL_CHECKPOINT_ID, fpn=True, head=satlaspretrain_models.Head.CLASSIFY, num_categories=2)

Use these identifiers with weights_manager.get_pretrained_model() to load specific pretrained weights. Checkpoints are categorized by sensor type, image mode (Single-image vs Multi-image), and spectral configuration.

Sentinel-2

• Swin-v2-Base (RGB): Sentinel2_SwinB_SI_RGB (SI), Sentinel2_SwinB_MI_RGB (MI)
• Swin-v2-Tiny (RGB): Sentinel2_SwinT_SI_RGB (SI), Sentinel2_SwinT_MI_RGB (MI)
• Resnet50 (RGB): Sentinel2_Resnet50_SI_RGB (SI), Sentinel2_Resnet50_MI_RGB (MI)
• Resnet152 (RGB): Sentinel2_Resnet152_SI_RGB (SI), Sentinel2_Resnet152_MI_RGB (MI)
• Swin-v2-Base (MS - 9 bands): Sentinel2_SwinB_SI_MS (SI), Sentinel2_SwinB_MI_MS (MI)
• Swin-v2-Tiny (MS - 9 bands): Sentinel2_SwinT_SI_MS (SI), Sentinel2_SwinT_MI_MS (MI)
• Resnet50 (MS - 9 bands): Sentinel2_Resnet50_SI_MS (SI), Sentinel2_Resnet50_MI_MS (MI)
• Resnet152 (MS - 9 bands): Sentinel2_Resnet152_SI_MS (SI), Sentinel2_Resnet152_MI_MS (MI)

Sentinel-1 (VH+VV)

• Swin-v2-Base: Sentinel1_SwinB_SI (SI), Sentinel1_SwinB_MI (MI)

Landsat 8/9 (All bands)

• Swin-v2-Base: Landsat_SwinB_SI (SI), Landsat_SwinB_MI (MI)

Aerial (0.5-2m/px RGB)

• Swin-v2-Base: Aerial_SwinB_SI (SI), Aerial_SwinB_MI (MI)