File size: 3,827 Bytes
b631fc1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | # Copyright 2025 The Galileo Authors and The HuggingFace Inc. team.
"""Custom Galileo feature extraction pipeline for trust_remote_code loading."""
from __future__ import annotations
from typing import Any, Union
import torch
from transformers.pipelines.base import GenericTensor, Pipeline, build_pipeline_init_args
from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
from transformers.utils import add_end_docstrings
@add_end_docstrings(
build_pipeline_init_args(has_processor=True),
"""
s1 (`torch.Tensor` or `numpy.ndarray`, *optional*):
Sentinel-1 tensor with shape `(H, W, T, C)` or `(H, W, C)`.
s2 (`torch.Tensor` or `numpy.ndarray`, *optional*):
Sentinel-2 tensor with shape `(H, W, T, C)` or `(H, W, C)`.
normalize (`bool`, *optional*, defaults to `True`):
Whether to apply pretraining normalization statistics.
patch_size (`int`, *optional*):
Spatial patch size for the flexible patch embedder (1-8).
pool (`bool`, *optional*, defaults to `False`):
Whether to return pooled features instead of sequence hidden states.
""",
)
class GalileoImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
"""
Galileo multimodal remote sensing feature extraction pipeline.
Load with:
```python
>>> from transformers import pipeline
>>> pipe = pipeline(
... task="galileo-feature-extraction",
... model="/path/to/galileo-nano-patch8",
... trust_remote_code=True,
... )
>>> features = pipe(s2=s2_array, pool=True, return_tensors=True)
```
"""
_load_processor = True
_load_image_processor = False
_load_feature_extractor = False
_load_tokenizer = False
def _sanitize_parameters(
self,
processor_kwargs=None,
image_processor_kwargs=None,
return_tensors=None,
pool=None,
normalize=None,
patch_size=None,
**kwargs,
):
preprocess_params = {}
if processor_kwargs is not None:
preprocess_params.update(processor_kwargs)
if image_processor_kwargs is not None:
preprocess_params.update(image_processor_kwargs)
for key in (
"s1",
"s2",
"era5",
"tc",
"viirs",
"srtm",
"dw",
"wc",
"landscan",
"latlon",
"months",
):
if key in kwargs:
preprocess_params[key] = kwargs.pop(key)
if normalize is not None:
preprocess_params["normalize"] = normalize
if patch_size is not None:
preprocess_params["patch_size"] = patch_size
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=None, timeout=None, **processor_kwargs) -> dict[str, GenericTensor]:
del image, timeout
model_inputs = self.processor(return_tensors="pt", **processor_kwargs)
model_inputs = model_inputs.to(self.dtype)
return model_inputs
def _forward(self, model_inputs):
patch_size = model_inputs.pop("patch_size", None)
if torch.is_tensor(patch_size):
patch_size = int(patch_size.item())
return self.model(patch_size=patch_size, **model_inputs)
def __call__(self, *args: Union[Any, list[Any]], **kwargs: Any) -> list[Any]:
if not args:
args = (None,)
return super().__call__(*args, **kwargs)
__all__ = ["GalileoImageFeatureExtractionPipeline"]
|