Upload SatMAE++ transformers checkpoints with model card metadata

README.md

@@ -0,0 +1,119 @@
+---
+license: mit
+language:
+- en
+tags:
+- remote-sensing
+- earth-observation
+- self-supervised-learning
+- satellite
+- multispectral
+- feature-extraction
+- vision
+- satmae
+- satmae-pp
+- vit
+- mae
+- transformers
+library_name: transformers
+pipeline_tag: feature-extraction
+datasets:
+- fMoW
+---
+
+# SatMAE++ Transformers Models
+
+Hugging Face–compatible checkpoints converted from the official [SatMAE++](https://arxiv.org/abs/2403.05419) pretrained weights. Each subfolder is a standalone model repo layout (`config.json`, `model.safetensors`, preprocessor, and remote code) for feature extraction on FMoW satellite imagery.
+
+## Model Description
+
+These models are ViT-Large encoders pretrained with multi-scale masked autoencoding (SatMAE++) on [fMoW-RGB](https://github.com/fMoW/dataset) and [fMoW-Sentinel](https://github.com/sustainlab-group/SatMAE) imagery.
+
+This collection currently bundles **2 converted pretrain checkpoints**:
+
+- **FMoW-RGB:** vanilla ViT-L/16, 3-channel BGR, 224×224
+- **FMoW-Sentinel:** grouped-channel ViT-L/8, 10-band multispectral, 96×96
+
+All folders ship self-contained remote code (`modeling_satmae_pp.py`, processor, pipeline) and load with `trust_remote_code=True`.
+
+**Developed by:** [techmn / SatMAE++](https://github.com/techmn/satmae_pp)  
+**Converted for Hugging Face by:** BiliSakura  
+**License (weights):** MIT  
+**Original paper:** [Rethinking Transformers Pre-training for Multi-Spectral Satellite Imagery](https://arxiv.org/abs/2403.05419) (CVPR 2024)
+
+## Available checkpoints
+
+| Folder | Dataset | Encoder | Channels | Image | Patch | Legacy file |
+|--------|---------|---------|----------|-------|-------|-------------|
+| `satmae-pp-vit-large-patch16-fmow-rgb-pretrain` | FMoW-RGB | vanilla ViT | 3 (BGR) | 224 | 16 | `checkpoint_ViT-L_pretrain_fmow_rgb.pth` |
+| `satmae-pp-vit-large-patch8-fmow-sentinel-pretrain` | FMoW-Sentinel | group-channel ViT | 10 | 96 | 8 | `checkpoint_ViT-L_pretrain_fmow_sentinel.pth` |
+
+Legacy `.pth` filename mapping is in [`conversion_manifest.json`](conversion_manifest.json).
+
+## Usage
+
+```python
+from transformers import pipeline
+import numpy as np
+
+REPO = "BiliSakura/SATMAE-PP-transformers"
+SUBFOLDER = "satmae-pp-vit-large-patch8-fmow-sentinel-pretrain"
+
+pipe = pipeline(
+    task="satmae-pp-feature-extraction",
+    model=REPO,
+    trust_remote_code=True,
+    model_kwargs={"subfolder": SUBFOLDER},
+)
+
+# FMoW-Sentinel: 10 bands after dropping bands 0, 9, 10
+image = np.random.randint(0, 255, (96, 96, 10), dtype=np.uint8)
+features = pipe(image, pool=True, return_tensors=True)
+print(features.shape)  # [1, 1024]
+```
+
+FMoW-RGB (BGR channel order):
+
+```python
+SUBFOLDER = "satmae-pp-vit-large-patch16-fmow-rgb-pretrain"
+pipe = pipeline(
+    task="satmae-pp-feature-extraction",
+    model=REPO,
+    trust_remote_code=True,
+    model_kwargs={"subfolder": SUBFOLDER},
+)
+image = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
+features = pipe(image, pool=True, return_tensors=True)
+print(features.shape)  # [1, 1024]
+```
+
+Load components directly:
+
+```python
+from transformers import AutoModel, AutoImageProcessor
+
+model = AutoModel.from_pretrained(REPO, subfolder=SUBFOLDER, trust_remote_code=True)
+processor = AutoImageProcessor.from_pretrained(REPO, subfolder=SUBFOLDER, trust_remote_code=True)
+```
+
+## Normalization
+
+The bundled image processor applies per-channel FMoW mean/std normalization by default (`do_normalize=True`). FMoW-RGB models expect **BGR** channel order; the processor swaps RGB→BGR when `channel_order="bgr"`.
+
+For FMoW-Sentinel, inputs should be 10-band reflectance arrays (bands 0, 9, 10 already dropped) with FMoW-Sentinel statistics baked into the preprocessor config.
+
+## Dependencies
+
+- `transformers`, `torch`, `timm`, `safetensors`
+- `opencv-python` (multispectral resize with more than 4 channels)
+
+## Citation
+
+```bibtex
+@inproceedings{satmaepp2024rethinking,
+  title={Rethinking Transformers Pre-training for Multi-Spectral Satellite Imagery},
+  author={Mubashir Noman and Muzammal Naseer and Hisham Cholakkal and Rao Muhammad Anwar and Salman Khan and Fahad Shahbaz Khan},
+  year={2024},
+  booktitle={CVPR}
+}
+```

satmae-pp-vit-large-patch16-fmow-rgb-pretrain/config.json

@@ -0,0 +1,71 @@
+{
+  "architectures": [
+    "SatMAEppModel"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "channel_embed_dim": 256,
+  "channel_groups": [
+    [
+      0,
+      1,
+      2,
+      6
+    ],
+    [
+      3,
+      4,
+      5,
+      7
+    ],
+    [
+      8,
+      9
+    ]
+  ],
+  "channel_order": "bgr",
+  "checkpoint_stage": "pretrain",
+  "dataset": "fmow_rgb",
+  "dtype": "float32",
+  "encoder_type": "vanilla",
+  "global_pool": true,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 1024,
+  "id2label": {},
+  "image_mean": [
+    0.4182007312774658,
+    0.4214799106121063,
+    0.3991275727748871
+  ],
+  "image_size": 224,
+  "image_std": [
+    0.28774282336235046,
+    0.27541765570640564,
+    0.2764017581939697
+  ],
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "mlp_ratio": 4.0,
+  "model_type": "satmae_pp",
+  "num_attention_heads": 16,
+  "num_channels": 3,
+  "num_hidden_layers": 24,
+  "patch_size": 16,
+  "qkv_bias": true,
+  "transformers_version": "5.0.0",
+  "auto_map": {
+    "AutoConfig": "modeling_satmae_pp.SatMAEppConfig",
+    "AutoModel": "modeling_satmae_pp.SatMAEppModel",
+    "AutoModelForImageClassification": "modeling_satmae_pp.SatMAEppForImageClassification"
+  },
+  "custom_pipelines": {
+    "satmae-pp-feature-extraction": {
+      "impl": "pipeline_satmae_pp.SatMAEppImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  }
+}

satmae-pp-vit-large-patch16-fmow-rgb-pretrain/image_processing_satmae_pp.py

@@ -0,0 +1,189 @@
+# Copyright 2024 SatMAE++ Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for SatMAE++ models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class SatMAEppImageProcessor(BaseImageProcessor):
+    """
+    Image processor for SatMAE++ satellite encoders.
+
+    FMoW-RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models)
+    to swap the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 224, "width": 224}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["SatMAEppImageProcessor"]

satmae-pp-vit-large-patch16-fmow-rgb-pretrain/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc037acaba43f824bb725ffa30469084956c334fae10a110d7fd28ff62211c33
+size 1213234544

satmae-pp-vit-large-patch16-fmow-rgb-pretrain/modeling_satmae_pp.py

@@ -0,0 +1,285 @@
+# Copyright 2024 SatMAE++ Authors and The HuggingFace Inc. team.
+"""Self-contained SatMAE++ model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Optional
+
+import numpy as np
+import torch
+from timm.models.vision_transformer import Block, PatchEmbed
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+FMOW_RGB_MEAN = [0.4182007312774658, 0.4214799106121063, 0.3991275727748871]
+FMOW_RGB_STD = [0.28774282336235046, 0.27541765570640564, 0.2764017581939697]
+FMOW_SENTINEL_MEAN_10 = [
+    1184.3824625, 1120.77120066, 1136.26026392, 1263.73947144, 1645.40315151,
+    1846.87040806, 1762.59530783, 1972.62420416, 1732.16362238, 1247.91870117,
+]
+FMOW_SENTINEL_STD_10 = [
+    650.2842772, 712.12507725, 965.23119807, 948.9819932, 1108.06650639,
+    1258.36394548, 1233.1492281, 1364.38688993, 1310.36996126, 1087.6020813,
+]
+DEFAULT_CHANNEL_GROUPS = [[0, 1, 2, 6], [3, 4, 5, 7], [8, 9]]
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> np.ndarray:
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0).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: int, grid: np.ndarray) -> np.ndarray:
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    omega = np.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+
+class SatMAEppConfig(PreTrainedConfig):
+    model_type = "satmae_pp"
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        num_hidden_layers: int = 24,
+        num_attention_heads: int = 16,
+        intermediate_size: int | None = None,
+        hidden_act: str = "gelu",
+        hidden_dropout_prob: float = 0.0,
+        attention_probs_dropout_prob: float = 0.0,
+        initializer_range: float = 0.02,
+        layer_norm_eps: float = 1e-6,
+        image_size: int = 224,
+        patch_size: int = 16,
+        num_channels: int = 3,
+        qkv_bias: bool = True,
+        mlp_ratio: float = 4.0,
+        global_pool: bool = True,
+        encoder_type: str = "vanilla",
+        channel_embed_dim: int = 256,
+        channel_groups: list[list[int]] | None = None,
+        channel_order: str = "bgr",
+        dataset: str = "fmow_rgb",
+        checkpoint_stage: str = "finetune",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.qkv_bias = qkv_bias
+        self.mlp_ratio = mlp_ratio
+        self.global_pool = global_pool
+        self.encoder_type = encoder_type
+        self.channel_embed_dim = channel_embed_dim
+        self.channel_groups = channel_groups if channel_groups is not None else list(DEFAULT_CHANNEL_GROUPS)
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.checkpoint_stage = checkpoint_stage
+        self.num_labels = num_labels
+        self.intermediate_size = int(hidden_size * mlp_ratio) if intermediate_size is None else intermediate_size
+        if image_mean is None or image_std is None:
+            if encoder_type == "group_channel":
+                self.image_mean = FMOW_SENTINEL_MEAN_10
+                self.image_std = FMOW_SENTINEL_STD_10
+            else:
+                self.image_mean = FMOW_RGB_MEAN
+                self.image_std = FMOW_RGB_STD
+        else:
+            self.image_mean = image_mean
+            self.image_std = image_std
+
+
+class SatMAEppPreTrainedModel(PreTrainedModel):
+    config_class = SatMAEppConfig
+    config: SatMAEppConfig
+    base_model_prefix = "satmae_pp"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Block"]
+
+
+class SatMAEppModel(SatMAEppPreTrainedModel):
+    def __init__(self, config: SatMAEppConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        if config.encoder_type == "group_channel":
+            self._init_group_channel_encoder(config)
+        else:
+            self._init_vanilla_encoder(config)
+        self.post_init()
+
+    def _init_vanilla_encoder(self, config: SatMAEppConfig) -> None:
+        image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0]
+        norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
+        self.patch_embed = PatchEmbed(image_size, config.patch_size, config.num_channels, config.hidden_size)
+        self.num_patches = self.patch_embed.num_patches
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, config.hidden_size))
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.num_patches**0.5), cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+        self.blocks = nn.ModuleList([
+            Block(config.hidden_size, config.num_attention_heads, config.mlp_ratio, qkv_bias=config.qkv_bias, norm_layer=norm_layer)
+            for _ in range(config.num_hidden_layers)
+        ])
+        self.global_pool = config.global_pool
+        if self.global_pool:
+            self.fc_norm = norm_layer(config.hidden_size)
+            self.norm = None
+        else:
+            self.fc_norm = None
+            self.norm = norm_layer(config.hidden_size)
+
+    def _init_group_channel_encoder(self, config: SatMAEppConfig) -> None:
+        image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0]
+        norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
+        self.channel_groups = tuple(tuple(group) for group in config.channel_groups)
+        self.patch_embed = nn.ModuleList([
+            PatchEmbed(image_size, config.patch_size, len(group), config.hidden_size)
+            for group in self.channel_groups
+        ])
+        self.num_patches = self.patch_embed[0].num_patches
+        pos_dim = config.hidden_size - config.channel_embed_dim
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, pos_dim))
+        pos_embed = get_2d_sincos_pos_embed(pos_dim, int(self.num_patches**0.5), cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+        num_groups = len(self.channel_groups)
+        self.channel_embed = nn.Parameter(torch.zeros(1, num_groups, config.channel_embed_dim))
+        chan_embed = get_1d_sincos_pos_embed_from_grid(self.channel_embed.shape[-1], np.arange(num_groups))
+        self.channel_embed.data.copy_(torch.from_numpy(chan_embed).float().unsqueeze(0))
+        self.channel_cls_embed = nn.Parameter(torch.zeros(1, 1, config.channel_embed_dim))
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+        self.blocks = nn.ModuleList([
+            Block(config.hidden_size, config.num_attention_heads, config.mlp_ratio, qkv_bias=config.qkv_bias, norm_layer=norm_layer)
+            for _ in range(config.num_hidden_layers)
+        ])
+        self.global_pool = config.global_pool
+        if self.global_pool:
+            self.fc_norm = norm_layer(config.hidden_size)
+            self.norm = None
+        else:
+            self.fc_norm = None
+            self.norm = norm_layer(config.hidden_size)
+
+    def _forward_vanilla(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size = pixel_values.shape[0]
+        patch_tokens = self.patch_embed(pixel_values)
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        hidden_states = torch.cat((cls_tokens, patch_tokens), dim=1) + self.pos_embed
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        if self.global_pool:
+            pooled_output = self.fc_norm(hidden_states[:, 1:, :].mean(dim=1))
+        else:
+            hidden_states = self.norm(hidden_states)
+            pooled_output = hidden_states[:, 0]
+        return hidden_states, pooled_output
+
+    def _forward_group_channel(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size = pixel_values.shape[0]
+        group_tokens = [self.patch_embed[i](pixel_values[:, group, :, :]) for i, group in enumerate(self.channel_groups)]
+        hidden_states = torch.stack(group_tokens, dim=1)
+        channel_embed = self.channel_embed.unsqueeze(2).expand(-1, -1, self.pos_embed[:, 1:, :].shape[1], -1)
+        pos_embed = self.pos_embed[:, 1:, :].unsqueeze(1).expand(-1, channel_embed.shape[1], -1, -1)
+        hidden_states = (hidden_states + torch.cat((pos_embed, channel_embed), dim=-1)).view(batch_size, -1, hidden_states.shape[-1])
+        cls_pos_channel = torch.cat((self.pos_embed[:, :1, :], self.channel_cls_embed), dim=-1)
+        hidden_states = torch.cat((cls_pos_channel + self.cls_token.expand(batch_size, -1, -1), hidden_states), dim=1)
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        if self.global_pool:
+            pooled_output = self.fc_norm(hidden_states[:, 1:, :].mean(dim=1))
+        else:
+            hidden_states = self.norm(hidden_states)
+            pooled_output = hidden_states[:, 0]
+        return hidden_states, pooled_output
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        if self.config.encoder_type == "group_channel":
+            last_hidden_state, pooled_output = self._forward_group_channel(pixel_values)
+        else:
+            last_hidden_state, pooled_output = self._forward_vanilla(pixel_values)
+        if not self.add_pooling_layer:
+            pooled_output = None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class SatMAEppForImageClassification(SatMAEppPreTrainedModel):
+    def __init__(self, config: SatMAEppConfig):
+        super().__init__(config)
+        self.satmae_pp = SatMAEppModel(config, add_pooling_layer=True)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.satmae_pp(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
+
+
+__all__ = [
+    "SatMAEppConfig",
+    "SatMAEppForImageClassification",
+    "SatMAEppModel",
+    "SatMAEppPreTrainedModel",
+]

satmae-pp-vit-large-patch16-fmow-rgb-pretrain/pipeline_satmae_pp.py

@@ -0,0 +1,79 @@
+# Copyright 2024 SatMAE++ Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""SatMAE++ image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class SatMAEppImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    SatMAE++ image feature extraction pipeline.
+
+    This pipeline wraps [`SatMAEppModel`] for RGB and multispectral satellite feature extraction.
+
+    Example:
+
+    ```python
+    >>> from transformers import pipeline
+
+    >>> pipe = pipeline(
+    ...     task="image-feature-extraction",
+    ...     model="/path/to/satmae-pp-vit-large-patch16-fmow-rgb-finetune",
+    ...     trust_remote_code=True,
+    ... )
+    >>> features = pipe(image_array, pool=True, return_tensors=True)
+    ```
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["SatMAEppImageFeatureExtractionPipeline"]

satmae-pp-vit-large-patch16-fmow-rgb-pretrain/preprocessor_config.json

@@ -0,0 +1,25 @@
+{
+  "image_processor_type": "SatMAEppImageProcessor",
+  "size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_resize": true,
+  "do_rescale": false,
+  "do_normalize": true,
+  "do_convert_rgb": false,
+  "channel_order": "bgr",
+  "image_mean": [
+    0.4182007312774658,
+    0.4214799106121063,
+    0.3991275727748871
+  ],
+  "image_std": [
+    0.28774282336235046,
+    0.27541765570640564,
+    0.2764017581939697
+  ],
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_satmae_pp.SatMAEppImageProcessor"
+  }
+}

satmae-pp-vit-large-patch8-fmow-sentinel-pretrain/config.json

@@ -0,0 +1,85 @@
+{
+  "architectures": [
+    "SatMAEppModel"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "channel_embed_dim": 256,
+  "channel_groups": [
+    [
+      0,
+      1,
+      2,
+      6
+    ],
+    [
+      3,
+      4,
+      5,
+      7
+    ],
+    [
+      8,
+      9
+    ]
+  ],
+  "channel_order": "rgb",
+  "checkpoint_stage": "pretrain",
+  "dataset": "fmow_sentinel",
+  "dtype": "float32",
+  "encoder_type": "group_channel",
+  "global_pool": true,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 1024,
+  "id2label": {},
+  "image_mean": [
+    1184.3824625,
+    1120.77120066,
+    1136.26026392,
+    1263.73947144,
+    1645.40315151,
+    1846.87040806,
+    1762.59530783,
+    1972.62420416,
+    1732.16362238,
+    1247.91870117
+  ],
+  "image_size": 96,
+  "image_std": [
+    650.2842772,
+    712.12507725,
+    965.23119807,
+    948.9819932,
+    1108.06650639,
+    1258.36394548,
+    1233.1492281,
+    1364.38688993,
+    1310.36996126,
+    1087.6020813
+  ],
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "mlp_ratio": 4.0,
+  "model_type": "satmae_pp",
+  "num_attention_heads": 16,
+  "num_channels": 10,
+  "num_hidden_layers": 24,
+  "patch_size": 8,
+  "qkv_bias": true,
+  "transformers_version": "5.0.0",
+  "auto_map": {
+    "AutoConfig": "modeling_satmae_pp.SatMAEppConfig",
+    "AutoModel": "modeling_satmae_pp.SatMAEppModel",
+    "AutoModelForImageClassification": "modeling_satmae_pp.SatMAEppForImageClassification"
+  },
+  "custom_pipelines": {
+    "satmae-pp-feature-extraction": {
+      "impl": "pipeline_satmae_pp.SatMAEppImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  }
+}

satmae-pp-vit-large-patch8-fmow-sentinel-pretrain/image_processing_satmae_pp.py

@@ -0,0 +1,189 @@
+# Copyright 2024 SatMAE++ Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for SatMAE++ models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class SatMAEppImageProcessor(BaseImageProcessor):
+    """
+    Image processor for SatMAE++ satellite encoders.
+
+    FMoW-RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models)
+    to swap the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 224, "width": 224}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["SatMAEppImageProcessor"]

satmae-pp-vit-large-patch8-fmow-sentinel-pretrain/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fb26efc283ba33e345aef826afffbe4e6b0e3f1eb00f1a3852e36ca5ba3e8aa
+size 1212361656

satmae-pp-vit-large-patch8-fmow-sentinel-pretrain/modeling_satmae_pp.py

@@ -0,0 +1,285 @@
+# Copyright 2024 SatMAE++ Authors and The HuggingFace Inc. team.
+"""Self-contained SatMAE++ model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Optional
+
+import numpy as np
+import torch
+from timm.models.vision_transformer import Block, PatchEmbed
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+FMOW_RGB_MEAN = [0.4182007312774658, 0.4214799106121063, 0.3991275727748871]
+FMOW_RGB_STD = [0.28774282336235046, 0.27541765570640564, 0.2764017581939697]
+FMOW_SENTINEL_MEAN_10 = [
+    1184.3824625, 1120.77120066, 1136.26026392, 1263.73947144, 1645.40315151,
+    1846.87040806, 1762.59530783, 1972.62420416, 1732.16362238, 1247.91870117,
+]
+FMOW_SENTINEL_STD_10 = [
+    650.2842772, 712.12507725, 965.23119807, 948.9819932, 1108.06650639,
+    1258.36394548, 1233.1492281, 1364.38688993, 1310.36996126, 1087.6020813,
+]
+DEFAULT_CHANNEL_GROUPS = [[0, 1, 2, 6], [3, 4, 5, 7], [8, 9]]
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> np.ndarray:
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0).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: int, grid: np.ndarray) -> np.ndarray:
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    omega = np.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+
+class SatMAEppConfig(PreTrainedConfig):
+    model_type = "satmae_pp"
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        num_hidden_layers: int = 24,
+        num_attention_heads: int = 16,
+        intermediate_size: int | None = None,
+        hidden_act: str = "gelu",
+        hidden_dropout_prob: float = 0.0,
+        attention_probs_dropout_prob: float = 0.0,
+        initializer_range: float = 0.02,
+        layer_norm_eps: float = 1e-6,
+        image_size: int = 224,
+        patch_size: int = 16,
+        num_channels: int = 3,
+        qkv_bias: bool = True,
+        mlp_ratio: float = 4.0,
+        global_pool: bool = True,
+        encoder_type: str = "vanilla",
+        channel_embed_dim: int = 256,
+        channel_groups: list[list[int]] | None = None,
+        channel_order: str = "bgr",
+        dataset: str = "fmow_rgb",
+        checkpoint_stage: str = "finetune",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.qkv_bias = qkv_bias
+        self.mlp_ratio = mlp_ratio
+        self.global_pool = global_pool
+        self.encoder_type = encoder_type
+        self.channel_embed_dim = channel_embed_dim
+        self.channel_groups = channel_groups if channel_groups is not None else list(DEFAULT_CHANNEL_GROUPS)
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.checkpoint_stage = checkpoint_stage
+        self.num_labels = num_labels
+        self.intermediate_size = int(hidden_size * mlp_ratio) if intermediate_size is None else intermediate_size
+        if image_mean is None or image_std is None:
+            if encoder_type == "group_channel":
+                self.image_mean = FMOW_SENTINEL_MEAN_10
+                self.image_std = FMOW_SENTINEL_STD_10
+            else:
+                self.image_mean = FMOW_RGB_MEAN
+                self.image_std = FMOW_RGB_STD
+        else:
+            self.image_mean = image_mean
+            self.image_std = image_std
+
+
+class SatMAEppPreTrainedModel(PreTrainedModel):
+    config_class = SatMAEppConfig
+    config: SatMAEppConfig
+    base_model_prefix = "satmae_pp"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Block"]
+
+
+class SatMAEppModel(SatMAEppPreTrainedModel):
+    def __init__(self, config: SatMAEppConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        if config.encoder_type == "group_channel":
+            self._init_group_channel_encoder(config)
+        else:
+            self._init_vanilla_encoder(config)
+        self.post_init()
+
+    def _init_vanilla_encoder(self, config: SatMAEppConfig) -> None:
+        image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0]
+        norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
+        self.patch_embed = PatchEmbed(image_size, config.patch_size, config.num_channels, config.hidden_size)
+        self.num_patches = self.patch_embed.num_patches
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, config.hidden_size))
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.num_patches**0.5), cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+        self.blocks = nn.ModuleList([
+            Block(config.hidden_size, config.num_attention_heads, config.mlp_ratio, qkv_bias=config.qkv_bias, norm_layer=norm_layer)
+            for _ in range(config.num_hidden_layers)
+        ])
+        self.global_pool = config.global_pool
+        if self.global_pool:
+            self.fc_norm = norm_layer(config.hidden_size)
+            self.norm = None
+        else:
+            self.fc_norm = None
+            self.norm = norm_layer(config.hidden_size)
+
+    def _init_group_channel_encoder(self, config: SatMAEppConfig) -> None:
+        image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0]
+        norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
+        self.channel_groups = tuple(tuple(group) for group in config.channel_groups)
+        self.patch_embed = nn.ModuleList([
+            PatchEmbed(image_size, config.patch_size, len(group), config.hidden_size)
+            for group in self.channel_groups
+        ])
+        self.num_patches = self.patch_embed[0].num_patches
+        pos_dim = config.hidden_size - config.channel_embed_dim
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, pos_dim))
+        pos_embed = get_2d_sincos_pos_embed(pos_dim, int(self.num_patches**0.5), cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+        num_groups = len(self.channel_groups)
+        self.channel_embed = nn.Parameter(torch.zeros(1, num_groups, config.channel_embed_dim))
+        chan_embed = get_1d_sincos_pos_embed_from_grid(self.channel_embed.shape[-1], np.arange(num_groups))
+        self.channel_embed.data.copy_(torch.from_numpy(chan_embed).float().unsqueeze(0))
+        self.channel_cls_embed = nn.Parameter(torch.zeros(1, 1, config.channel_embed_dim))
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+        self.blocks = nn.ModuleList([
+            Block(config.hidden_size, config.num_attention_heads, config.mlp_ratio, qkv_bias=config.qkv_bias, norm_layer=norm_layer)
+            for _ in range(config.num_hidden_layers)
+        ])
+        self.global_pool = config.global_pool
+        if self.global_pool:
+            self.fc_norm = norm_layer(config.hidden_size)
+            self.norm = None
+        else:
+            self.fc_norm = None
+            self.norm = norm_layer(config.hidden_size)
+
+    def _forward_vanilla(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size = pixel_values.shape[0]
+        patch_tokens = self.patch_embed(pixel_values)
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        hidden_states = torch.cat((cls_tokens, patch_tokens), dim=1) + self.pos_embed
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        if self.global_pool:
+            pooled_output = self.fc_norm(hidden_states[:, 1:, :].mean(dim=1))
+        else:
+            hidden_states = self.norm(hidden_states)
+            pooled_output = hidden_states[:, 0]
+        return hidden_states, pooled_output
+
+    def _forward_group_channel(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size = pixel_values.shape[0]
+        group_tokens = [self.patch_embed[i](pixel_values[:, group, :, :]) for i, group in enumerate(self.channel_groups)]
+        hidden_states = torch.stack(group_tokens, dim=1)
+        channel_embed = self.channel_embed.unsqueeze(2).expand(-1, -1, self.pos_embed[:, 1:, :].shape[1], -1)
+        pos_embed = self.pos_embed[:, 1:, :].unsqueeze(1).expand(-1, channel_embed.shape[1], -1, -1)
+        hidden_states = (hidden_states + torch.cat((pos_embed, channel_embed), dim=-1)).view(batch_size, -1, hidden_states.shape[-1])
+        cls_pos_channel = torch.cat((self.pos_embed[:, :1, :], self.channel_cls_embed), dim=-1)
+        hidden_states = torch.cat((cls_pos_channel + self.cls_token.expand(batch_size, -1, -1), hidden_states), dim=1)
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        if self.global_pool:
+            pooled_output = self.fc_norm(hidden_states[:, 1:, :].mean(dim=1))
+        else:
+            hidden_states = self.norm(hidden_states)
+            pooled_output = hidden_states[:, 0]
+        return hidden_states, pooled_output
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        if self.config.encoder_type == "group_channel":
+            last_hidden_state, pooled_output = self._forward_group_channel(pixel_values)
+        else:
+            last_hidden_state, pooled_output = self._forward_vanilla(pixel_values)
+        if not self.add_pooling_layer:
+            pooled_output = None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class SatMAEppForImageClassification(SatMAEppPreTrainedModel):
+    def __init__(self, config: SatMAEppConfig):
+        super().__init__(config)
+        self.satmae_pp = SatMAEppModel(config, add_pooling_layer=True)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.satmae_pp(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
+
+
+__all__ = [
+    "SatMAEppConfig",
+    "SatMAEppForImageClassification",
+    "SatMAEppModel",
+    "SatMAEppPreTrainedModel",
+]

satmae-pp-vit-large-patch8-fmow-sentinel-pretrain/pipeline_satmae_pp.py

@@ -0,0 +1,79 @@
+# Copyright 2024 SatMAE++ Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""SatMAE++ image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class SatMAEppImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    SatMAE++ image feature extraction pipeline.
+
+    This pipeline wraps [`SatMAEppModel`] for RGB and multispectral satellite feature extraction.
+
+    Example:
+
+    ```python
+    >>> from transformers import pipeline
+
+    >>> pipe = pipeline(
+    ...     task="image-feature-extraction",
+    ...     model="/path/to/satmae-pp-vit-large-patch16-fmow-rgb-finetune",
+    ...     trust_remote_code=True,
+    ... )
+    >>> features = pipe(image_array, pool=True, return_tensors=True)
+    ```
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["SatMAEppImageFeatureExtractionPipeline"]

satmae-pp-vit-large-patch8-fmow-sentinel-pretrain/preprocessor_config.json

@@ -0,0 +1,39 @@
+{
+  "image_processor_type": "SatMAEppImageProcessor",
+  "size": {
+    "height": 96,
+    "width": 96
+  },
+  "do_resize": true,
+  "do_rescale": false,
+  "do_normalize": true,
+  "do_convert_rgb": false,
+  "channel_order": "rgb",
+  "image_mean": [
+    1184.3824625,
+    1120.77120066,
+    1136.26026392,
+    1263.73947144,
+    1645.40315151,
+    1846.87040806,
+    1762.59530783,
+    1972.62420416,
+    1732.16362238,
+    1247.91870117
+  ],
+  "image_std": [
+    650.2842772,
+    712.12507725,
+    965.23119807,
+    948.9819932,
+    1108.06650639,
+    1258.36394548,
+    1233.1492281,
+    1364.38688993,
+    1310.36996126,
+    1087.6020813
+  ],
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_satmae_pp.SatMAEppImageProcessor"
+  }
+}