| |
| |
| |
| |
| |
| """Prithvi-EO-2.0 spatiotemporal feature extraction pipeline.""" |
|
|
| from typing import Any, Union |
|
|
| from transformers.pipelines.base import GenericTensor, Pipeline, build_pipeline_init_args |
| 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 to return the CLS token (`True`) or the full token sequence (`False`). |
| temporal_coords (`list[list[float]]`, *optional*): |
| Year and Julian day for each frame, shape `(T, 2)`. |
| location_coords (`list[float]`, *optional*): |
| Center latitude and longitude, shape `(2,)`. |
| """, |
| ) |
| class PrithviFeatureExtractionPipeline(Pipeline): |
| """ |
| Prithvi-EO-2.0 spatiotemporal feature extraction pipeline. |
| |
| Example: |
| |
| ```python |
| >>> from transformers import pipeline |
| >>> import numpy as np |
| |
| >>> pipe = pipeline( |
| ... task="prithvi-eo-feature-extraction", |
| ... model="prithvi-eo-v2-300m-tl", |
| ... trust_remote_code=True, |
| ... ) |
| >>> frames = [np.random.rand(224, 224, 6).astype("float32") for _ in range(4)] |
| >>> features = pipe(frames, pool=True, return_tensors=True) |
| ``` |
| """ |
|
|
| def _sanitize_parameters( |
| self, |
| image_processor_kwargs=None, |
| return_tensors=None, |
| pool=None, |
| temporal_coords=None, |
| location_coords=None, |
| **kwargs, |
| ): |
| preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs) |
| if "timeout" in kwargs: |
| preprocess_params["timeout"] = kwargs["timeout"] |
| if temporal_coords is not None: |
| preprocess_params["temporal_coords"] = temporal_coords |
| if location_coords is not None: |
| preprocess_params["location_coords"] = location_coords |
|
|
| 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 _forward(self, model_inputs): |
| temporal_coords = model_inputs.pop("temporal_coords", None) |
| location_coords = model_inputs.pop("location_coords", None) |
| return self.model( |
| temporal_coords=temporal_coords, |
| location_coords=location_coords, |
| **model_inputs, |
| ) |
|
|
| def postprocess(self, model_outputs, pool=False, return_tensors=False): |
| if pool: |
| output = model_outputs.pooler_output |
| else: |
| output = model_outputs.last_hidden_state |
|
|
| if return_tensors: |
| return output |
| return output.tolist() |
|
|
| def _is_temporal_frame_list(self, value) -> bool: |
| if not isinstance(value, list) or len(value) == 0: |
| return False |
| if not all(hasattr(item, "shape") for item in value): |
| return False |
| shapes = [tuple(item.shape) for item in value] |
| return all(len(shape) == 3 for shape in shapes) and len(set(shapes)) == 1 |
|
|
| def __call__( |
| self, |
| *args: Union[str, Any, list[Any]], |
| **kwargs: Any, |
| ) -> list[Any]: |
| if len(args) == 1 and self._is_temporal_frame_list(args[0]): |
| results = super().__call__([args[0]], **kwargs) |
| return results[0] if len(results) == 1 else results |
| results = super().__call__(*args, **kwargs) |
| if len(args) == 1 and not isinstance(args[0], list) and len(results) == 1: |
| return results[0] |
| return results |
|
|
|
|
| __all__ = ["PrithviFeatureExtractionPipeline"] |
|
|
