Image Feature Extraction
Transformers
Safetensors
English
vision
sar
remote-sensing
synthetic-aperture-radar
masked-autoencoder
model-hub
Instructions to use BiliSakura/SARMAE-transformers with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use BiliSakura/SARMAE-transformers with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("image-feature-extraction", model="BiliSakura/SARMAE-transformers")# Load model directly from transformers import AutoModel model = AutoModel.from_pretrained("BiliSakura/SARMAE-transformers", device_map="auto") - Notebooks
- Google Colab
- Kaggle
Upload SARMAE ViT-L transformers-format pretrain checkpoint
Browse files- vit-large-patch16-pretrain/README.md +128 -0
- vit-large-patch16-pretrain/config.json +50 -0
- vit-large-patch16-pretrain/image_processing_sarmae.py +164 -0
- vit-large-patch16-pretrain/model.safetensors +3 -0
- vit-large-patch16-pretrain/modeling_sarmae.py +200 -0
- vit-large-patch16-pretrain/pipeline_sarmae.py +51 -0
- vit-large-patch16-pretrain/preprocessor_config.json +26 -0
vit-large-patch16-pretrain/README.md
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: cc-by-nc-4.0
|
| 3 |
+
language: en
|
| 4 |
+
tags:
|
| 5 |
+
- vision
|
| 6 |
+
- image-feature-extraction
|
| 7 |
+
- sar
|
| 8 |
+
- remote-sensing
|
| 9 |
+
- synthetic-aperture-radar
|
| 10 |
+
- masked-autoencoder
|
| 11 |
+
- transformers
|
| 12 |
+
library_name: transformers
|
| 13 |
+
pipeline_tag: image-feature-extraction
|
| 14 |
+
datasets:
|
| 15 |
+
- Wenquandan777/SAR-1M
|
| 16 |
+
arxiv: 2512.16635
|
| 17 |
+
base_model: Wenquandan777/SARMAE
|
| 18 |
+
model-index:
|
| 19 |
+
- name: sarmae-vit-large-patch16-pretrain
|
| 20 |
+
results: []
|
| 21 |
+
---
|
| 22 |
+
|
| 23 |
+
# sarmae-vit-large-patch16-pretrain
|
| 24 |
+
|
| 25 |
+
SARMAE (ViT-L, patch 16) encoder checkpoint converted to native Hugging Face Transformers format.
|
| 26 |
+
|
| 27 |
+
SARMAE is a Noise-Aware Masked Autoencoder for self-supervised SAR representation learning, pretrained on [SAR-1M](https://huggingface.co/datasets/Wenquandan777/SAR-1M) with Speckle-Aware Representation Enhancement (SARE) and Semantic Anchor Representation Constraint (SARC).
|
| 28 |
+
|
| 29 |
+
- **Paper:** [2512.16635](https://arxiv.org/abs/2512.16635)
|
| 30 |
+
- **Legacy weights:** [Wenquandan777/SARMAE](https://huggingface.co/Wenquandan777/SARMAE)
|
| 31 |
+
- **Stage:** `pretrain`
|
| 32 |
+
- **Input:** 3 x 224 x 224 (single-channel SAR is repeated to 3 channels)
|
| 33 |
+
- **Architecture:** 24 layers, hidden size 1024, 16 heads
|
| 34 |
+
|
| 35 |
+
## Model specifications
|
| 36 |
+
|
| 37 |
+
| Property | Value |
|
| 38 |
+
|----------|-------|
|
| 39 |
+
| Model type | `sarmae` |
|
| 40 |
+
| Backbone | ViT-L |
|
| 41 |
+
| Patch size | 16 |
|
| 42 |
+
| Image size | 224 |
|
| 43 |
+
| Hidden size | 1024 |
|
| 44 |
+
| Layers | 24 |
|
| 45 |
+
| Attention heads | 16 |
|
| 46 |
+
| Global pooling | `True` |
|
| 47 |
+
| Normalization mean | `[0.485, 0.456, 0.406]` |
|
| 48 |
+
| Normalization std | `[0.229, 0.224, 0.225]` |
|
| 49 |
+
|
| 50 |
+
## Intended use
|
| 51 |
+
|
| 52 |
+
- SAR image feature extraction for downstream classification, detection, and segmentation
|
| 53 |
+
- Initializing OpenMMLab backbones (`mmrotate`, `mmseg`) after weight porting
|
| 54 |
+
- Research and non-commercial use under [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/)
|
| 55 |
+
|
| 56 |
+
## Quick start
|
| 57 |
+
|
| 58 |
+
Install dependencies:
|
| 59 |
+
|
| 60 |
+
```bash
|
| 61 |
+
pip install transformers timm torch torchvision safetensors
|
| 62 |
+
```
|
| 63 |
+
|
| 64 |
+
### Feature extraction with `transformers.pipeline`
|
| 65 |
+
|
| 66 |
+
```python
|
| 67 |
+
from transformers import pipeline
|
| 68 |
+
|
| 69 |
+
pipe = pipeline(
|
| 70 |
+
task="image-feature-extraction",
|
| 71 |
+
model="BiliSakura/SARMAE-transformers", subfolder="vit-large-patch16-pretrain",
|
| 72 |
+
trust_remote_code=True,
|
| 73 |
+
)
|
| 74 |
+
features = pipe(sar_image, pool=True, return_tensors=True)
|
| 75 |
+
print(features.shape)
|
| 76 |
+
```
|
| 77 |
+
|
| 78 |
+
### Direct model loading
|
| 79 |
+
|
| 80 |
+
```python
|
| 81 |
+
from transformers import AutoModel, AutoImageProcessor
|
| 82 |
+
|
| 83 |
+
model = AutoModel.from_pretrained("BiliSakura/SARMAE-transformers", subfolder="vit-large-patch16-pretrain", trust_remote_code=True)
|
| 84 |
+
processor = AutoImageProcessor.from_pretrained("BiliSakura/SARMAE-transformers", subfolder="vit-large-patch16-pretrain", trust_remote_code=True)
|
| 85 |
+
|
| 86 |
+
inputs = processor(images=sar_image, return_tensors="pt")
|
| 87 |
+
outputs = model(**inputs)
|
| 88 |
+
pooled_features = outputs.pooler_output
|
| 89 |
+
```
|
| 90 |
+
|
| 91 |
+
### Local checkout
|
| 92 |
+
|
| 93 |
+
```python
|
| 94 |
+
pipe = pipeline(
|
| 95 |
+
task="image-feature-extraction",
|
| 96 |
+
model="./sarmae-vit-large-patch16-pretrain",
|
| 97 |
+
trust_remote_code=True,
|
| 98 |
+
)
|
| 99 |
+
```
|
| 100 |
+
|
| 101 |
+
## Preprocessing
|
| 102 |
+
|
| 103 |
+
- Resize to 224x224
|
| 104 |
+
- Scale pixel values to `[0, 1]` (`rescale_factor=1/255`)
|
| 105 |
+
- Repeat grayscale SAR to 3 channels when `repeat_grayscale_channels=true`
|
| 106 |
+
- Normalize with ImageNet mean/std (same as SARMAE fine-tuning code)
|
| 107 |
+
|
| 108 |
+
## Training data
|
| 109 |
+
|
| 110 |
+
Pretrained on **SAR-1M**, a million-scale SAR dataset with paired optical anchors for a subset of samples.
|
| 111 |
+
|
| 112 |
+
## Citation
|
| 113 |
+
|
| 114 |
+
```bibtex
|
| 115 |
+
@misc{liu2025sarmaemaskedautoencodersar,
|
| 116 |
+
title={SARMAE: Masked Autoencoder for SAR Representation Learning},
|
| 117 |
+
author={Danxu Liu and Di Wang and Hebaixu Wang and Haoyang Chen and Wentao Jiang and Yilin Cheng and Haonan Guo and Wei Cui and Jing Zhang},
|
| 118 |
+
year={2025},
|
| 119 |
+
eprint={2512.16635},
|
| 120 |
+
archivePrefix={arXiv},
|
| 121 |
+
primaryClass={cs.CV},
|
| 122 |
+
url={https://arxiv.org/abs/2512.16635},
|
| 123 |
+
}
|
| 124 |
+
```
|
| 125 |
+
|
| 126 |
+
## License
|
| 127 |
+
|
| 128 |
+
This model is released under [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/).
|
vit-large-patch16-pretrain/config.json
ADDED
|
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"SarmaeModel"
|
| 4 |
+
],
|
| 5 |
+
"attention_probs_dropout_prob": 0.0,
|
| 6 |
+
"checkpoint_stage": "pretrain",
|
| 7 |
+
"dtype": "float32",
|
| 8 |
+
"global_pool": true,
|
| 9 |
+
"hidden_act": "gelu",
|
| 10 |
+
"hidden_dropout_prob": 0.0,
|
| 11 |
+
"hidden_size": 1024,
|
| 12 |
+
"id2label": {},
|
| 13 |
+
"image_mean": [
|
| 14 |
+
0.485,
|
| 15 |
+
0.456,
|
| 16 |
+
0.406
|
| 17 |
+
],
|
| 18 |
+
"image_size": 224,
|
| 19 |
+
"image_std": [
|
| 20 |
+
0.229,
|
| 21 |
+
0.224,
|
| 22 |
+
0.225
|
| 23 |
+
],
|
| 24 |
+
"initializer_range": 0.02,
|
| 25 |
+
"intermediate_size": 4096,
|
| 26 |
+
"label2id": {},
|
| 27 |
+
"layer_norm_eps": 1e-06,
|
| 28 |
+
"mlp_ratio": 4.0,
|
| 29 |
+
"model_type": "sarmae",
|
| 30 |
+
"num_attention_heads": 16,
|
| 31 |
+
"num_channels": 3,
|
| 32 |
+
"num_hidden_layers": 24,
|
| 33 |
+
"patch_size": 16,
|
| 34 |
+
"qkv_bias": true,
|
| 35 |
+
"repeat_grayscale_channels": true,
|
| 36 |
+
"transformers_version": "5.0.0",
|
| 37 |
+
"auto_map": {
|
| 38 |
+
"AutoConfig": "modeling_sarmae.SarmaeConfig",
|
| 39 |
+
"AutoModel": "modeling_sarmae.SarmaeModel",
|
| 40 |
+
"AutoModelForImageClassification": "modeling_sarmae.SarmaeForImageClassification"
|
| 41 |
+
},
|
| 42 |
+
"custom_pipelines": {
|
| 43 |
+
"sarmae-feature-extraction": {
|
| 44 |
+
"impl": "pipeline_sarmae.SarmaeImageFeatureExtractionPipeline",
|
| 45 |
+
"pt": [
|
| 46 |
+
"AutoModel"
|
| 47 |
+
]
|
| 48 |
+
}
|
| 49 |
+
}
|
| 50 |
+
}
|
vit-large-patch16-pretrain/image_processing_sarmae.py
ADDED
|
@@ -0,0 +1,164 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 SARMAE Authors and The HuggingFace Inc. team.
|
| 2 |
+
"""Image processor for SARMAE models (self-contained for trust_remote_code)."""
|
| 3 |
+
|
| 4 |
+
from typing import Optional, Union
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
|
| 8 |
+
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
|
| 9 |
+
from transformers.image_transforms import resize, to_channel_dimension_format
|
| 10 |
+
from transformers.image_utils import (
|
| 11 |
+
ChannelDimension,
|
| 12 |
+
ImageInput,
|
| 13 |
+
PILImageResampling,
|
| 14 |
+
infer_channel_dimension_format,
|
| 15 |
+
to_numpy_array,
|
| 16 |
+
valid_images,
|
| 17 |
+
validate_preprocess_arguments,
|
| 18 |
+
)
|
| 19 |
+
from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
logger = logging.get_logger(__name__)
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def _repeat_grayscale_channels(image: np.ndarray, target_channels: int, input_data_format: ChannelDimension) -> np.ndarray:
|
| 26 |
+
if input_data_format == ChannelDimension.FIRST:
|
| 27 |
+
num_channels = image.shape[0]
|
| 28 |
+
if num_channels == target_channels:
|
| 29 |
+
return image
|
| 30 |
+
if num_channels == 1:
|
| 31 |
+
return np.repeat(image, target_channels, axis=0)
|
| 32 |
+
return image[:target_channels]
|
| 33 |
+
num_channels = image.shape[-1]
|
| 34 |
+
if num_channels == target_channels:
|
| 35 |
+
return image
|
| 36 |
+
if num_channels == 1:
|
| 37 |
+
return np.repeat(image, target_channels, axis=-1)
|
| 38 |
+
return image[..., :target_channels]
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def _prepare_image_batch(images: ImageInput) -> list:
|
| 42 |
+
if isinstance(images, np.ndarray):
|
| 43 |
+
images = [images]
|
| 44 |
+
elif not isinstance(images, (list, tuple)):
|
| 45 |
+
images = [images]
|
| 46 |
+
|
| 47 |
+
prepared = []
|
| 48 |
+
for image in images:
|
| 49 |
+
array = to_numpy_array(image)
|
| 50 |
+
if array.ndim == 2:
|
| 51 |
+
array = np.expand_dims(array, axis=-1)
|
| 52 |
+
prepared.append(array)
|
| 53 |
+
return prepared
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
class SarmaeImageProcessor(BaseImageProcessor):
|
| 57 |
+
model_input_names = ["pixel_values"]
|
| 58 |
+
|
| 59 |
+
def __init__(
|
| 60 |
+
self,
|
| 61 |
+
do_resize: bool = True,
|
| 62 |
+
size: Optional[dict[str, int]] = None,
|
| 63 |
+
resample: PILImageResampling = PILImageResampling.BILINEAR,
|
| 64 |
+
do_rescale: bool = True,
|
| 65 |
+
rescale_factor: float = 1 / 255.0,
|
| 66 |
+
do_normalize: bool = True,
|
| 67 |
+
image_mean: Optional[Union[float, list[float]]] = None,
|
| 68 |
+
image_std: Optional[Union[float, list[float]]] = None,
|
| 69 |
+
do_convert_rgb: bool = False,
|
| 70 |
+
repeat_grayscale_channels: bool = True,
|
| 71 |
+
**kwargs,
|
| 72 |
+
):
|
| 73 |
+
super().__init__(**kwargs)
|
| 74 |
+
size = size if size is not None else {"height": 224, "width": 224}
|
| 75 |
+
self.do_resize = do_resize
|
| 76 |
+
self.size = size
|
| 77 |
+
self.resample = resample
|
| 78 |
+
self.do_rescale = do_rescale
|
| 79 |
+
self.rescale_factor = rescale_factor
|
| 80 |
+
self.do_normalize = do_normalize
|
| 81 |
+
self.image_mean = image_mean
|
| 82 |
+
self.image_std = image_std
|
| 83 |
+
self.do_convert_rgb = do_convert_rgb
|
| 84 |
+
self.repeat_grayscale_channels = repeat_grayscale_channels
|
| 85 |
+
|
| 86 |
+
@filter_out_non_signature_kwargs()
|
| 87 |
+
def preprocess(
|
| 88 |
+
self,
|
| 89 |
+
images: ImageInput,
|
| 90 |
+
do_resize: Optional[bool] = None,
|
| 91 |
+
size: Optional[dict[str, int]] = None,
|
| 92 |
+
resample: Optional[PILImageResampling] = None,
|
| 93 |
+
do_rescale: Optional[bool] = None,
|
| 94 |
+
rescale_factor: Optional[float] = None,
|
| 95 |
+
do_normalize: Optional[bool] = None,
|
| 96 |
+
image_mean: Optional[Union[float, list[float]]] = None,
|
| 97 |
+
image_std: Optional[Union[float, list[float]]] = None,
|
| 98 |
+
return_tensors: Optional[Union[str, TensorType]] = None,
|
| 99 |
+
data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
|
| 100 |
+
input_data_format: Optional[Union[str, ChannelDimension]] = None,
|
| 101 |
+
do_convert_rgb: Optional[bool] = None,
|
| 102 |
+
repeat_grayscale_channels: Optional[bool] = None,
|
| 103 |
+
):
|
| 104 |
+
do_resize = do_resize if do_resize is not None else self.do_resize
|
| 105 |
+
size = get_size_dict(size if size is not None else self.size, default_to_square=True)
|
| 106 |
+
resample = resample if resample is not None else self.resample
|
| 107 |
+
do_rescale = do_rescale if do_rescale is not None else self.do_rescale
|
| 108 |
+
rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
|
| 109 |
+
do_normalize = do_normalize if do_normalize is not None else self.do_normalize
|
| 110 |
+
image_mean = image_mean if image_mean is not None else self.image_mean
|
| 111 |
+
image_std = image_std if image_std is not None else self.image_std
|
| 112 |
+
do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
|
| 113 |
+
repeat_grayscale_channels = (
|
| 114 |
+
repeat_grayscale_channels if repeat_grayscale_channels is not None else self.repeat_grayscale_channels
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
if do_normalize and (image_mean is None or image_std is None):
|
| 118 |
+
raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
|
| 119 |
+
|
| 120 |
+
images = _prepare_image_batch(images)
|
| 121 |
+
if not valid_images(images):
|
| 122 |
+
raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
|
| 123 |
+
|
| 124 |
+
validate_preprocess_arguments(
|
| 125 |
+
do_rescale=do_rescale,
|
| 126 |
+
rescale_factor=rescale_factor,
|
| 127 |
+
do_normalize=do_normalize,
|
| 128 |
+
image_mean=image_mean,
|
| 129 |
+
image_std=image_std,
|
| 130 |
+
do_resize=do_resize,
|
| 131 |
+
size=size,
|
| 132 |
+
resample=resample,
|
| 133 |
+
)
|
| 134 |
+
|
| 135 |
+
processed_images = []
|
| 136 |
+
for image in images:
|
| 137 |
+
image = to_numpy_array(image)
|
| 138 |
+
if do_convert_rgb:
|
| 139 |
+
image = self._convert_image_to_rgb(image)
|
| 140 |
+
if input_data_format is None:
|
| 141 |
+
try:
|
| 142 |
+
input_data_format = infer_channel_dimension_format(image)
|
| 143 |
+
except ValueError:
|
| 144 |
+
input_data_format = ChannelDimension.LAST
|
| 145 |
+
if repeat_grayscale_channels:
|
| 146 |
+
image = _repeat_grayscale_channels(image, target_channels=3, input_data_format=input_data_format)
|
| 147 |
+
if do_resize:
|
| 148 |
+
image = resize(
|
| 149 |
+
image,
|
| 150 |
+
size=(size["height"], size["width"]),
|
| 151 |
+
resample=resample,
|
| 152 |
+
input_data_format=input_data_format,
|
| 153 |
+
)
|
| 154 |
+
if do_rescale:
|
| 155 |
+
image = image * rescale_factor
|
| 156 |
+
if do_normalize:
|
| 157 |
+
image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
|
| 158 |
+
image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
|
| 159 |
+
processed_images.append(image)
|
| 160 |
+
|
| 161 |
+
return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
|
| 162 |
+
|
| 163 |
+
|
| 164 |
+
__all__ = ["SarmaeImageProcessor"]
|
vit-large-patch16-pretrain/model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e80a1a46202d5f2cfd64955101f2fc81fbbe476954ce9a52fd597c95bb419d84
|
| 3 |
+
size 1213234544
|
vit-large-patch16-pretrain/modeling_sarmae.py
ADDED
|
@@ -0,0 +1,200 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 SARMAE Authors and The HuggingFace Inc. team.
|
| 2 |
+
"""Self-contained SARMAE model and config for trust_remote_code loading."""
|
| 3 |
+
|
| 4 |
+
from __future__ import annotations
|
| 5 |
+
|
| 6 |
+
from functools import partial
|
| 7 |
+
from typing import Optional
|
| 8 |
+
|
| 9 |
+
import numpy as np
|
| 10 |
+
import torch
|
| 11 |
+
from timm.models.vision_transformer import Block, PatchEmbed
|
| 12 |
+
from torch import nn
|
| 13 |
+
|
| 14 |
+
from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
|
| 15 |
+
from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
|
| 16 |
+
from transformers.modeling_utils import PreTrainedModel
|
| 17 |
+
from transformers.processing_utils import Unpack
|
| 18 |
+
from transformers.utils import TransformersKwargs, logging
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
logger = logging.get_logger(__name__)
|
| 22 |
+
|
| 23 |
+
IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
| 24 |
+
IMAGENET_STD = [0.229, 0.224, 0.225]
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> np.ndarray:
|
| 28 |
+
grid_h = np.arange(grid_size, dtype=np.float32)
|
| 29 |
+
grid_w = np.arange(grid_size, dtype=np.float32)
|
| 30 |
+
grid = np.meshgrid(grid_w, grid_h)
|
| 31 |
+
grid = np.stack(grid, axis=0).reshape([2, 1, grid_size, grid_size])
|
| 32 |
+
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
|
| 33 |
+
if cls_token:
|
| 34 |
+
pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
|
| 35 |
+
return pos_embed
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
|
| 39 |
+
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
|
| 40 |
+
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
|
| 41 |
+
return np.concatenate([emb_h, emb_w], axis=1)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
|
| 45 |
+
omega = np.arange(embed_dim // 2, dtype=np.float32)
|
| 46 |
+
omega /= embed_dim / 2.0
|
| 47 |
+
omega = 1.0 / 10000**omega
|
| 48 |
+
pos = pos.reshape(-1)
|
| 49 |
+
out = np.einsum("m,d->md", pos, omega)
|
| 50 |
+
return np.concatenate([np.sin(out), np.cos(out)], axis=1)
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
class SarmaeConfig(PreTrainedConfig):
|
| 54 |
+
model_type = "sarmae"
|
| 55 |
+
|
| 56 |
+
def __init__(
|
| 57 |
+
self,
|
| 58 |
+
hidden_size: int = 768,
|
| 59 |
+
num_hidden_layers: int = 12,
|
| 60 |
+
num_attention_heads: int = 12,
|
| 61 |
+
intermediate_size: int | None = None,
|
| 62 |
+
hidden_act: str = "gelu",
|
| 63 |
+
hidden_dropout_prob: float = 0.0,
|
| 64 |
+
attention_probs_dropout_prob: float = 0.0,
|
| 65 |
+
initializer_range: float = 0.02,
|
| 66 |
+
layer_norm_eps: float = 1e-6,
|
| 67 |
+
image_size: int = 224,
|
| 68 |
+
patch_size: int = 16,
|
| 69 |
+
num_channels: int = 3,
|
| 70 |
+
qkv_bias: bool = True,
|
| 71 |
+
mlp_ratio: float = 4.0,
|
| 72 |
+
global_pool: bool = True,
|
| 73 |
+
repeat_grayscale_channels: bool = True,
|
| 74 |
+
checkpoint_stage: str = "pretrain",
|
| 75 |
+
image_mean: list[float] | None = None,
|
| 76 |
+
image_std: list[float] | None = None,
|
| 77 |
+
num_labels: int = 0,
|
| 78 |
+
**kwargs,
|
| 79 |
+
):
|
| 80 |
+
super().__init__(**kwargs)
|
| 81 |
+
self.hidden_size = hidden_size
|
| 82 |
+
self.num_hidden_layers = num_hidden_layers
|
| 83 |
+
self.num_attention_heads = num_attention_heads
|
| 84 |
+
self.hidden_act = hidden_act
|
| 85 |
+
self.hidden_dropout_prob = hidden_dropout_prob
|
| 86 |
+
self.attention_probs_dropout_prob = attention_probs_dropout_prob
|
| 87 |
+
self.initializer_range = initializer_range
|
| 88 |
+
self.layer_norm_eps = layer_norm_eps
|
| 89 |
+
self.image_size = image_size
|
| 90 |
+
self.patch_size = patch_size
|
| 91 |
+
self.num_channels = num_channels
|
| 92 |
+
self.qkv_bias = qkv_bias
|
| 93 |
+
self.mlp_ratio = mlp_ratio
|
| 94 |
+
self.global_pool = global_pool
|
| 95 |
+
self.repeat_grayscale_channels = repeat_grayscale_channels
|
| 96 |
+
self.checkpoint_stage = checkpoint_stage
|
| 97 |
+
self.num_labels = num_labels
|
| 98 |
+
self.intermediate_size = int(hidden_size * mlp_ratio) if intermediate_size is None else intermediate_size
|
| 99 |
+
self.image_mean = image_mean if image_mean is not None else IMAGENET_MEAN
|
| 100 |
+
self.image_std = image_std if image_std is not None else IMAGENET_STD
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
class SarmaePreTrainedModel(PreTrainedModel):
|
| 104 |
+
config_class = SarmaeConfig
|
| 105 |
+
config: SarmaeConfig
|
| 106 |
+
base_model_prefix = "sarmae"
|
| 107 |
+
main_input_name = "pixel_values"
|
| 108 |
+
input_modalities = ("image",)
|
| 109 |
+
supports_gradient_checkpointing = True
|
| 110 |
+
_no_split_modules = ["Block"]
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
class SarmaeModel(SarmaePreTrainedModel):
|
| 114 |
+
def __init__(self, config: SarmaeConfig, add_pooling_layer: bool = True):
|
| 115 |
+
super().__init__(config)
|
| 116 |
+
self.config = config
|
| 117 |
+
self.add_pooling_layer = add_pooling_layer
|
| 118 |
+
|
| 119 |
+
image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0]
|
| 120 |
+
norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
|
| 121 |
+
self.patch_embed = PatchEmbed(image_size, config.patch_size, config.num_channels, config.hidden_size)
|
| 122 |
+
self.num_patches = self.patch_embed.num_patches
|
| 123 |
+
self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
|
| 124 |
+
self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, config.hidden_size))
|
| 125 |
+
pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.num_patches**0.5), cls_token=True)
|
| 126 |
+
self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
|
| 127 |
+
self.blocks = nn.ModuleList([
|
| 128 |
+
Block(config.hidden_size, config.num_attention_heads, config.mlp_ratio, qkv_bias=config.qkv_bias, norm_layer=norm_layer)
|
| 129 |
+
for _ in range(config.num_hidden_layers)
|
| 130 |
+
])
|
| 131 |
+
self.global_pool = config.global_pool
|
| 132 |
+
if self.global_pool:
|
| 133 |
+
self.fc_norm = norm_layer(config.hidden_size)
|
| 134 |
+
self.norm = None
|
| 135 |
+
else:
|
| 136 |
+
self.fc_norm = None
|
| 137 |
+
self.norm = norm_layer(config.hidden_size)
|
| 138 |
+
self.post_init()
|
| 139 |
+
|
| 140 |
+
def forward(
|
| 141 |
+
self,
|
| 142 |
+
pixel_values: Optional[torch.Tensor] = None,
|
| 143 |
+
return_dict: Optional[bool] = None,
|
| 144 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 145 |
+
) -> BaseModelOutputWithPooling:
|
| 146 |
+
if pixel_values is None:
|
| 147 |
+
raise ValueError("You must specify `pixel_values`")
|
| 148 |
+
pixel_values = pixel_values.to(dtype=self.dtype)
|
| 149 |
+
if return_dict is None:
|
| 150 |
+
return_dict = self.config.use_return_dict
|
| 151 |
+
|
| 152 |
+
batch_size = pixel_values.shape[0]
|
| 153 |
+
patch_tokens = self.patch_embed(pixel_values)
|
| 154 |
+
cls_tokens = self.cls_token.expand(batch_size, -1, -1)
|
| 155 |
+
hidden_states = torch.cat((cls_tokens, patch_tokens), dim=1) + self.pos_embed
|
| 156 |
+
for block in self.blocks:
|
| 157 |
+
hidden_states = block(hidden_states)
|
| 158 |
+
if self.global_pool:
|
| 159 |
+
pooled_output = self.fc_norm(hidden_states[:, 1:, :].mean(dim=1))
|
| 160 |
+
else:
|
| 161 |
+
hidden_states = self.norm(hidden_states)
|
| 162 |
+
pooled_output = hidden_states[:, 0]
|
| 163 |
+
if not self.add_pooling_layer:
|
| 164 |
+
pooled_output = None
|
| 165 |
+
if not return_dict:
|
| 166 |
+
return (hidden_states, pooled_output)
|
| 167 |
+
return BaseModelOutputWithPooling(last_hidden_state=hidden_states, pooler_output=pooled_output)
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class SarmaeForImageClassification(SarmaePreTrainedModel):
|
| 171 |
+
def __init__(self, config: SarmaeConfig):
|
| 172 |
+
super().__init__(config)
|
| 173 |
+
self.sarmae = SarmaeModel(config, add_pooling_layer=True)
|
| 174 |
+
self.classifier = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
|
| 175 |
+
self.post_init()
|
| 176 |
+
|
| 177 |
+
def forward(
|
| 178 |
+
self,
|
| 179 |
+
pixel_values: Optional[torch.Tensor] = None,
|
| 180 |
+
labels: Optional[torch.Tensor] = None,
|
| 181 |
+
return_dict: Optional[bool] = None,
|
| 182 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 183 |
+
) -> ImageClassifierOutput:
|
| 184 |
+
outputs = self.sarmae(pixel_values=pixel_values, return_dict=True, **kwargs)
|
| 185 |
+
logits = self.classifier(outputs.pooler_output)
|
| 186 |
+
loss = None
|
| 187 |
+
if labels is not None:
|
| 188 |
+
loss = self.loss_function(labels, logits, self.config, **kwargs)
|
| 189 |
+
if not return_dict:
|
| 190 |
+
output = (logits,) + outputs[1:]
|
| 191 |
+
return ((loss,) + output) if loss is not None else output
|
| 192 |
+
return ImageClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
__all__ = [
|
| 196 |
+
"SarmaeConfig",
|
| 197 |
+
"SarmaeForImageClassification",
|
| 198 |
+
"SarmaeModel",
|
| 199 |
+
"SarmaePreTrainedModel",
|
| 200 |
+
]
|
vit-large-patch16-pretrain/pipeline_sarmae.py
ADDED
|
@@ -0,0 +1,51 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 SARMAE Authors and The HuggingFace Inc. team.
|
| 2 |
+
"""SARMAE image feature extraction pipeline (self-contained for trust_remote_code)."""
|
| 3 |
+
|
| 4 |
+
from typing import Any, Union
|
| 5 |
+
|
| 6 |
+
from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
|
| 7 |
+
from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
|
| 8 |
+
from transformers.utils import add_end_docstrings, is_vision_available
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
if is_vision_available():
|
| 12 |
+
from transformers.image_utils import load_image
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
@add_end_docstrings(
|
| 16 |
+
build_pipeline_init_args(has_image_processor=True),
|
| 17 |
+
"""
|
| 18 |
+
pool (`bool`, *optional*, defaults to `False`):
|
| 19 |
+
Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
|
| 20 |
+
""",
|
| 21 |
+
)
|
| 22 |
+
class SarmaeImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
|
| 23 |
+
def _sanitize_parameters(
|
| 24 |
+
self,
|
| 25 |
+
image_processor_kwargs=None,
|
| 26 |
+
return_tensors=None,
|
| 27 |
+
pool=None,
|
| 28 |
+
**kwargs,
|
| 29 |
+
):
|
| 30 |
+
preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
|
| 31 |
+
if "timeout" in kwargs:
|
| 32 |
+
preprocess_params["timeout"] = kwargs["timeout"]
|
| 33 |
+
postprocess_params = {}
|
| 34 |
+
if pool is not None:
|
| 35 |
+
postprocess_params["pool"] = pool
|
| 36 |
+
if return_tensors is not None:
|
| 37 |
+
postprocess_params["return_tensors"] = return_tensors
|
| 38 |
+
return preprocess_params, {}, postprocess_params
|
| 39 |
+
|
| 40 |
+
def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
|
| 41 |
+
if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
|
| 42 |
+
image = load_image(image, timeout=timeout)
|
| 43 |
+
model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
|
| 44 |
+
model_inputs = model_inputs.to(self.dtype)
|
| 45 |
+
return model_inputs
|
| 46 |
+
|
| 47 |
+
def __call__(self, *args: Union[str, Any, list[Any]], **kwargs: Any) -> list[Any]:
|
| 48 |
+
return super().__call__(*args, **kwargs)
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
__all__ = ["SarmaeImageFeatureExtractionPipeline"]
|
vit-large-patch16-pretrain/preprocessor_config.json
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"image_processor_type": "SarmaeImageProcessor",
|
| 3 |
+
"size": {
|
| 4 |
+
"height": 224,
|
| 5 |
+
"width": 224
|
| 6 |
+
},
|
| 7 |
+
"do_resize": true,
|
| 8 |
+
"do_rescale": true,
|
| 9 |
+
"rescale_factor": 0.00392156862745098,
|
| 10 |
+
"do_normalize": true,
|
| 11 |
+
"do_convert_rgb": false,
|
| 12 |
+
"repeat_grayscale_channels": true,
|
| 13 |
+
"image_mean": [
|
| 14 |
+
0.485,
|
| 15 |
+
0.456,
|
| 16 |
+
0.406
|
| 17 |
+
],
|
| 18 |
+
"image_std": [
|
| 19 |
+
0.229,
|
| 20 |
+
0.224,
|
| 21 |
+
0.225
|
| 22 |
+
],
|
| 23 |
+
"auto_map": {
|
| 24 |
+
"AutoImageProcessor": "image_processing_sarmae.SarmaeImageProcessor"
|
| 25 |
+
}
|
| 26 |
+
}
|