| |
| """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"] |
|
|
