diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..595a75e31493e7c68f84370be42049134536b849 --- /dev/null +++ b/README.md @@ -0,0 +1,164 @@ +--- +license: mit +language: +- en +tags: +- remote-sensing +- earth-observation +- self-supervised-learning +- satellite +- multispectral +- feature-extraction +- convnext +- mae +- mmearth +- mp-mae +- transformers +library_name: transformers +pipeline_tag: feature-extraction +--- + +# MMEarth Transformers Models + +Hugging Face–compatible checkpoints converted from the official [MMEarth](https://arxiv.org/abs/2405.02771) MP-MAE pretrained weights. Each subfolder is a standalone model repo layout (`config.json`, `model.safetensors`, preprocessor, and remote code) for geospatial feature extraction. + +## Model Description + +These models are ConvNeXt V2 encoders pretrained with Multi Pretext Masked Autoencoding (MP-MAE) on the [MMEarth](https://github.com/vishalned/MMEarth-data) multi-modal geospatial dataset. Checkpoints cover different pretext task configurations (all modalities, S2-only, RGB/BGR, image-level, pixel-level) and model sizes (atto, tiny). + +All folders ship self-contained remote code (`modeling_mmearth.py`, processor, pipeline) and load with `trust_remote_code=True`. + +**Developed by:** [MMEarth Authors](https://github.com/vishalned/MMEarth-train) +**Converted for Hugging Face by:** BiliSakura +**License (weights):** MIT +**Original paper:** [MMEarth: Exploring Multi-Modal Pretext Tasks For Geospatial Representation Learning](https://arxiv.org/abs/2405.02771) (ECCV 2024) + +## Available checkpoints (10 models) + +| Folder | Input | Size | Dataset | Loss | Image | Patch | Ch | +|--------|-------|------|---------|------|-------|-------|----| +| `mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8` | all_mod | atto | 1M_64 | uncertainty | 56 | 8 | 12 | +| `mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8` | all_mod | atto | 1M_64 | unweighted | 56 | 8 | 12 | +| `mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16` | all_mod | atto | 1M_128 | uncertainty | 112 | 16 | 12 | +| `mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16` | all_mod | atto | 100k_128 | uncertainty | 112 | 16 | 12 | +| `mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8` | all_mod | tiny | 1M_64 | uncertainty | 56 | 8 | 12 | +| `mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8` | S2 | atto | 1M_64 | uncertainty | 56 | 8 | 12 | +| `mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8` | rgb (BGR) | atto | 1M_64 | uncertainty | 56 | 8 | 3 | +| `mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16` | rgb (BGR) | atto | 1M_128 | uncertainty | 112 | 16 | 3 | +| `mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8` | img_mod | atto | 1M_64 | uncertainty | 56 | 8 | 12 | +| `mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8` | pix_mod | atto | 1M_64 | uncertainty | 56 | 8 | 12 | + +Legacy `.pth` filename mapping is in [`conversion_manifest.json`](conversion_manifest.json). + +## Usage + +Processors default to **`do_resize: false`**. Inputs keep native height and width. Apply per-band MMEarth normalization when you have dataset statistics (`image_mean` / `image_std`). + +```python +from transformers import pipeline +import numpy as np + +MODEL = "/path/to/MMEarth-transformers/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8" + +pipe = pipeline( + task="mmearth-feature-extraction", + model=MODEL, + trust_remote_code=True, +) + +# RGB/BGR: 3 bands at native size (56×56 for this checkpoint) +image = np.random.rand(56, 56, 3).astype(np.float32) * 1000 +features = pipe(image, pool=True, return_tensors=True) +print(features.shape) # torch.Size([1, 320]) +``` + +12-band Sentinel-2 (all_mod / S2 checkpoints): + +```python +MODEL = "/path/to/MMEarth-transformers/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8" +pipe = pipeline(task="mmearth-feature-extraction", model=MODEL, trust_remote_code=True) + +image = np.random.rand(56, 56, 12).astype(np.float32) * 1000 +features = pipe(image, pool=True, return_tensors=True) +print(features.shape) # torch.Size([1, 320]) +``` + +Dense spatial token map: + +```python +tokens = pipe(image, pool=False, return_tensors=True) +print(tokens.shape) # [1, num_patches, hidden_size] +``` + +To resize to the pretraining reference size: + +```python +features = pipe(image, pool=True, return_tensors=True, image_processor_kwargs={"do_resize": True}) +``` + +Load components directly: + +```python +from transformers import AutoModel, AutoImageProcessor + +model = AutoModel.from_pretrained(MODEL, trust_remote_code=True) +processor = AutoImageProcessor.from_pretrained(MODEL, trust_remote_code=True) +``` + +## Custom pipeline + +Each checkpoint registers a custom pipeline in `config.json`: + +```json +"custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": ["AutoModel"] + } +} +``` + +This follows the [HuggingFace custom pipeline pattern](https://huggingface.co/docs/transformers/add_new_pipeline): remote code ships with the model folder, and `trust_remote_code=True` loads `MMEarthImageFeatureExtractionPipeline`, which extends the standard `ImageFeatureExtractionPipeline` with numpy array and file path support. + +The built-in `image-feature-extraction` task also works: + +```python +pipe = pipeline(task="image-feature-extraction", model=MODEL, trust_remote_code=True) +``` + +## Normalization + +MMEarth pretraining normalizes each band with dataset-specific mean/std from `data_*_band_stats.json`. The converted preprocessor defaults to `do_normalize: false` because band statistics are not embedded in the legacy checkpoints. Provide your own `image_mean` / `image_std` when preprocessing: + +```python +features = pipe( + image, + pool=True, + return_tensors=True, + image_processor_kwargs={ + "do_normalize": True, + "image_mean": [...], # one value per channel + "image_std": [...], + }, +) +``` + +RGB checkpoints were trained with **BGR** channel order (bands B4, B3, B2). The processor swaps RGB→BGR when `channel_order="bgr"`. + +## Dependencies + +- `transformers`, `torch`, `timm`, `safetensors` +- `opencv-python` (multispectral resize with more than 4 channels when `do_resize=True`) + +## Citation + +```bibtex +@inproceedings{nedungadi2024mmearth, + title={MMEarth: Exploring multi-modal pretext tasks for geospatial representation learning}, + author={Nedungadi, Vishal and Kariryaa, Ankit and Oehmcke, Stefan and Belongie, Serge and Igel, Christian and Lang, Nico}, + booktitle={European Conference on Computer Vision}, + pages={164--182}, + year={2024}, + organization={Springer} +} +``` diff --git a/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/config.json b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/config.json new file mode 100644 index 0000000000000000000000000000000000000000..42ad092d9f115f2570cd1341f2119354d9e1f580 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/config.json @@ -0,0 +1,70 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B1", + "B2", + "B3", + "B4", + "B5", + "B6", + "B7", + "B8A", + "B8", + "B9", + "B11", + "B12" + ], + "channel_order": "rgb", + "checkpoint_stage": "pretrain", + "dataset": "100k_128", + "depths": [ + 2, + 2, + 6, + 2 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 320, + "hidden_sizes": [ + 40, + 80, + 160, + 320 + ], + "id2label": {}, + "image_mean": null, + "image_size": 112, + "image_std": null, + "input_modality": "all_mod", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "uncertainty", + "model_size": "atto", + "model_type": "mmearth", + "num_channels": 12, + "num_stages": 4, + "patch_size": 16, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-all_mod_atto_100k_128_uncertainty_112-16" +} diff --git a/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/image_processing_mmearth.py b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/model.safetensors b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..baf0cc12a97d115521f81299cf7e7d4d869d4297 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:73ad3615ab31506d9e2829912385adcb82140d3ab9c07784d37db8af5f2caaae +size 13573176 diff --git a/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/modeling_mmearth.py b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/pipeline_mmearth.py b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/preprocessor_config.json b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..8f512286989d7a41c650f9c85a1db92fc9ee921c --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 112, + "width": 112 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "rgb", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +} diff --git a/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/config.json b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/config.json new file mode 100644 index 0000000000000000000000000000000000000000..0bea07af5df04ae2a5a79c82bc1f034ea13a215d --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/config.json @@ -0,0 +1,70 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B1", + "B2", + "B3", + "B4", + "B5", + "B6", + "B7", + "B8A", + "B8", + "B9", + "B11", + "B12" + ], + "channel_order": "rgb", + "checkpoint_stage": "pretrain", + "dataset": "1M_128", + "depths": [ + 2, + 2, + 6, + 2 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 320, + "hidden_sizes": [ + 40, + 80, + 160, + 320 + ], + "id2label": {}, + "image_mean": null, + "image_size": 112, + "image_std": null, + "input_modality": "all_mod", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "uncertainty", + "model_size": "atto", + "model_type": "mmearth", + "num_channels": 12, + "num_stages": 4, + "patch_size": 16, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-all_mod_atto_1M_128_uncertainty_112-16" +} diff --git a/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/image_processing_mmearth.py b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/model.safetensors b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..2474d5cf65b90377ec259b578743b1e6d25e955a --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f2f0e7cecf8534809972730b4960c07078387fc72f97e3d255ce20286803815b +size 13573176 diff --git a/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/modeling_mmearth.py b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/pipeline_mmearth.py b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/preprocessor_config.json b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..8f512286989d7a41c650f9c85a1db92fc9ee921c --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 112, + "width": 112 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "rgb", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +} diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/config.json b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/config.json new file mode 100644 index 0000000000000000000000000000000000000000..7359b3cbb06a1aec38a3ee97a8cb960c8e66bf42 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/config.json @@ -0,0 +1,70 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B1", + "B2", + "B3", + "B4", + "B5", + "B6", + "B7", + "B8A", + "B8", + "B9", + "B11", + "B12" + ], + "channel_order": "rgb", + "checkpoint_stage": "pretrain", + "dataset": "1M_64", + "depths": [ + 2, + 2, + 6, + 2 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 320, + "hidden_sizes": [ + 40, + 80, + 160, + 320 + ], + "id2label": {}, + "image_mean": null, + "image_size": 56, + "image_std": null, + "input_modality": "all_mod", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "uncertainty", + "model_size": "atto", + "model_type": "mmearth", + "num_channels": 12, + "num_stages": 4, + "patch_size": 8, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-all_mod_atto_1M_64_uncertainty_56-8" +} diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/model.safetensors b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..ce5cb15ce1c4875b2555bef55e0be4c39e948853 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cf0cfcaad31d7cd1cf4f7e730bec625d0f4d769eec25bf9aa3bed1e0bb105648 +size 13572696 diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/modeling_mmearth.py b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/preprocessor_config.json b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..ed53bb0fb1dcaf78a697525ba4006afbb22de2b7 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 56, + "width": 56 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "rgb", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +} diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/config.json b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae1142aa689ed86aa4569f0ca80442b8e7393a07 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/config.json @@ -0,0 +1,70 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B1", + "B2", + "B3", + "B4", + "B5", + "B6", + "B7", + "B8A", + "B8", + "B9", + "B11", + "B12" + ], + "channel_order": "rgb", + "checkpoint_stage": "pretrain", + "dataset": "1M_64", + "depths": [ + 2, + 2, + 6, + 2 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 320, + "hidden_sizes": [ + 40, + 80, + 160, + 320 + ], + "id2label": {}, + "image_mean": null, + "image_size": 56, + "image_std": null, + "input_modality": "all_mod", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "unweighted", + "model_size": "atto", + "model_type": "mmearth", + "num_channels": 12, + "num_stages": 4, + "patch_size": 8, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-all_mod_atto_1M_64_unweighted_56-8" +} diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/image_processing_mmearth.py b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/model.safetensors b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..2a1c86e02726536798a86f6467f728680ec23d03 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:232078eaefd5615e3d4e8b9d7b3c0aa5836611b1bea18c3bfaa2f8b3aa397560 +size 13572696 diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/modeling_mmearth.py b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/pipeline_mmearth.py b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/preprocessor_config.json b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..ed53bb0fb1dcaf78a697525ba4006afbb22de2b7 --- /dev/null +++ b/mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 56, + "width": 56 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "rgb", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +} diff --git a/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/config.json b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/config.json new file mode 100644 index 0000000000000000000000000000000000000000..4eb4c60e3da5b809d90896969afbe8c00709ccd5 --- /dev/null +++ b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/config.json @@ -0,0 +1,70 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B1", + "B2", + "B3", + "B4", + "B5", + "B6", + "B7", + "B8A", + "B8", + "B9", + "B11", + "B12" + ], + "channel_order": "rgb", + "checkpoint_stage": "pretrain", + "dataset": "1M_64", + "depths": [ + 2, + 2, + 6, + 2 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 320, + "hidden_sizes": [ + 40, + 80, + 160, + 320 + ], + "id2label": {}, + "image_mean": null, + "image_size": 56, + "image_std": null, + "input_modality": "img_mod", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "uncertainty", + "model_size": "atto", + "model_type": "mmearth", + "num_channels": 12, + "num_stages": 4, + "patch_size": 8, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-img_mod_atto_1M_64_uncertainty_56-8" +} diff --git a/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/model.safetensors b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..33d9450603bc1c77faee8119c7b535821a3d2577 --- /dev/null +++ b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1ec8508296b7288978f784efa4db9e9df14c52e9e74bb0d618c6eb4d28f45c98 +size 13572696 diff --git a/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/modeling_mmearth.py b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/preprocessor_config.json b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..ed53bb0fb1dcaf78a697525ba4006afbb22de2b7 --- /dev/null +++ b/mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 56, + "width": 56 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "rgb", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +} diff --git a/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/config.json b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/config.json new file mode 100644 index 0000000000000000000000000000000000000000..d8617bb2851b016702d76be7725b4f00b02528d8 --- /dev/null +++ b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/config.json @@ -0,0 +1,70 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B1", + "B2", + "B3", + "B4", + "B5", + "B6", + "B7", + "B8A", + "B8", + "B9", + "B11", + "B12" + ], + "channel_order": "rgb", + "checkpoint_stage": "pretrain", + "dataset": "1M_64", + "depths": [ + 2, + 2, + 6, + 2 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 320, + "hidden_sizes": [ + 40, + 80, + 160, + 320 + ], + "id2label": {}, + "image_mean": null, + "image_size": 56, + "image_std": null, + "input_modality": "pix_mod", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "uncertainty", + "model_size": "atto", + "model_type": "mmearth", + "num_channels": 12, + "num_stages": 4, + "patch_size": 8, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-pix_mod_atto_1M_64_uncertainty_56-8" +} diff --git a/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/model.safetensors b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..bee4efb5454fece86351bbfae4d9755ee99f16ce --- /dev/null +++ b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:527f40016c44e6183103af0ee0e0668cab617ccafe29868e1ea78570871f0905 +size 13572696 diff --git a/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/modeling_mmearth.py b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/preprocessor_config.json b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..ed53bb0fb1dcaf78a697525ba4006afbb22de2b7 --- /dev/null +++ b/mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 56, + "width": 56 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "rgb", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +} diff --git a/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/config.json b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/config.json new file mode 100644 index 0000000000000000000000000000000000000000..c19c0f9b63c93eff8d9a8201604c4d23117c648c --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/config.json @@ -0,0 +1,61 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B4", + "B3", + "B2" + ], + "channel_order": "bgr", + "checkpoint_stage": "pretrain", + "dataset": "1M_128", + "depths": [ + 2, + 2, + 6, + 2 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 320, + "hidden_sizes": [ + 40, + 80, + 160, + 320 + ], + "id2label": {}, + "image_mean": null, + "image_size": 112, + "image_std": null, + "input_modality": "rgb", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "uncertainty", + "model_size": "atto", + "model_type": "mmearth", + "num_channels": 3, + "num_stages": 4, + "patch_size": 16, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-rgb_atto_1M_128_uncertainty_112-16" +} diff --git a/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/image_processing_mmearth.py b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/model.safetensors b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e49aa8eab6098393489594c020011f89385eecf5 --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0eb3fccae67ad295325778b6cf730834438b2d24ee0f71226aa7d4da9c8351aa +size 13560216 diff --git a/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/modeling_mmearth.py b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/pipeline_mmearth.py b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/preprocessor_config.json b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..44d968d56c79f40147ffdc46e8fa4bdf7b529abd --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 112, + "width": 112 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "bgr", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +} diff --git a/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/config.json b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ad9ed2f781b8b4bcdd68979cb3e6ed2146172847 --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/config.json @@ -0,0 +1,61 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B4", + "B3", + "B2" + ], + "channel_order": "bgr", + "checkpoint_stage": "pretrain", + "dataset": "1M_64", + "depths": [ + 2, + 2, + 6, + 2 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 320, + "hidden_sizes": [ + 40, + 80, + 160, + 320 + ], + "id2label": {}, + "image_mean": null, + "image_size": 56, + "image_std": null, + "input_modality": "rgb", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "uncertainty", + "model_size": "atto", + "model_type": "mmearth", + "num_channels": 3, + "num_stages": 4, + "patch_size": 8, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-rgb_atto_1M_64_uncertainty_56-8" +} diff --git a/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/image_processing_mmearth.py b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/model.safetensors b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..608a5e754df1bb45719f3b5ae7287a74ed2f1c98 --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:72548211ba30d3d39b82cd43c286a1c1ec94d46ae3ac7af873b16594bf0e052f +size 13559736 diff --git a/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/modeling_mmearth.py b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/pipeline_mmearth.py b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/preprocessor_config.json b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..b6c3a01882a2eff22e6e704a6c4ca3d0ee371aa8 --- /dev/null +++ b/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 56, + "width": 56 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "bgr", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +} diff --git a/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/config.json b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/config.json new file mode 100644 index 0000000000000000000000000000000000000000..296c51fe9b7d7434756f34e6dc691d48c1fe696d --- /dev/null +++ b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/config.json @@ -0,0 +1,70 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B1", + "B2", + "B3", + "B4", + "B5", + "B6", + "B7", + "B8A", + "B8", + "B9", + "B11", + "B12" + ], + "channel_order": "rgb", + "checkpoint_stage": "pretrain", + "dataset": "1M_64", + "depths": [ + 2, + 2, + 6, + 2 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 320, + "hidden_sizes": [ + 40, + 80, + 160, + 320 + ], + "id2label": {}, + "image_mean": null, + "image_size": 56, + "image_std": null, + "input_modality": "S2", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "uncertainty", + "model_size": "atto", + "model_type": "mmearth", + "num_channels": 12, + "num_stages": 4, + "patch_size": 8, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-S2_atto_1M_64_uncertainty_56-8" +} diff --git a/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/image_processing_mmearth.py b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/model.safetensors b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6ea619e652c01c012806f38f78fd0b4edcc08931 --- /dev/null +++ b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f60576aef5ff8fb71fb1d63544737afbe9e5711709ba7f2acb3afe180a6abbcb +size 13572696 diff --git a/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/modeling_mmearth.py b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/pipeline_mmearth.py b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/preprocessor_config.json b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..ed53bb0fb1dcaf78a697525ba4006afbb22de2b7 --- /dev/null +++ b/mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 56, + "width": 56 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "rgb", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +} diff --git a/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/config.json b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/config.json new file mode 100644 index 0000000000000000000000000000000000000000..e94a2aead1391dc037c0db5fbad8d88bd28c18e6 --- /dev/null +++ b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/config.json @@ -0,0 +1,70 @@ +{ + "architectures": [ + "MMEarthModel" + ], + "band_names": [ + "B1", + "B2", + "B3", + "B4", + "B5", + "B6", + "B7", + "B8A", + "B8", + "B9", + "B11", + "B12" + ], + "channel_order": "rgb", + "checkpoint_stage": "pretrain", + "dataset": "1M_64", + "depths": [ + 3, + 3, + 9, + 3 + ], + "do_rescale": false, + "drop_path_rate": 0.0, + "dtype": "float32", + "hidden_act": "gelu", + "hidden_size": 768, + "hidden_sizes": [ + 96, + 192, + 384, + 768 + ], + "id2label": {}, + "image_mean": null, + "image_size": 56, + "image_std": null, + "input_modality": "all_mod", + "label2id": {}, + "layer_norm_eps": 1e-06, + "loss_aggr": "uncertainty", + "model_size": "tiny", + "model_type": "mmearth", + "num_channels": 12, + "num_stages": 4, + "patch_size": 8, + "rescale_factor": 1.0, + "transformers_version": "5.0.0", + "use_orig_stem": false, + "auto_map": { + "AutoConfig": "modeling_mmearth.MMEarthConfig", + "AutoModel": "modeling_mmearth.MMEarthModel", + "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification" + }, + "custom_pipelines": { + "mmearth-feature-extraction": { + "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + }, + "legacy_checkpoint": "checkpoint-199.pth", + "legacy_source": "pt-all_mod_tiny_1M_64_uncertainty_56-8" +} diff --git a/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd7e8248c0a4263afea443be4ad165e1e375374 --- /dev/null +++ b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py @@ -0,0 +1,189 @@ +# Copyright 2024 MMEarth 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 MMEarth 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 MMEarthImageProcessor(BaseImageProcessor): + """ + Image processor for MMEarth ConvNeXt V2 encoders. + + 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": 112, "width": 112} + 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__ = ["MMEarthImageProcessor"] diff --git a/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/model.safetensors b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a783e97efda8b9d897c48c89e376dab831db9749 --- /dev/null +++ b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8f29593a7f2417250b0fbfbb7fec99a00e71553cd80c970b97e875f66da5862d +size 111509184 diff --git a/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/modeling_mmearth.py b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/modeling_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..0a718d4c4c6659eedf39ce815ae729ead294f8ef --- /dev/null +++ b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/modeling_mmearth.py @@ -0,0 +1,292 @@ +# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team. +"""Self-contained MMEarth model and config for trust_remote_code loading.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ +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__) + +MODEL_SIZE_PRESETS = { + "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]}, + "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]}, + "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]}, + "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]}, + "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]}, + "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]}, + "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]}, + "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]}, +} + + +class MMEarthConfig(PreTrainedConfig): + model_type = "mmearth" + + def __init__( + self, + depths: list[int] | None = None, + hidden_sizes: list[int] | None = None, + num_channels: int = 12, + image_size: int = 112, + patch_size: int = 16, + drop_path_rate: float = 0.0, + layer_norm_eps: float = 1e-6, + hidden_act: str = "gelu", + use_orig_stem: bool = False, + model_size: str = "atto", + input_modality: str = "all_mod", + channel_order: str = "rgb", + dataset: str = "1M_128", + loss_aggr: str = "uncertainty", + checkpoint_stage: str = "pretrain", + image_mean: list[float] | None = None, + image_std: list[float] | None = None, + do_rescale: bool = False, + rescale_factor: float = 1.0, + num_labels: int = 0, + **kwargs, + ): + super().__init__(**kwargs) + preset = MODEL_SIZE_PRESETS[model_size] + self.model_size = model_size + self.input_modality = input_modality + self.channel_order = channel_order + self.dataset = dataset + self.loss_aggr = loss_aggr + self.checkpoint_stage = checkpoint_stage + self.num_channels = num_channels + self.image_size = image_size + self.patch_size = patch_size + self.drop_path_rate = drop_path_rate + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.use_orig_stem = use_orig_stem + self.num_labels = num_labels + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.depths = depths if depths is not None else preset["depths"] + self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"] + self.hidden_size = self.hidden_sizes[-1] + self.num_stages = len(self.depths) + self.image_mean = image_mean + self.image_std = image_std + + +class MMEarthLayerNorm(nn.Module): + def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + return self.weight[:, None, None] * x + self.bias[:, None, None] + + +class MMEarthGRN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4) + return self.gamma * (x * nx) + self.beta + x + + +class MMEarthConvNeXtBlock(nn.Module): + def __init__(self, dim: int, drop_path: float = 0.0): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.norm = MMEarthLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.grn = MMEarthGRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_tensor = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return input_tensor + self.drop_path(x) + + +class MMEarthPreTrainedModel(PreTrainedModel): + config_class = MMEarthConfig + config: MMEarthConfig + base_model_prefix = "mmearth" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["MMEarthConvNeXtBlock"] + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, (nn.Conv2d, nn.Linear)): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class MMEarthModel(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True): + super().__init__(config) + self.config = config + self.add_pooling_layer = add_pooling_layer + depths = config.depths + dims = config.hidden_sizes + patch_size = config.patch_size + num_stages = len(depths) + self.downsample_layers = nn.ModuleList() + if config.use_orig_stem: + self.stem_orig = nn.Sequential( + nn.Conv2d( + config.num_channels, + dims[0], + kernel_size=patch_size // (2 ** (num_stages - 1)), + stride=patch_size // (2 ** (num_stages - 1)), + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + self.initial_conv = None + self.stem = None + else: + self.stem_orig = None + self.initial_conv = nn.Sequential( + nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + nn.GELU(), + ) + stem_kernel = patch_size // (2 ** (num_stages - 1)) + self.stem = nn.Sequential( + nn.Conv2d( + dims[0], + dims[0], + kernel_size=stem_kernel, + stride=stem_kernel, + padding=stem_kernel // 2, + groups=dims[0], + ), + MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"), + ) + for i in range(3): + self.downsample_layers.append( + nn.Sequential( + MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), + ) + ) + dp_rates = [ + x.item() + for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu")) + ] + cur = 0 + self.stages = nn.ModuleList() + for i in range(num_stages): + stage = nn.Sequential( + *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps) + self.post_init() + + def _forward_stem(self, x: torch.Tensor) -> torch.Tensor: + if self.config.use_orig_stem: + return self.stem_orig(x) + x = self.initial_conv(x) + return self.stem(x) + + def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor: + x = self._forward_stem(pixel_values) + x = self.stages[0](x) + for i in range(3): + x = self.downsample_layers[i](x) + x = self.stages[i + 1](x) + return x + + 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 + spatial_features = self.forward_features(pixel_values) + last_hidden_state = spatial_features.flatten(2).transpose(1, 2) + pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None + if not return_dict: + return (last_hidden_state, pooled_output) + return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) + + +class MMEarthForImageClassification(MMEarthPreTrainedModel): + def __init__(self, config: MMEarthConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.mmearth = MMEarthModel(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.mmearth(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__ = [ + "MMEarthConfig", + "MMEarthForImageClassification", + "MMEarthModel", + "MMEarthPreTrainedModel", +] diff --git a/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py new file mode 100644 index 0000000000000000000000000000000000000000..98d8212a3309c06e0b6f3853a1469f585fa38e48 --- /dev/null +++ b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py @@ -0,0 +1,68 @@ +# Copyright 2024 MMEarth 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. +"""MMEarth 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 MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): + """ + MMEarth image feature extraction pipeline. + + This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction. + It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to + standard image inputs. + """ + + 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__ = ["MMEarthImageFeatureExtractionPipeline"] diff --git a/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/preprocessor_config.json b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..ed53bb0fb1dcaf78a697525ba4006afbb22de2b7 --- /dev/null +++ b/mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8/preprocessor_config.json @@ -0,0 +1,18 @@ +{ + "image_processor_type": "MMEarthImageProcessor", + "size": { + "height": 56, + "width": 56 + }, + "do_resize": false, + "do_rescale": false, + "rescale_factor": 1.0, + "do_normalize": false, + "do_convert_rgb": false, + "channel_order": "rgb", + "image_mean": null, + "image_std": null, + "auto_map": { + "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor" + } +}