Package: satlaspretrain-models
SatlasPretrain Models: Foundation models for satellite and aerial imagery.SatlasPretrainis a large-scale pre-training dataset for remote sensing image understanding. This work
was published at ICCV 2023. Details and download links for the dataset can be foundhere.This repository containssatlaspretrain_models, a lightweight library to easily load pretrained SatlasPretrain models for:Sentinel-2Sentinel-1Landsat 8/90.5-2 m/pixel aerial imageryThese models can be fine-tuned on downstream tasks that use these image sources, leading to faster training
and improved performance compared to training from other initializations.Model Structure and UsageThe SatlasPretrain models consist of three main components: backbone, feature pyramid network (FPN), and prediction head.For models trained onmulti-imageinput, the backbone is applied on each individual image, and then max pooling is applied
in the temporal dimension, i.e., across the multiple aligned images.Single-imagemodels input an individual image.This package allows you to load the backbone or backbone+FPN using a model checkpoint ID from the tables below.MODEL_CHECKPOINT_ID="Sentinel2_SwinB_SI_RGB"model=weights_manager.get_pretrained_model(MODEL_CHECKPOINT_ID)model=weights_manager.get_pretrained_model(MODEL_CHECKPOINT_ID,fpn=True)The output of the model is the multi-scale feature map (either from the backbone or from the FPN).For a complete fine-tuning example,see our tutorial on fine-tuning the pre-trained model on EuroSAT.Installationconda create --name satlaspretrain python==3.9
conda activate satlaspretrain
pip install satlaspretrain-modelsAvailable Pretrained ModelsThe tables below list available model checkpoint IDs (likeSentinel2_SwinB_SI_RGB).
Checkpoints are released underODC-BY.
This package will download model checkpoints automatically, but you can download them directly using the links below if desired.Sentinel-2 Pretrained ModelsSingle-image, RGBMulti-image, RGBSwin-v2-BaseSentinel2_SwinB_SI_RGBSentinel2_SwinB_MI_RGBSwin-v2-TinySentinel2_SwinT_SI_RGBResnet50Sentinel2_Resnet50_SI_RGBSentinel2_Resnet50_MI_RGBResnet152Sentinel2_Resnet152_SI_RGBSentinel2_Resnet152_MI_RGBSingle-image, MSMulti-image, MSSwin-v2-BaseSentinel2_SwinB_SI_MSSentinel2_SwinB_MI_MSResnet50Sentinel2_Resnet50_SI_MSResnet152Sentinel2_Resnet152_SI_MSSentinel-1 Pretrained ModelsSingle-image, VH+VVMulti-image, VH+VVSwin-v2-BaseSentinel1_SwinB_SISentinel1_SwinB_MILandsat 8/9 Pretrained ModelsSingle-image, all bandsMulti-image, all bandsSwin-v2-BaseLandsat_SwinB_SILandsat_SwinB_MIAerial (0.5-2m/px high-res imagery) Pretrained ModelsSingle-image, RGBMulti-image, RGBSwin-v2-BaseAerial_SwinB_SIAerial_SwinB_MISingle-imagemodels learn strong representations for individual satellite or aerial images, whilemulti-imagemodels use multiple image captures of the same location for added robustness when making predictions about static objects. Inmulti-imagemodels, feature maps from the backbone are passed through temporal max pooling, so the backbone itself is still applied on individual images, but is trained to provide strong representations after the temporal max pooling step. SeeModelArchitecture.mdfor more details.Sentinel-2RGBmodels input the B4, B3, and B2 bands only, while the multi-spectral (MS) models input 9 bands. The aerial (0.5-2m/px high-res imagery) models inputRGBNAIP and other high-res images, and we have found them to be effective on aerial imagery from a variety of sources. Landsat models input B1-B11 (all bands). Sentinel-1 models inputVV and VHbands.
SeeNormalization.mdfor details on how pixel values should be normalized for input to the pre-trained models.Usage ExamplesFirst initialize aWeightsinstance:importsatlaspretrain_modelsimporttorchweights_manager=satlaspretrain_models.Weights()Then choose amodel_identifierfrom the tables above to specify the pretrained model you want to load.
Below are examples showing how to load in a few of the available models.Pretrained single-image Sentinel-2 RGB model, backbone only:model=weights_manager.get_pretrained_model(model_identifier="Sentinel2_SwinB_SI_RGB")# Expected input is a portion of a Sentinel-2 L1C TCI image.# The 0-255 pixel values should be divided by 255 so they are 0-1.# tensor = tci_image[None, :, :, :] / 255tensor=torch.zeros((1,3,512,512),dtype=torch.float32)# Since we only loaded the backbone, it outputs feature maps from the Swin-v2-Base backbone.output=model(tensor)print([feature_map.shapeforfeature_mapinoutput])# [torch.Size([1, 128, 128, 128]), torch.Size([1, 256, 64, 64]), torch.Size([1, 512, 32, 32]), torch.Size([1, 1024, 16, 16])]Pretrained single-image Sentinel-1 model, backbone+FPNmodel=weights_manager.get_pretrained_model("Sentinel1_SwinB_SI",fpn=True)# Expected input is a portion of a Sentinel-1 vh+vv image (in that order).# The 16-bit pixel values should be divided by 255 and clipped to 0-1 (any pixel values greater than 255 become 1).# tensor = torch.clip(torch.stack([vh_image, vv_image], dim=0)[None, :, :, :] / 255, 0, 1)tensor=torch.zeros((1,2,512,512),dtype=torch.float32)# The model outputs feature maps from the FPN.output=model(tensor)print([feature_map.shapeforfeature_mapinoutput])# [torch.Size([1, 128, 128, 128]), torch.Size([1, 128, 64, 64]), torch.Size([1, 128, 32, 32]), torch.Size([1, 128, 16, 16])]Prediction headsAlthough the checkpoints include prediction head parameters, these heads are task-specific, so loading the head parameters is not supported in this repository.
Computing outputs from the pre-trained prediction heads is supportedin the dataset codebase.For convenience when fine-tuning on certain types of tasks, though,satlaspretrain_modelssupports attaching certain heads (initialized randomly) to the pre-trained model:# 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)The following head architectures are available:Segmentation: U-Net Decoder w/ Cross Entropy lossDetection: Faster R-CNN DecoderInstance Segmentation: Mask R-CNN DecoderRegression: U-Net Decoder w/ L1 lossClassification: Pooling + Linear layersMulti-label Classification: Pooling + Linear layersPretrained multi-image aerial model, backbone + FPN + classification head:# num_categories is the number of categories to predict.# All heads are randomly initialized and provided only for convenience for fine-tuning.model=weights_manager.get_pretrained_model("Aerial_SwinB_MI",fpn=True,head=satlaspretrain_models.Head.CLASSIFY,num_categories=2)# Expected input is 8-bit (0-255) aerial images at 0.5 - 2 m/pixel.# The 0-255 pixel values should be divided by 255 so they are 0-1.# This multi-image model is trained to input 4 images but should perform well with different numbers of images.# tensor = torch.stack([rgb_image1, rgb_image2], dim=0)[None, :, :, :, :] / 255tensor=torch.zeros((1,4,3,512,512),dtype=torch.float32)# The head needs to be fine-tuned on a downstream classification task.# It outputs classification probabilities.model.eval()output=model(tensor.reshape(1,4*3,512,512))print(output)# tensor([[0.0266, 0.9734]])Pretrained multi-image Landsat model, backbone + FPN + detection head# num_categories is the number of bounding box detection categories.# All heads are randomly initialized and provided only for convenience for fine-tuning.model=weights_manager.get_pretrained_model("Landsat_SwinB_MI",fpn=True,head=satlaspretrain_models.Head.DETECT,num_categories=5)# Expected input is Landsat B1-B11 stacked in order.# This multi-image model is trained to input 8 images but should perform well with different numbers of images.# The 16-bit pixel values are normalized as follows:# landsat_images = torch.stack([landsat_image1, landsat_image2], dim=0)# tensor = torch.clip(landsat_images[None, :, :, :, :]-4000)/16320, 0, 1)tensor=torch.zeros((1,8,11,512,512),dtype=torch.float32)# The head needs to be fine-tuned on a downstream object detection task.# It outputs bounding box detections.model.eval()output=model(tensor.reshape(1,8*11,512,512))print(output)#[{'boxes': tensor([[ 67.0772, 239.2646, 95.6874, 16.3644], ...]),# 'labels': tensor([3, ...]),# 'scores': tensor([0.5443, ...])}]DemosWe provide ademoshowing how to finetune a
SatlasPretrain Sentinel-2 model on the EuroSAT classification task.TestsThere are tests to test loading pretrained models and one to test randomly initialized models.To run the tests, run the following command from the root directory:pytest tests/ContactIf you have any questions, please emailsatlas@allenai.orgor open an issuehere.