# Copyright 2024 SatMAE++ 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.
"""SatMAE++ 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 SatMAEppImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
    """
    SatMAE++ image feature extraction pipeline.

    This pipeline wraps [`SatMAEppModel`] for RGB and multispectral satellite feature extraction.

    Example:

    ```python
    >>> from transformers import pipeline

    >>> pipe = pipeline(
    ...     task="image-feature-extraction",
    ...     model="/path/to/satmae-pp-vit-large-patch16-fmow-rgb-finetune",
    ...     trust_remote_code=True,
    ... )
    >>> features = pipe(image_array, pool=True, return_tensors=True)
    ```
    """

    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__ = ["SatMAEppImageFeatureExtractionPipeline"]