| |
| """DeCUR image feature extraction pipeline.""" |
|
|
| from typing import Any, Union |
|
|
| from transformers.pipelines.base import GenericTensor, build_pipeline_init_args |
| from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline |
| 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 or not to return the pooled output. If `False`, the model will return the raw hidden states. |
| """, |
| ) |
| class DeCURImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): |
|
|
| def _sanitize_parameters( |
| self, |
| image_processor_kwargs=None, |
| return_tensors=None, |
| pool=None, |
| **kwargs, |
| ): |
| preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs) |
| if "timeout" in kwargs: |
| preprocess_params["timeout"] = kwargs["timeout"] |
|
|
| 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 __call__( |
| self, |
| *args: Union[str, Any, list[Any]], |
| **kwargs: Any, |
| ) -> list[Any]: |
| return super().__call__(*args, **kwargs) |
|
|
|
|
| __all__ = ["DeCURImageFeatureExtractionPipeline"] |
|
|
