Hugging Face transformers model files

config.json

@@ -0,0 +1,24 @@
+{   
+    "model_type": "SMARTIES-v1-ViT-B", 
+    "auto_map": {
+      "AutoModel": "smarties_model.SMARTIESHF",
+      "AutoConfig": "smarties_config.SMARTIESConfig"
+    },
+    "transformers_weights": "smarties-v1-vitb.safetensors",
+    "spectrum_specs": null,
+    "img_size": 224,
+    "patch_size": 16,
+    "embed_dim": 768,
+    "depth": 12,
+    "num_heads": 12,
+    "mlp_ratio": 4.0,
+    "qkv_bias": true,
+    "norm_eps": 1e-6,
+    "global_pool": false,
+    "pos_drop_rate": 0.0,
+    "norm_layer_eps": 1e-6,
+    "mixed_precision": "no",
+    "decoder_embed_dim": 512,
+    "decoder_depth": 8,
+    "decoder_num_heads": 16
+}

smarties_config.py

@@ -0,0 +1,48 @@
+from transformers import PretrainedConfig
+import os
+import yaml
+import requests
+from functools import partial
+import torch.nn as nn
+
+class SMARTIESConfig(PretrainedConfig):
+    model_type = "SMARTIES-v1-ViT-B"
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        norm_eps=1e-6,
+        spectrum_specs=None,
+        global_pool=False, 
+        norm_layer_eps=1e-6,
+        mixed_precision='no',
+        decoder_embed_dim=512, 
+        decoder_depth=8, 
+        decoder_num_heads=16,
+        pos_drop_rate=0.0,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+        self.qkv_bias = qkv_bias
+        self.norm_eps = norm_eps
+        self.spectrum_specs = spectrum_specs
+        self.global_pool = global_pool
+        self.pos_drop_rate = pos_drop_rate
+        self.num_heads = self.num_heads
+        self.norm_layer_eps = norm_layer_eps
+        self.mixed_precision = mixed_precision
+        self.decoder_embed_dim = decoder_embed_dim
+        self.decoder_depth = decoder_depth
+        self.decoder_num_heads = decoder_num_heads

smarties_model.py

@@ -0,0 +1,328 @@
+from functools import partial
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.utils import cached_file
+from .smarties_config import SMARTIESConfig
+from functools import partial
+import numpy as np
+from timm.models.vision_transformer import Block
+import os 
+import yaml
+
+class SpectrumRangeProjection(nn.Module):
+    """Patch Embedding of a sensor without patchify"""
+    def __init__(
+            self,
+            spectral_range,
+            spectrum_spec,
+            patch_size,
+            embed_dim,
+            bias=True
+    ):
+        super().__init__()
+        self.spectral_range = spectral_range
+        self.name = spectrum_spec['name']
+        self.min_wavelength = spectrum_spec['min_wavelength']
+        self.max_wavelength = spectrum_spec['max_wavelength']
+        self.sensors = spectrum_spec['sensors']
+        self.nb_pixels = patch_size**2
+        self.proj = nn.Linear(self.nb_pixels, embed_dim, bias=bias)
+    
+    def forward(self, x):
+        return self.proj(x.view(-1, self.nb_pixels)) 
+
+class SpectrumRangeProjectionAvg(nn.Module):
+    """Patch Embedding of a sensor without patchify"""
+    def __init__(
+            self,
+            spectrum_projections,
+            spectrum_spec,
+            embed_dim
+    ):
+        super().__init__()
+        self.min_wavelength = spectrum_spec['min_wavelength']
+        self.max_wavelength = spectrum_spec['max_wavelength']
+        self.central_lambda = 0.5*(float(self.min_wavelength) + float(self.max_wavelength))
+        self.spectrum_projections = spectrum_projections
+        self.weights = []
+        for spectrum_proj in self.spectrum_projections:
+            central_lambda = 0.5*(float(spectrum_proj.min_wavelength) + float(spectrum_proj.max_wavelength))
+            self.weights.append(abs(self.central_lambda-central_lambda))
+        self.weights = np.array(self.weights) / sum(self.weights)
+        self.embed_dim = embed_dim
+    
+    def forward(self, x):
+        out = 0. #torch.zeros((len(x),self.embed_dim))
+        for i, spectrum_proj in enumerate(self.spectrum_projections):
+            out += spectrum_proj(x) * self.weights[i]
+        return out
+    
+
+class SpectrumAwareProjection(nn.Module):
+    """Patch Embedding of a sensor without patchify"""
+    def __init__(
+            self,
+            spectrum_specs,
+            patch_size,
+            embed_dim,
+            bias=True
+    ):
+        super().__init__()
+        self.nb_pixels = patch_size**2
+    
+        self.spectrum_embeds = torch.nn.ModuleList()
+        for spectral_range in sorted(spectrum_specs,key=lambda key:spectrum_specs[key]['projection_idx']):
+            if ((spectrum_specs[spectral_range]['projection_idx'] != -1) and (len(spectrum_specs[spectral_range]['agg_projections']) == 0)) :
+                self.spectrum_embeds.append(SpectrumRangeProjection(
+                    spectral_range, spectrum_specs[spectral_range], patch_size, embed_dim
+                ))
+
+        for spectral_range in sorted(spectrum_specs,key=lambda key:spectrum_specs[key]['projection_idx']): 
+            if ((spectrum_specs[spectral_range]['projection_idx'] != -1) and (len(spectrum_specs[spectral_range]['agg_projections']) > 0)):
+                self.spectrum_embeds.append(
+                    SpectrumRangeProjectionAvg(
+                        [self.spectrum_embeds[agg_proj_idx] for agg_proj_idx in spectrum_specs[spectral_range]['agg_projections']], 
+                        spectrum_specs[spectral_range],
+                        embed_dim))
+                
+    def forward(self, x, projection_idx):
+        return self.spectrum_embeds[projection_idx](x)
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=float)
+    grid_w = np.arange(grid_size, dtype=float)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+def get_dtype(mixed_precision):
+    if mixed_precision == 'no':
+        return torch.float32
+    elif mixed_precision == 'bf16':
+        return torch.bfloat16
+    elif mixed_precision == 'fp16':
+        return torch.float16
+    else:
+        raise NotImplementedError
+
+class SMARTIESHF(PreTrainedModel):
+    config_class = SMARTIESConfig
+    def __init__(self, config: SMARTIESConfig):
+        super().__init__(config)
+        try:
+            if config.spectrum_specs is None:
+                spectrum_path = cached_file(
+                    config.name_or_path, 
+                    "spectrum_specs.yaml"
+                )
+                with open(spectrum_path, "r") as f:
+                    config.spectrum_specs = yaml.safe_load(f)
+        except Exception as e:
+                raise RuntimeError(
+                "spectrum_specs couldn't be loaded from spectrum_specs.yaml. " \
+                "Please load yaml file yourself and provide the argument spectrum_specs with the loaded file."
+            ) from e
+        self.model_dtype = get_dtype(config.mixed_precision)
+        self.embed_dim = config.embed_dim
+        self.decoder_embed_dim = config.decoder_embed_dim
+        self.projection_conversion = {i: config.spectrum_specs[i]['projection_idx'] for i in config.spectrum_specs}
+        self.sensor_band_specs = {
+            'S2': [
+                'aerosol', 
+                'blue_1', 
+                'green_2', 
+                'red_2', 
+                'red_edge_1', 
+                'red_edge_2', 
+                'near_infrared_2', 
+                'near_infrared_1', 
+                'near_infrared_3', 
+                'short_wave_infrared_1', 
+                'short_wave_infrared_3', 
+                'short_wave_infrared_4'
+            ],
+            'S1': [
+                'microwave_1', 
+                'microwave_2'
+            ],
+            'RGB': [
+                'red_1', 
+                'green_1', 
+                'blue_3'
+            ]
+        }
+        self.sensor_projection_specs = {}
+        for sensor_name in self.sensor_band_specs:
+            self.sensor_projection_specs[sensor_name] = np.array(
+                [self.projection_conversion[i] for i in self.sensor_band_specs[sensor_name]])
+
+        self.patch_size = config.patch_size
+        self.pos_drop = nn.Dropout(p=config.pos_drop_rate)
+        self.nb_patch_length = int(config.img_size / self.patch_size)
+        self.num_patches = self.nb_patch_length**2
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, self.embed_dim), requires_grad=False)  # fixed sin-cos embedding
+
+        self.spectrum_projection = SpectrumAwareProjection(
+            spectrum_specs=config.spectrum_specs,
+            patch_size=self.patch_size,
+            embed_dim=self.embed_dim
+        )
+
+        pos_embed = get_2d_sincos_pos_embed(
+            self.pos_embed.shape[-1],
+            self.nb_patch_length,
+            cls_token=True,
+        )
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+        self.projection_scaler = 12
+        self.norm_layer = partial(nn.LayerNorm, eps=config.norm_layer_eps)
+
+        self.blocks = nn.ModuleList([
+            Block(self.embed_dim, config.num_heads, config.mlp_ratio, qkv_bias=config.qkv_bias, norm_layer=self.norm_layer)
+            for i in range(config.depth)])
+        self.norm = self.norm_layer(self.embed_dim)
+        self.global_pool = config.global_pool
+        if self.global_pool:
+            self.fc_norm = self.norm_layer(self.embed_dim)
+
+        # decoder specifics
+        self.decoder_embed = nn.Linear(self.embed_dim, self.decoder_embed_dim, bias=True)
+
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, self.decoder_embed_dim))
+
+        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, self.decoder_embed_dim), requires_grad=False)  # fixed sin-cos embedding
+        self.projection_scaler = 12
+
+        self.decoder_blocks = nn.ModuleList([
+            Block(self.decoder_embed_dim, config.decoder_num_heads, config.mlp_ratio, qkv_bias=True, norm_layer=self.norm_layer)
+            for i in range(config.decoder_depth)])
+
+        self.decoder_norm = self.norm_layer(self.decoder_embed_dim)
+        self.decoder_preds = torch.nn.ModuleList()
+        for band_idx in sorted(config.spectrum_specs, key=lambda key: config.spectrum_specs[key]['projection_idx']):
+            if ((config.spectrum_specs[band_idx]['projection_idx'] != -1) and (len(config.spectrum_specs[band_idx]['agg_projections']) == 0)):
+                self.decoder_preds.append(nn.Linear(self.decoder_embed_dim, self.patch_size**2, bias=True))
+    
+    def tensor_patchify(self, imgs):
+        """
+        imgs: (N, nb_bands, H, W)
+        x: (N, L, patch_size**2 *nb_bands)
+        """
+        p = self.patch_size
+        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0
+
+        h = w = imgs.shape[2] // p
+        x = imgs.reshape(shape=(imgs.shape[0], imgs.shape[1], h, p, w, p))
+        x = torch.einsum('nchpwq->nhwpqc', x)
+        x = x.reshape(shape=(imgs.shape[0], h, w, p, p, imgs.shape[1])).permute(0,1,2,5,3,4)
+        return x
+    
+    def forward_encoder(self, imgs, proj_indices, is_patchify, all_tokens):
+        if is_patchify:
+            img_patches = self.tensor_patchify(imgs)
+        else:
+            img_patches = imgs
+        B, nb_patch_h, nb_patch_w, nb_bands, _, _ = img_patches.shape
+        device = img_patches.device
+
+        img_spectrum_embeds = torch.zeros((B, nb_patch_h, nb_patch_w, nb_bands, self.embed_dim), device=device, dtype=self.model_dtype)
+
+        for projection_idx in torch.unbind(torch.unique(proj_indices)):
+            mask = (proj_indices==projection_idx)
+            img_spectrum_embeds[mask] = self.spectrum_projection(img_patches[mask], projection_idx) 
+
+        img_embeddings = self.projection_scaler*img_spectrum_embeds.mean(dim=3)
+        img_embeddings = img_embeddings.reshape(-1,nb_patch_h*nb_patch_w,self.embed_dim)
+
+        cls_tokens = self.cls_token.expand(
+            B, -1, -1
+        )
+        x = torch.cat((cls_tokens, img_embeddings), dim=1)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        if all_tokens:
+            return self.norm(x) # B, L, embed_dim (L=1+patch_size**2)
+
+        if self.global_pool:
+            x = x[:, 1:, :].mean(dim=1)
+            outcome = self.fc_norm(x)
+        else:
+            x = self.norm(x)
+            outcome = x[:, 0]
+
+        return outcome
+
+    def forward(self, imgs, is_patchify=True, sensor_type='S2', bands=None, proj_indices=None, all_tokens=False):
+        if proj_indices is None:
+            if bands is None:
+                assert sensor_type in self.sensor_band_specs.keys(), f"Sensor type {sensor_type} not recognized. Available types: {list(self.sensor_band_specs.keys())}. Otherwise provide bands."
+                proj_indices = self.sensor_projection_specs[sensor_type]
+            else:
+                proj_indices = []
+                for i in bands:
+                    if i in self.projection_conversion.keys():
+                        proj_indices.append(self.projection_conversion[i])
+                assert len(proj_indices) > 0, \
+                "No valid bands provided. Please check the bands to be aligned with the spectrum_specs definition \
+                (default version can be accessed at https://github.com/gsumbul/SMARTIES/blob/main/config/electromagnetic_spectrum.yaml)."
+            proj_indices = np.array(proj_indices)
+            proj_indices = torch.as_tensor(np.tile(proj_indices.reshape(
+                1,1,1,-1), (imgs.shape[0], self.nb_patch_length, self.nb_patch_length, 1)).astype(np.int32), device=imgs.device)
+            
+        return self.forward_encoder(imgs, proj_indices, is_patchify=is_patchify, all_tokens=all_tokens)
+    

spectrum_specs.yaml

@@ -0,0 +1,140 @@
+aerosol:
+  min_wavelength: 422
+  max_wavelength: 463
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B01 (aerosol)'
+  projection_idx: 0
+  agg_projections: []
+blue_1:
+  min_wavelength: 427
+  max_wavelength: 558
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B02 (blue)'
+  projection_idx: 1
+  agg_projections: []
+blue_2:
+  min_wavelength: 452
+  max_wavelength: 512
+  sensors: ['Landsat8-L2']
+  name: 'B2 (blue)' 
+  projection_idx: 18
+  agg_projections: [0, 1]
+blue_3:
+  min_wavelength: 430
+  max_wavelength: 545
+  sensors: ['RGB']
+  name: 'blue'
+  projection_idx: 2
+  agg_projections: []
+green_1:
+  min_wavelength: 466
+  max_wavelength: 620
+  sensors: ['RGB']
+  name: 'green'
+  projection_idx: 3
+  agg_projections: []
+green_2:
+  min_wavelength: 524
+  max_wavelength: 595
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B03 (green)'
+  projection_idx: 4
+  agg_projections: []
+red_1:
+  min_wavelength: 590
+  max_wavelength: 710
+  sensors: ['RGB']
+  name: 'red'
+  projection_idx: 5
+  agg_projections: []
+red_2:
+  min_wavelength: 634
+  max_wavelength: 696
+  sensors: ['SENTINEL2']
+  name: 'B04 (red)'
+  projection_idx: 6
+  agg_projections: []
+red_edge_1:
+  min_wavelength: 689
+  max_wavelength: 719
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B05 (red edge 1)'
+  projection_idx: 7
+  agg_projections: []
+red_edge_2:
+  min_wavelength: 726
+  max_wavelength: 755
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B06 (red edge 2)'
+  projection_idx: 8
+  agg_projections: []
+near_infrared_1:
+  min_wavelength: 728
+  max_wavelength: 938
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B08 (NIR 1)'
+  projection_idx: 9
+  agg_projections: []
+near_infrared_2:
+  min_wavelength: 761
+  max_wavelength: 802
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B07 (NIR 2)'
+  projection_idx: 10
+  agg_projections: []
+near_infrared_3:
+  min_wavelength: 843
+  max_wavelength: 886
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B8A (NIR 3)'
+  projection_idx: 11
+  agg_projections: []
+short_wave_infrared_1:
+  min_wavelength: 923
+  max_wavelength: 964
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B09 (SWIR water vapour)'
+  projection_idx: 12
+  agg_projections: []
+short_wave_infrared_2:
+  min_wavelength: 1345
+  max_wavelength: 1406
+  sensors: ['SENTINEL2-L1C']
+  name: 'B10 (SWIR circus)'
+  projection_idx: -1
+  agg_projections: []
+short_wave_infrared_3:
+  min_wavelength: 1516
+  max_wavelength: 1704
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B11 (SWIR 1)'
+  projection_idx: 13
+  agg_projections: []
+short_wave_infrared_4:
+  min_wavelength: 2002
+  max_wavelength: 2376
+  sensors: ['SENTINEL2-L1C', 'SENTINEL2-L2A']
+  name: 'B12 (SWIR 2)' 
+  projection_idx: 14
+  agg_projections: []
+thermal_infrared_1:
+  min_wavelength: 10600
+  max_wavelength: 11190
+  sensors: ['Landsat8-L2']
+  name: 'B10 (surface temperature)' 
+  projection_idx: 17
+  agg_projections: [14, 15]
+microwave_1:
+  min_wavelength: 5.5e7
+  max_wavelength: 5.6e7
+  sensors: ['SENTINEL1-GRD']
+  name: 'VV' 
+  projection_idx: 15
+  agg_projections: []
+microwave_2:
+  min_wavelength: 5.5e7
+  max_wavelength: 5.6e7
+  sensors: ['SENTINEL1-GRD']
+  name: 'VH' 
+  projection_idx: 16
+  agg_projections: []