File size: 4,466 Bytes
d324dd8 | 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 117 118 119 120 121 122 123 | # Copyright (c) IBM Corp. 2024. All rights reserved.
# Copyright 2024 Prithvi-EO-2.0 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.
"""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"]
|