Upload 12 files

README.md

@@ -0,0 +1,59 @@
+# DOFA Transformers Models
+
+Self-contained HuggingFace model checkpoints for [DOFA](https://arxiv.org/abs/2403.15356).
+
+Each checkpoint subfolder ships remote code for model, processor, and pipeline loading with `trust_remote_code=True`.
+
+Sentinel-2 9-band defaults (`default_wavelengths`, `default_image_mean`, `default_image_std`) are baked into `config.json` and `preprocessor_config.json`.
+
+## Available checkpoints
+
+| Folder | Hidden size | Layers | Heads |
+|--------|-------------|--------|-------|
+| `dofa-base-patch16-224/` | 768 | 12 | 12 |
+| `dofa-large-patch16-224/` | 1024 | 24 | 16 |
+
+## Usage
+
+```python
+from transformers import pipeline
+
+MODEL = "/path/to/DOFA-transformers/dofa-base-patch16-224"
+
+pipe = pipeline(
+    task="dofa-feature-extraction",
+    model=MODEL,
+    trust_remote_code=True,
+)
+
+features = pipe(image_array, pool=True, return_tensors=True)
+```
+
+Override Sentinel-2 defaults for other sensors:
+
+```python
+features = pipe(
+    image_array,
+    wavelengths=[...],
+    image_mean=[...],
+    image_std=[...],
+    pool=True,
+    return_tensors=True,
+)
+```
+
+## Test CLI
+
+```bash
+conda activate rsgen
+python test_dofa.py
+python test_dofa.py --model dofa-large-patch16-224
+python test_dofa.py --model dofa-base-patch16-224 --no-pool
+```
+
+## Dependencies
+
+- `transformers`
+- `timm`
+- `torch`
+- `opencv-python` (only when resizing inputs with more than 4 channels)

dofa-base-patch16-224/config.json

@@ -0,0 +1,73 @@
+{
+  "architectures": [
+    "DOFAModel"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "default_wavelengths": [
+    0.665,
+    0.56,
+    0.49,
+    0.705,
+    0.74,
+    0.783,
+    0.842,
+    1.61,
+    2.19
+  ],
+  "dtype": "float32",
+  "global_pool": true,
+  "head_dropout": 0.0,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 768,
+  "id2label": {},
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "mlp_ratio": 4.0,
+  "model_type": "dofa",
+  "num_attention_heads": 12,
+  "num_channels": 9,
+  "num_hidden_layers": 12,
+  "patch_size": 16,
+  "qkv_bias": true,
+  "transformers_version": "5.0.0",
+  "wv_planes": 128,
+  "auto_map": {
+    "AutoConfig": "modeling_dofa.DOFAConfig",
+    "AutoModel": "modeling_dofa.DOFAModel",
+    "AutoModelForImageClassification": "modeling_dofa.DOFAForImageClassification"
+  },
+  "custom_pipelines": {
+    "dofa-feature-extraction": {
+      "impl": "pipeline_dofa.DOFAImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "default_image_mean": [
+    114.11,
+    114.82,
+    126.64,
+    84.34,
+    97.85,
+    103.94,
+    101.44,
+    72.33,
+    56.67
+  ],
+  "default_image_std": [
+    77.84,
+    69.97,
+    67.42,
+    64.57,
+    61.73,
+    61.34,
+    60.3,
+    47.89,
+    42.56
+  ]
+}

dofa-base-patch16-224/image_processing_dofa.py

@@ -0,0 +1,180 @@
+# Copyright 2024 The DOFA Authors and The HuggingFace Inc. team.
+"""Image processor for DOFA."""
+
+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
+
+
+class DOFAImageProcessor(BaseImageProcessor):
+    model_input_names = ["pixel_values", "wavelengths"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.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,
+        default_wavelengths: Optional[list[float]] = None,
+        **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.default_wavelengths = default_wavelengths
+
+    @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,
+        wavelengths: Optional[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,
+    ):
+        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
+        wavelengths = wavelengths if wavelengths is not None else self.default_wavelengths
+
+        if wavelengths is None:
+            raise ValueError(
+                "DOFA requires `wavelengths` (one value per channel, in micrometers). "
+                "Pass them to the processor call or set `default_wavelengths` in the processor config."
+            )
+
+        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. "
+                "Provide modality-specific statistics for your input bands."
+            )
+
+        if do_normalize and (
+            len(image_mean) != len(wavelengths) or len(image_std) != len(wavelengths)
+        ):
+            raise ValueError(
+                f"`image_mean` and `image_std` must match the number of wavelengths/channels "
+                f"({len(wavelengths)}), got mean={len(image_mean)} std={len(image_std)}."
+            )
+
+        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
+
+            num_channels = image.shape[-1] if input_data_format == ChannelDimension.LAST else image.shape[0]
+
+            if do_resize:
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(image, size=size, 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,
+            "wavelengths": np.asarray(wavelengths, dtype=np.float32),
+        }
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["DOFAImageProcessor"]

dofa-base-patch16-224/modeling_dofa.py

@@ -0,0 +1,368 @@
+# Copyright 2024 The DOFA Authors and The HuggingFace Inc. team.
+"""Self-contained DOFA model, config, and dynamic patch embedding."""
+
+from functools import partial
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+from timm.models.vision_transformer import Block
+
+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__)
+
+
+class DOFAConfig(PreTrainedConfig):
+    model_type = "dofa"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=None,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.0,
+        attention_probs_dropout_prob=0.0,
+        initializer_range=0.02,
+        layer_norm_eps=1e-6,
+        image_size=224,
+        patch_size=16,
+        num_channels=3,
+        qkv_bias=True,
+        wv_planes=128,
+        mlp_ratio=4.0,
+        global_pool=True,
+        default_wavelengths=None,
+        default_image_mean=None,
+        default_image_std=None,
+        head_dropout=0.0,
+        num_labels=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.wv_planes = wv_planes
+        self.mlp_ratio = mlp_ratio
+        self.global_pool = global_pool
+        self.default_wavelengths = default_wavelengths
+        self.default_image_mean = default_image_mean
+        self.default_image_std = default_image_std
+        self.head_dropout = head_dropout
+        self.num_labels = num_labels
+
+        if intermediate_size is None:
+            self.intermediate_size = int(hidden_size * mlp_ratio)
+        else:
+            self.intermediate_size = intermediate_size
+
+
+def get_1d_sincos_pos_embed_from_grid_torch(embed_dim, pos):
+    assert embed_dim % 2 == 0
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=pos.device)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+
+    pos = pos.reshape(-1)
+    out = torch.einsum("m,d->md", pos, omega)
+
+    emb_sin = torch.sin(out)
+    emb_cos = torch.cos(out)
+
+    return torch.cat([emb_sin, emb_cos], dim=1)
+
+
+class TransformerWeightGenerator(nn.Module):
+    def __init__(self, input_dim, output_dim, embed_dim, num_heads=4, num_layers=1):
+        super().__init__()
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=input_dim,
+            nhead=num_heads,
+            activation="gelu",
+            norm_first=False,
+            batch_first=False,
+            dropout=False,
+        )
+        self.transformer_encoder = nn.TransformerEncoder(
+            encoder_layer, num_layers=num_layers, enable_nested_tensor=False
+        )
+
+        self.fc_weight = nn.Linear(input_dim, output_dim)
+        self.fc_bias = nn.Linear(input_dim, embed_dim)
+        self.wt_num = 128
+        self.weight_tokens = nn.Parameter(torch.empty([self.wt_num, input_dim]))
+        self.bias_token = nn.Parameter(torch.empty([1, input_dim]))
+
+        torch.nn.init.normal_(self.weight_tokens, std=0.02)
+        torch.nn.init.normal_(self.bias_token, std=0.02)
+
+    def forward(self, x):
+        pos_wave = x
+        x = torch.cat([self.weight_tokens, pos_wave], dim=0)
+        x = torch.cat([x, self.bias_token], dim=0)
+        transformer_output = self.transformer_encoder(x)
+        weights = self.fc_weight(transformer_output[self.wt_num : -1] + pos_wave)
+        bias = self.fc_bias(transformer_output[-1])
+        return weights, bias
+
+
+class FCResLayer(nn.Module):
+    def __init__(self, linear_size=128):
+        super().__init__()
+        self.nonlin1 = nn.ReLU(inplace=True)
+        self.nonlin2 = nn.ReLU(inplace=True)
+        self.w1 = nn.Linear(linear_size, linear_size)
+        self.w2 = nn.Linear(linear_size, linear_size)
+
+    def forward(self, x):
+        y = self.w1(x)
+        y = self.nonlin1(y)
+        y = self.w2(y)
+        y = self.nonlin2(y)
+        return x + y
+
+
+class DOFADynamicPatchEmbed(nn.Module):
+    def __init__(self, wv_planes, inter_dim=128, kernel_size=16, embed_dim=768):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.wv_planes = wv_planes
+        self.embed_dim = embed_dim
+        self._num_kernel = self.kernel_size * self.kernel_size * self.embed_dim
+        self.inter_dim = inter_dim
+        self.patch_size = (kernel_size, kernel_size)
+
+        self.weight_generator = TransformerWeightGenerator(wv_planes, self._num_kernel, embed_dim)
+        self.scaler = 0.01
+        self.fclayer = FCResLayer(wv_planes)
+        self.weight_generator.apply(self._weight_init)
+        self.fclayer.apply(self._weight_init)
+
+    def _weight_init(self, module):
+        if isinstance(module, nn.Linear):
+            init.xavier_uniform_(module.weight)
+            module.bias.data.fill_(0.01)
+
+    def forward(self, pixel_values, wavelengths):
+        inplanes = wavelengths.size(0)
+        waves = get_1d_sincos_pos_embed_from_grid_torch(self.wv_planes, wavelengths * 1000)
+        waves = self.fclayer(waves)
+        weight, bias = self.weight_generator(waves)
+
+        dynamic_weight = weight.view(inplanes, self.kernel_size, self.kernel_size, self.embed_dim)
+        dynamic_weight = dynamic_weight.permute([3, 0, 1, 2])
+
+        if bias is not None:
+            bias = bias.view([self.embed_dim]) * self.scaler
+
+        weights = dynamic_weight * self.scaler
+        dynamic_out = F.conv2d(
+            pixel_values, weights, bias=bias, stride=self.kernel_size, padding=1, dilation=1
+        )
+        return dynamic_out.flatten(2).transpose(1, 2), waves
+
+
+def _prepare_wavelengths(wavelengths, pixel_values, default_wavelengths=None):
+    if wavelengths is None:
+        if default_wavelengths is None:
+            raise ValueError(
+                "DOFA requires per-channel wavelengths. Pass `wavelengths` to the model or image processor, "
+                "or set `default_wavelengths` in the model config."
+            )
+        wavelengths = default_wavelengths
+
+    if not torch.is_tensor(wavelengths):
+        wavelengths = torch.tensor(wavelengths, dtype=torch.float32, device=pixel_values.device)
+    else:
+        wavelengths = wavelengths.to(device=pixel_values.device, dtype=torch.float32)
+
+    if wavelengths.dim() == 2 and wavelengths.shape[0] == 1:
+        wavelengths = wavelengths.squeeze(0)
+    elif wavelengths.dim() == 2 and wavelengths.shape[0] == pixel_values.shape[0]:
+        if not torch.allclose(wavelengths[0], wavelengths):
+            raise ValueError("DOFA currently expects identical wavelengths for all items in a batch.")
+        wavelengths = wavelengths[0]
+
+    if wavelengths.dim() != 1:
+        raise ValueError("`wavelengths` must be a 1D tensor/list with one value per input channel.")
+
+    if wavelengths.shape[0] != pixel_values.shape[1]:
+        raise ValueError(
+            f"Expected {pixel_values.shape[1]} wavelengths for {pixel_values.shape[1]} channels, "
+            f"but got {wavelengths.shape[0]}."
+        )
+
+    return wavelengths
+
+
+class DOFAPreTrainedModel(PreTrainedModel):
+    config_class = DOFAConfig
+    base_model_prefix = "dofa"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Block"]
+    _supports_sdpa = False
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, DOFAModel):
+            if hasattr(module, "pos_embed"):
+                nn.init.trunc_normal_(module.pos_embed, std=self.config.initializer_range)
+            if hasattr(module, "cls_token"):
+                nn.init.trunc_normal_(module.cls_token, std=self.config.initializer_range)
+
+
+class DOFAModel(DOFAPreTrainedModel):
+    def __init__(self, config: DOFAConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+
+        image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0]
+        self.patch_embed = DOFADynamicPatchEmbed(
+            wv_planes=config.wv_planes,
+            inter_dim=config.wv_planes,
+            kernel_size=config.patch_size,
+            embed_dim=config.hidden_size,
+        )
+        self.num_patches = (image_size // config.patch_size) ** 2
+        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), requires_grad=False
+        )
+
+        norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
+        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)
+
+        self.add_pooling_layer = add_pooling_layer
+        self.post_init()
+
+    def forward_features(self, pixel_values, wavelengths):
+        patch_tokens, _ = self.patch_embed(pixel_values, wavelengths)
+        patch_tokens = patch_tokens + self.pos_embed[:, 1:, :]
+
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(pixel_values.shape[0], -1, -1)
+        hidden_states = torch.cat((cls_tokens, patch_tokens), 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:
+            pooled_output = self.norm(hidden_states)[:, 0]
+
+        return hidden_states, pooled_output
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        wavelengths: Optional[Union[torch.Tensor, list]] = 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
+
+        wavelengths = _prepare_wavelengths(wavelengths, pixel_values, self.config.default_wavelengths)
+        last_hidden_state, pooled_output = self.forward_features(pixel_values, wavelengths)
+
+        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 DOFAForImageClassification(DOFAPreTrainedModel):
+    def __init__(self, config: DOFAConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.dofa = DOFAModel(config, add_pooling_layer=True)
+        self.head_drop = nn.Dropout(config.head_dropout)
+        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,
+        wavelengths: Optional[Union[torch.Tensor, list]] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.dofa(
+            pixel_values=pixel_values,
+            wavelengths=wavelengths,
+            return_dict=True,
+            **kwargs,
+        )
+        pooled_output = self.head_drop(outputs.pooler_output)
+        logits = self.classifier(pooled_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__ = ["DOFAConfig", "DOFAForImageClassification", "DOFAModel", "DOFAPreTrainedModel"]

dofa-base-patch16-224/pipeline_dofa.py

@@ -0,0 +1,95 @@
+# Copyright 2024 The DOFA Authors and The HuggingFace Inc. team.
+"""Custom DOFA feature extraction pipeline for trust_remote_code loading."""
+
+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),
+    """
+        wavelengths (`list[float]`, *optional*):
+            Per-channel wavelengths in micrometers. Defaults to Sentinel-2 bands from config.
+        image_mean (`list[float]`, *optional*):
+            Per-channel normalization mean values.
+        image_std (`list[float]`, *optional*):
+            Per-channel normalization standard deviation values.
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether to return pooled features instead of sequence hidden states.
+    """,
+)
+class DOFAImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    DOFA wavelength-conditioned multispectral feature extraction pipeline.
+
+    This pipeline can be loaded from [`pipeline`] using the task identifier `"dofa-feature-extraction"`
+    with `trust_remote_code=True`.
+
+    Example:
+
+    ```python
+    >>> from transformers import pipeline
+
+    >>> pipe = pipeline(
+    ...     task="dofa-feature-extraction",
+    ...     model="/path/to/dofa-base-patch16-224",
+    ...     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,
+        wavelengths=None,
+        image_mean=None,
+        image_std=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+
+        if wavelengths is not None:
+            preprocess_params["wavelengths"] = wavelengths
+        if image_mean is not None:
+            preprocess_params["image_mean"] = image_mean
+        if image_std is not None:
+            preprocess_params["image_std"] = image_std
+        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]:
+        """
+        Extract features from multispectral inputs.
+
+        Accepts numpy arrays (HWC), file paths, URLs, or PIL images in addition to the
+        standard image inputs supported by [`ImageFeatureExtractionPipeline`].
+        """
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["DOFAImageFeatureExtractionPipeline"]

dofa-base-patch16-224/preprocessor_config.json

@@ -0,0 +1,48 @@
+{
+  "image_processor_type": "DOFAImageProcessor",
+  "size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_resize": true,
+  "do_rescale": true,
+  "do_normalize": true,
+  "do_convert_rgb": false,
+  "rescale_factor": 0.00392156862745098,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_dofa.DOFAImageProcessor"
+  },
+  "default_wavelengths": [
+    0.665,
+    0.56,
+    0.49,
+    0.705,
+    0.74,
+    0.783,
+    0.842,
+    1.61,
+    2.19
+  ],
+  "image_mean": [
+    114.11,
+    114.82,
+    126.64,
+    84.34,
+    97.85,
+    103.94,
+    101.44,
+    72.33,
+    56.67
+  ],
+  "image_std": [
+    77.84,
+    69.97,
+    67.42,
+    64.57,
+    61.73,
+    61.34,
+    60.3,
+    47.89,
+    42.56
+  ]
+}

dofa-large-patch16-224/config.json

@@ -0,0 +1,73 @@
+{
+  "architectures": [
+    "DOFAModel"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "default_wavelengths": [
+    0.665,
+    0.56,
+    0.49,
+    0.705,
+    0.74,
+    0.783,
+    0.842,
+    1.61,
+    2.19
+  ],
+  "dtype": "float32",
+  "global_pool": true,
+  "head_dropout": 0.0,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 1024,
+  "id2label": {},
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "mlp_ratio": 4.0,
+  "model_type": "dofa",
+  "num_attention_heads": 16,
+  "num_channels": 9,
+  "num_hidden_layers": 24,
+  "patch_size": 16,
+  "qkv_bias": true,
+  "transformers_version": "5.0.0",
+  "wv_planes": 128,
+  "auto_map": {
+    "AutoConfig": "modeling_dofa.DOFAConfig",
+    "AutoModel": "modeling_dofa.DOFAModel",
+    "AutoModelForImageClassification": "modeling_dofa.DOFAForImageClassification"
+  },
+  "custom_pipelines": {
+    "dofa-feature-extraction": {
+      "impl": "pipeline_dofa.DOFAImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "default_image_mean": [
+    114.11,
+    114.82,
+    126.64,
+    84.34,
+    97.85,
+    103.94,
+    101.44,
+    72.33,
+    56.67
+  ],
+  "default_image_std": [
+    77.84,
+    69.97,
+    67.42,
+    64.57,
+    61.73,
+    61.34,
+    60.3,
+    47.89,
+    42.56
+  ]
+}

dofa-large-patch16-224/image_processing_dofa.py

@@ -0,0 +1,180 @@
+# Copyright 2024 The DOFA Authors and The HuggingFace Inc. team.
+"""Image processor for DOFA."""
+
+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
+
+
+class DOFAImageProcessor(BaseImageProcessor):
+    model_input_names = ["pixel_values", "wavelengths"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.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,
+        default_wavelengths: Optional[list[float]] = None,
+        **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.default_wavelengths = default_wavelengths
+
+    @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,
+        wavelengths: Optional[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,
+    ):
+        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
+        wavelengths = wavelengths if wavelengths is not None else self.default_wavelengths
+
+        if wavelengths is None:
+            raise ValueError(
+                "DOFA requires `wavelengths` (one value per channel, in micrometers). "
+                "Pass them to the processor call or set `default_wavelengths` in the processor config."
+            )
+
+        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. "
+                "Provide modality-specific statistics for your input bands."
+            )
+
+        if do_normalize and (
+            len(image_mean) != len(wavelengths) or len(image_std) != len(wavelengths)
+        ):
+            raise ValueError(
+                f"`image_mean` and `image_std` must match the number of wavelengths/channels "
+                f"({len(wavelengths)}), got mean={len(image_mean)} std={len(image_std)}."
+            )
+
+        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
+
+            num_channels = image.shape[-1] if input_data_format == ChannelDimension.LAST else image.shape[0]
+
+            if do_resize:
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(image, size=size, 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,
+            "wavelengths": np.asarray(wavelengths, dtype=np.float32),
+        }
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["DOFAImageProcessor"]

dofa-large-patch16-224/modeling_dofa.py

@@ -0,0 +1,368 @@
+# Copyright 2024 The DOFA Authors and The HuggingFace Inc. team.
+"""Self-contained DOFA model, config, and dynamic patch embedding."""
+
+from functools import partial
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+from timm.models.vision_transformer import Block
+
+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__)
+
+
+class DOFAConfig(PreTrainedConfig):
+    model_type = "dofa"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=None,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.0,
+        attention_probs_dropout_prob=0.0,
+        initializer_range=0.02,
+        layer_norm_eps=1e-6,
+        image_size=224,
+        patch_size=16,
+        num_channels=3,
+        qkv_bias=True,
+        wv_planes=128,
+        mlp_ratio=4.0,
+        global_pool=True,
+        default_wavelengths=None,
+        default_image_mean=None,
+        default_image_std=None,
+        head_dropout=0.0,
+        num_labels=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.wv_planes = wv_planes
+        self.mlp_ratio = mlp_ratio
+        self.global_pool = global_pool
+        self.default_wavelengths = default_wavelengths
+        self.default_image_mean = default_image_mean
+        self.default_image_std = default_image_std
+        self.head_dropout = head_dropout
+        self.num_labels = num_labels
+
+        if intermediate_size is None:
+            self.intermediate_size = int(hidden_size * mlp_ratio)
+        else:
+            self.intermediate_size = intermediate_size
+
+
+def get_1d_sincos_pos_embed_from_grid_torch(embed_dim, pos):
+    assert embed_dim % 2 == 0
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=pos.device)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+
+    pos = pos.reshape(-1)
+    out = torch.einsum("m,d->md", pos, omega)
+
+    emb_sin = torch.sin(out)
+    emb_cos = torch.cos(out)
+
+    return torch.cat([emb_sin, emb_cos], dim=1)
+
+
+class TransformerWeightGenerator(nn.Module):
+    def __init__(self, input_dim, output_dim, embed_dim, num_heads=4, num_layers=1):
+        super().__init__()
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=input_dim,
+            nhead=num_heads,
+            activation="gelu",
+            norm_first=False,
+            batch_first=False,
+            dropout=False,
+        )
+        self.transformer_encoder = nn.TransformerEncoder(
+            encoder_layer, num_layers=num_layers, enable_nested_tensor=False
+        )
+
+        self.fc_weight = nn.Linear(input_dim, output_dim)
+        self.fc_bias = nn.Linear(input_dim, embed_dim)
+        self.wt_num = 128
+        self.weight_tokens = nn.Parameter(torch.empty([self.wt_num, input_dim]))
+        self.bias_token = nn.Parameter(torch.empty([1, input_dim]))
+
+        torch.nn.init.normal_(self.weight_tokens, std=0.02)
+        torch.nn.init.normal_(self.bias_token, std=0.02)
+
+    def forward(self, x):
+        pos_wave = x
+        x = torch.cat([self.weight_tokens, pos_wave], dim=0)
+        x = torch.cat([x, self.bias_token], dim=0)
+        transformer_output = self.transformer_encoder(x)
+        weights = self.fc_weight(transformer_output[self.wt_num : -1] + pos_wave)
+        bias = self.fc_bias(transformer_output[-1])
+        return weights, bias
+
+
+class FCResLayer(nn.Module):
+    def __init__(self, linear_size=128):
+        super().__init__()
+        self.nonlin1 = nn.ReLU(inplace=True)
+        self.nonlin2 = nn.ReLU(inplace=True)
+        self.w1 = nn.Linear(linear_size, linear_size)
+        self.w2 = nn.Linear(linear_size, linear_size)
+
+    def forward(self, x):
+        y = self.w1(x)
+        y = self.nonlin1(y)
+        y = self.w2(y)
+        y = self.nonlin2(y)
+        return x + y
+
+
+class DOFADynamicPatchEmbed(nn.Module):
+    def __init__(self, wv_planes, inter_dim=128, kernel_size=16, embed_dim=768):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.wv_planes = wv_planes
+        self.embed_dim = embed_dim
+        self._num_kernel = self.kernel_size * self.kernel_size * self.embed_dim
+        self.inter_dim = inter_dim
+        self.patch_size = (kernel_size, kernel_size)
+
+        self.weight_generator = TransformerWeightGenerator(wv_planes, self._num_kernel, embed_dim)
+        self.scaler = 0.01
+        self.fclayer = FCResLayer(wv_planes)
+        self.weight_generator.apply(self._weight_init)
+        self.fclayer.apply(self._weight_init)
+
+    def _weight_init(self, module):
+        if isinstance(module, nn.Linear):
+            init.xavier_uniform_(module.weight)
+            module.bias.data.fill_(0.01)
+
+    def forward(self, pixel_values, wavelengths):
+        inplanes = wavelengths.size(0)
+        waves = get_1d_sincos_pos_embed_from_grid_torch(self.wv_planes, wavelengths * 1000)
+        waves = self.fclayer(waves)
+        weight, bias = self.weight_generator(waves)
+
+        dynamic_weight = weight.view(inplanes, self.kernel_size, self.kernel_size, self.embed_dim)
+        dynamic_weight = dynamic_weight.permute([3, 0, 1, 2])
+
+        if bias is not None:
+            bias = bias.view([self.embed_dim]) * self.scaler
+
+        weights = dynamic_weight * self.scaler
+        dynamic_out = F.conv2d(
+            pixel_values, weights, bias=bias, stride=self.kernel_size, padding=1, dilation=1
+        )
+        return dynamic_out.flatten(2).transpose(1, 2), waves
+
+
+def _prepare_wavelengths(wavelengths, pixel_values, default_wavelengths=None):
+    if wavelengths is None:
+        if default_wavelengths is None:
+            raise ValueError(
+                "DOFA requires per-channel wavelengths. Pass `wavelengths` to the model or image processor, "
+                "or set `default_wavelengths` in the model config."
+            )
+        wavelengths = default_wavelengths
+
+    if not torch.is_tensor(wavelengths):
+        wavelengths = torch.tensor(wavelengths, dtype=torch.float32, device=pixel_values.device)
+    else:
+        wavelengths = wavelengths.to(device=pixel_values.device, dtype=torch.float32)
+
+    if wavelengths.dim() == 2 and wavelengths.shape[0] == 1:
+        wavelengths = wavelengths.squeeze(0)
+    elif wavelengths.dim() == 2 and wavelengths.shape[0] == pixel_values.shape[0]:
+        if not torch.allclose(wavelengths[0], wavelengths):
+            raise ValueError("DOFA currently expects identical wavelengths for all items in a batch.")
+        wavelengths = wavelengths[0]
+
+    if wavelengths.dim() != 1:
+        raise ValueError("`wavelengths` must be a 1D tensor/list with one value per input channel.")
+
+    if wavelengths.shape[0] != pixel_values.shape[1]:
+        raise ValueError(
+            f"Expected {pixel_values.shape[1]} wavelengths for {pixel_values.shape[1]} channels, "
+            f"but got {wavelengths.shape[0]}."
+        )
+
+    return wavelengths
+
+
+class DOFAPreTrainedModel(PreTrainedModel):
+    config_class = DOFAConfig
+    base_model_prefix = "dofa"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Block"]
+    _supports_sdpa = False
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, DOFAModel):
+            if hasattr(module, "pos_embed"):
+                nn.init.trunc_normal_(module.pos_embed, std=self.config.initializer_range)
+            if hasattr(module, "cls_token"):
+                nn.init.trunc_normal_(module.cls_token, std=self.config.initializer_range)
+
+
+class DOFAModel(DOFAPreTrainedModel):
+    def __init__(self, config: DOFAConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+
+        image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0]
+        self.patch_embed = DOFADynamicPatchEmbed(
+            wv_planes=config.wv_planes,
+            inter_dim=config.wv_planes,
+            kernel_size=config.patch_size,
+            embed_dim=config.hidden_size,
+        )
+        self.num_patches = (image_size // config.patch_size) ** 2
+        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), requires_grad=False
+        )
+
+        norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
+        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)
+
+        self.add_pooling_layer = add_pooling_layer
+        self.post_init()
+
+    def forward_features(self, pixel_values, wavelengths):
+        patch_tokens, _ = self.patch_embed(pixel_values, wavelengths)
+        patch_tokens = patch_tokens + self.pos_embed[:, 1:, :]
+
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(pixel_values.shape[0], -1, -1)
+        hidden_states = torch.cat((cls_tokens, patch_tokens), 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:
+            pooled_output = self.norm(hidden_states)[:, 0]
+
+        return hidden_states, pooled_output
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        wavelengths: Optional[Union[torch.Tensor, list]] = 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
+
+        wavelengths = _prepare_wavelengths(wavelengths, pixel_values, self.config.default_wavelengths)
+        last_hidden_state, pooled_output = self.forward_features(pixel_values, wavelengths)
+
+        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 DOFAForImageClassification(DOFAPreTrainedModel):
+    def __init__(self, config: DOFAConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.dofa = DOFAModel(config, add_pooling_layer=True)
+        self.head_drop = nn.Dropout(config.head_dropout)
+        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,
+        wavelengths: Optional[Union[torch.Tensor, list]] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.dofa(
+            pixel_values=pixel_values,
+            wavelengths=wavelengths,
+            return_dict=True,
+            **kwargs,
+        )
+        pooled_output = self.head_drop(outputs.pooler_output)
+        logits = self.classifier(pooled_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__ = ["DOFAConfig", "DOFAForImageClassification", "DOFAModel", "DOFAPreTrainedModel"]

dofa-large-patch16-224/pipeline_dofa.py

@@ -0,0 +1,95 @@
+# Copyright 2024 The DOFA Authors and The HuggingFace Inc. team.
+"""Custom DOFA feature extraction pipeline for trust_remote_code loading."""
+
+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),
+    """
+        wavelengths (`list[float]`, *optional*):
+            Per-channel wavelengths in micrometers. Defaults to Sentinel-2 bands from config.
+        image_mean (`list[float]`, *optional*):
+            Per-channel normalization mean values.
+        image_std (`list[float]`, *optional*):
+            Per-channel normalization standard deviation values.
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether to return pooled features instead of sequence hidden states.
+    """,
+)
+class DOFAImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    DOFA wavelength-conditioned multispectral feature extraction pipeline.
+
+    This pipeline can be loaded from [`pipeline`] using the task identifier `"dofa-feature-extraction"`
+    with `trust_remote_code=True`.
+
+    Example:
+
+    ```python
+    >>> from transformers import pipeline
+
+    >>> pipe = pipeline(
+    ...     task="dofa-feature-extraction",
+    ...     model="/path/to/dofa-base-patch16-224",
+    ...     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,
+        wavelengths=None,
+        image_mean=None,
+        image_std=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+
+        if wavelengths is not None:
+            preprocess_params["wavelengths"] = wavelengths
+        if image_mean is not None:
+            preprocess_params["image_mean"] = image_mean
+        if image_std is not None:
+            preprocess_params["image_std"] = image_std
+        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]:
+        """
+        Extract features from multispectral inputs.
+
+        Accepts numpy arrays (HWC), file paths, URLs, or PIL images in addition to the
+        standard image inputs supported by [`ImageFeatureExtractionPipeline`].
+        """
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["DOFAImageFeatureExtractionPipeline"]

dofa-large-patch16-224/preprocessor_config.json

@@ -0,0 +1,48 @@
+{
+  "image_processor_type": "DOFAImageProcessor",
+  "size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_resize": true,
+  "do_rescale": true,
+  "do_normalize": true,
+  "do_convert_rgb": false,
+  "rescale_factor": 0.00392156862745098,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_dofa.DOFAImageProcessor"
+  },
+  "default_wavelengths": [
+    0.665,
+    0.56,
+    0.49,
+    0.705,
+    0.74,
+    0.783,
+    0.842,
+    1.61,
+    2.19
+  ],
+  "image_mean": [
+    114.11,
+    114.82,
+    126.64,
+    84.34,
+    97.85,
+    103.94,
+    101.44,
+    72.33,
+    56.67
+  ],
+  "image_std": [
+    77.84,
+    69.97,
+    67.42,
+    64.57,
+    61.73,
+    61.34,
+    60.3,
+    47.89,
+    42.56
+  ]
+}

test_dofa.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""Quick CLI smoke test for a self-contained DOFA checkpoint folder."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import pipeline
+
+
+def parse_args() -> argparse.Namespace:
+    repo_root = Path(__file__).resolve().parent
+    parser = argparse.ArgumentParser(description="Run DOFA feature extraction on dummy or real input.")
+    parser.add_argument(
+        "--model",
+        type=Path,
+        default=repo_root / "dofa-base-patch16-224",
+        help="Path to a checkpoint folder (default: dofa-base-patch16-224)",
+    )
+    parser.add_argument(
+        "--image",
+        type=Path,
+        default=None,
+        help="Optional image path (.npy HWC array or image readable by rasterio/PIL)",
+    )
+    parser.add_argument(
+        "--pool",
+        action="store_true",
+        default=True,
+        help="Return pooled features (default: True)",
+    )
+    parser.add_argument(
+        "--no-pool",
+        action="store_true",
+        help="Return sequence features instead of pooled output",
+    )
+    parser.add_argument(
+        "--device",
+        default=None,
+        help="Torch device, e.g. cuda or cpu (default: auto)",
+    )
+    return parser.parse_args()
+
+
+def load_image(path: Path, num_channels: int) -> np.ndarray:
+    if path.suffix == ".npy":
+        array = np.load(path)
+        if array.ndim != 3:
+            raise ValueError(f"Expected HWC numpy array, got shape {array.shape}")
+        return array
+
+    try:
+        import rasterio
+    except ImportError as exc:
+        raise ImportError("Install rasterio to load geospatial images, or pass a .npy file.") from exc
+
+    with rasterio.open(path) as src:
+        array = src.read()
+        array = np.transpose(array, (1, 2, 0))
+    return array
+
+
+def main() -> None:
+    args = parse_args()
+    model_dir = args.model.resolve()
+    if not model_dir.is_dir():
+        raise SystemExit(f"Model folder not found: {model_dir}")
+
+    config_path = model_dir / "config.json"
+    with open(config_path, encoding="utf-8") as handle:
+        config = json.load(handle)
+
+    num_channels = config.get("num_channels") or len(config["default_wavelengths"])
+
+    if args.image is None:
+        image = np.random.randint(0, 255, (224, 224, num_channels), dtype=np.uint8)
+        source = f"random dummy array ({num_channels} channels)"
+    else:
+        image = load_image(args.image.resolve(), num_channels)
+        source = str(args.image)
+
+    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
+    pool = not args.no_pool
+
+    print(f"model:  {model_dir}")
+    print(f"input:  {source}")
+    print(f"device: {device}")
+    print(f"pool:   {pool}")
+
+    pipe = pipeline(
+        task="dofa-feature-extraction",
+        model=str(model_dir),
+        trust_remote_code=True,
+        device=device,
+    )
+
+    features = pipe(image, pool=pool, return_tensors=True)
+    print(f"output: {tuple(features.shape)} dtype={features.dtype}")
+    print("OK")
+
+
+if __name__ == "__main__":
+    main()