| |
| """Custom CROMA dual-modality feature extraction pipeline.""" |
|
|
| from typing import Any, Optional, Union |
|
|
| from transformers.image_utils import load_image |
| from transformers.pipelines.base import GenericTensor, Pipeline |
| from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline |
|
|
|
|
| class CROMAImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline): |
| """SAR-optical feature extraction for CROMA.""" |
|
|
| def _sanitize_parameters( |
| self, |
| image_processor_kwargs=None, |
| return_tensors=None, |
| pool=None, |
| sar_images=None, |
| optical_images=None, |
| use_8_bit=None, |
| output_mode=None, |
| **kwargs, |
| ): |
| preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs) |
|
|
| if sar_images is not None: |
| preprocess_params["sar_images"] = sar_images |
| if optical_images is not None: |
| preprocess_params["optical_images"] = optical_images |
| if use_8_bit is not None: |
| preprocess_params["use_8_bit"] = use_8_bit |
| 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 |
| if output_mode is not None: |
| postprocess_params["output_mode"] = output_mode |
|
|
| return preprocess_params, {}, postprocess_params |
|
|
| def preprocess( |
| self, |
| image=None, |
| sar_images=None, |
| optical_images=None, |
| timeout=None, |
| **image_processor_kwargs, |
| ) -> dict[str, GenericTensor]: |
| if sar_images is None: |
| sar_images = image_processor_kwargs.pop("sar_images", None) |
| if optical_images is None: |
| optical_images = image_processor_kwargs.pop("optical_images", None) |
|
|
| if sar_images is not None and not isinstance(sar_images, (list, tuple)) and not hasattr(sar_images, "shape"): |
| sar_images = load_image(sar_images, timeout=timeout) |
| if optical_images is not None and not isinstance(optical_images, (list, tuple)) and not hasattr( |
| optical_images, "shape" |
| ): |
| optical_images = load_image(optical_images, timeout=timeout) |
| if image is not None and optical_images is None and sar_images is None: |
| if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"): |
| image = load_image(image, timeout=timeout) |
| optical_images = image |
|
|
| model_inputs = self.image_processor.preprocess( |
| images=optical_images if optical_images is not None else sar_images, |
| sar_images=sar_images, |
| optical_images=optical_images, |
| return_tensors="pt", |
| **image_processor_kwargs, |
| ) |
| model_inputs = model_inputs.to(self.dtype) |
| return model_inputs |
|
|
| def postprocess(self, model_outputs, pool=None, return_tensors=False, output_mode="primary"): |
| pool = pool if pool is not None else False |
| output_mode = output_mode or "primary" |
|
|
| if output_mode == "sar": |
| outputs = model_outputs["sar_pooler_output"] if pool else model_outputs["sar_hidden_states"] |
| elif output_mode == "optical": |
| outputs = model_outputs["optical_pooler_output"] if pool else model_outputs["optical_hidden_states"] |
| elif output_mode == "joint": |
| outputs = model_outputs["joint_pooler_output"] if pool else model_outputs["joint_hidden_states"] |
| elif pool: |
| outputs = model_outputs["pooler_output"] |
| else: |
| outputs = model_outputs[0] |
|
|
| if return_tensors: |
| return outputs |
| return outputs.tolist() |
|
|
| def __call__( |
| self, |
| images: Optional[Union[str, Any, list[Any]]] = None, |
| sar_images: Optional[Union[str, Any, list[Any]]] = None, |
| optical_images: Optional[Union[str, Any, list[Any]]] = None, |
| **kwargs: Any, |
| ) -> list[Any]: |
| if images is None and sar_images is None and optical_images is None: |
| raise ValueError("Provide at least one of `images`, `sar_images`, or `optical_images`.") |
|
|
| if images is None: |
| images = optical_images if optical_images is not None else sar_images |
|
|
| return Pipeline.__call__(self, images, sar_images=sar_images, optical_images=optical_images, **kwargs) |
|
|
|
|
| __all__ = ["CROMAImageFeatureExtractionPipeline"] |
|
|
