Add additional pretrained models.

Author: Piper Wolters

README.md

@@ -54,16 +54,16 @@ This package will download model checkpoints automatically, but you can download
 |  | Single-image, RGB | Multi-image, RGB |
 | ---------- | ------------ | ------------ |
 | **Swin-v2-Base** | [Sentinel2_SwinB_SI_RGB](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swinb_si_rgb.pth?download=true) | [Sentinel2_SwinB_MI_RGB](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swinb_mi_rgb.pth?download=true) |
-| **Swin-v2-Tiny** | [Sentinel2_SwinT_SI_RGB](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_si_rgb.pth?download=true) | |
+| **Swin-v2-Tiny** | [Sentinel2_SwinT_SI_RGB](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_si_rgb.pth?download=true) | [Sentinel2_SwinT_MI_RGB](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_mi_rgb.pth?download=true) |
 | **Resnet50** | [Sentinel2_Resnet50_SI_RGB](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet50_si_rgb.pth?download=true) | [Sentinel2_Resnet50_MI_RGB](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet50_mi_rgb.pth?download=true) |
 | **Resnet152** | [Sentinel2_Resnet152_SI_RGB](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet152_si_rgb.pth?download=true) | [Sentinel2_Resnet152_MI_RGB](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet152_mi_rgb.pth?download=true) |
 
 |  | Single-image, MS | Multi-image, MS |
 | ---------- | ------------ | ------------ |
 | **Swin-v2-Base** | [Sentinel2_SwinB_SI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swinb_si_ms.pth?download=true) | [Sentinel2_SwinB_MI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swinb_mi_ms.pth?download=true) |
-| **Swin-v2-Tiny** | [Sentinel2_SwinT_SI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_si_ms.pth?download=true) | |
-| **Resnet50** | [Sentinel2_Resnet50_SI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_si_ms.pth?download=true) | |
-| **Resnet152** | [Sentinel2_Resnet152_SI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet152_si_ms.pth?download=true) | |
+| **Swin-v2-Tiny** | [Sentinel2_SwinT_SI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_si_ms.pth?download=true) | [Sentinel2_SwinT_MI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_mi_ms.pth?download=true) |
+| **Resnet50** | [Sentinel2_Resnet50_SI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_si_ms.pth?download=true) | [Sentinel2_Resnet50_MI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet50_mi_ms.pth?download=true) |
+| **Resnet152** | [Sentinel2_Resnet152_SI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet152_si_ms.pth?download=true) | [Sentinel2_Resnet152_MI_MS](https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet152_mi_ms.pth?download=true) |
 
 #### Sentinel-1 Pretrained Models
 |  | Single-image, VH+VV | Multi-image, VH+VV |

satlaspretrain_models/utils.py

@@ -89,6 +89,18 @@ class Head(Enum):
         'num_channels': 9,
         'multi_image': False
     },
+    'Sentinel2_SwinT_MI_RGB': {
+        'url': 'https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_mi_rgb.pth?download=true',
+        'backbone': Backbone.SWINT,
+        'num_channels': 3,
+        'multi_image': True
+    },
+    'Sentinel2_SwinT_MI_MS': {
+        'url': 'https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_swint_mi_ms.pth?download=true',
+        'backbone': Backbone.SWINT,
+        'num_channels': 9,
+        'multi_image': True
+    },
     'Sentinel2_Resnet50_SI_RGB': {
         'url': 'https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet50_si_rgb.pth?download=true',
         'backbone': Backbone.RESNET50,
@@ -107,6 +119,12 @@ class Head(Enum):
         'num_channels': 3,
         'multi_image': True
     },
+    'Sentinel2_Resnet50_MI_MS': {
+        'url': 'https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet50_mi_ms.pth?download=true',
+        'backbone': Backbone.RESNET50,
+        'num_channels': 9,
+        'multi_image': True
+    },
     'Sentinel2_Resnet152_SI_RGB': {
         'url': 'https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet152_si_rgb.pth?download=true',
         'backbone': Backbone.RESNET152,
@@ -125,6 +143,12 @@ class Head(Enum):
         'num_channels': 3,
         'multi_image': True
     },
+    'Sentinel2_Resnet152_MI_MS': {
+        'url': 'https://huggingface.co/allenai/satlas-pretrain/resolve/main/sentinel2_resnet152_mi_ms.pth?download=true',
+        'backbone': Backbone.RESNET152,
+        'num_channels': 9,
+        'multi_image': True
+    },
 }