

# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
  import jax
  import jaxlib

__UpperCAmelCase    				=     get_logger()

__UpperCAmelCase    				=     None





class A__    ( TensorFormatter[Mapping, '''jax.Array''', Mapping] ):



 """simple docstring"""







 def __init__(		self :	Any     ,						A_ :	Optional[int]=None     ,						A_ :	Union[str, Any]=None     ,						**A_ :	str       ):






      '''simple docstring'''




      super().__init__(features=A_       )
      import jax
      from jaxlib.xla_client import Device

      if isinstance(A_     ,						A_       ):
           raise ValueError(
               F'''Expected {device} to be a `str` not {type(A_       )}, as `jaxlib.xla_extension.Device` '''
               "is not serializable neither with `pickle` nor with `dill`. Instead you can surround "
               "the device with `str()` to get its string identifier that will be internally mapped "
               "to the actual `jaxlib.xla_extension.Device`."       )
      _lowerCAmelCase   :					int        =    device if isinstance(A_     ,						A_       ) else str(jax.devices()[0]       )
      # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
      # with `pickle` nor with `dill`, so we need to use a global variable instead
      global DEVICE_MAPPING
      if DEVICE_MAPPING is None:
           _lowerCAmelCase   :					Tuple        =    self._map_devices_to_str()
      if self.device not in list(DEVICE_MAPPING.keys()       ):
           logger.warning(
               F'''Device with string identifier {self.device} not listed among the available '''
               F'''devices: {list(DEVICE_MAPPING.keys()       )}, so falling back to the default '''
               F'''device: {str(jax.devices()[0]       )}.'''       )
           _lowerCAmelCase   :					List[str]        =    str(jax.devices()[0]       )
      _lowerCAmelCase   :					str        =    jnp_array_kwargs

 @staticmethod
 def    __magic_name__       (		):






      '''simple docstring'''




      import jax

      return {str(A_       ): device for device in jax.devices()}

 def    __magic_name__       (		self :	Dict     ,						A_ :	Dict       ):






      '''simple docstring'''




      import jax
      import jax.numpy as jnp

      if isinstance(A_     ,						A_       ) and column:
           if all(
               isinstance(A_     ,						jax.Array       ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column       ):
                return jnp.stack(A_     ,						axis=0       )
      return column

 def    __magic_name__       (		self :	int     ,						A_ :	int       ):






      '''simple docstring'''




      import jax
      import jax.numpy as jnp

      if isinstance(A_     ,						(str, bytes, type(A_       ))       ):
           return value
      elif isinstance(A_     ,						(np.character, np.ndarray)       ) and np.issubdtype(value.dtype     ,						np.character       ):
           return value.tolist()

      _lowerCAmelCase   :					List[str]        =    {}

      if isinstance(A_     ,						(np.number, np.ndarray)       ) and np.issubdtype(value.dtype     ,						np.integer       ):
           # the default int precision depends on the jax config
           # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
           if jax.config.jax_enable_xaa:
                _lowerCAmelCase   :					Dict        =    {"dtype": jnp.intaa}
           else:
                _lowerCAmelCase   :					Union[str, Any]        =    {"dtype": jnp.intaa}
      elif isinstance(A_     ,						(np.number, np.ndarray)       ) and np.issubdtype(value.dtype     ,						np.floating       ):
           _lowerCAmelCase   :					Tuple        =    {"dtype": jnp.floataa}
      elif config.PIL_AVAILABLE and "PIL" in sys.modules:
           import PIL.Image

           if isinstance(A_     ,						PIL.Image.Image       ):
                _lowerCAmelCase   :					Any        =    np.asarray(A_       )

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
      global DEVICE_MAPPING
      if DEVICE_MAPPING is None:
           _lowerCAmelCase   :					Union[str, Any]        =    self._map_devices_to_str()

      with jax.default_device(DEVICE_MAPPING[self.device]       ):
           # calling jnp.array on a np.ndarray does copy the data
           # see https://github.com/google/jax/issues/4486
           return jnp.array(A_     ,						**{**default_dtype, **self.jnp_array_kwargs}       )

 def    __magic_name__       (		self :	Tuple     ,						A_ :	int       ):






      '''simple docstring'''




      import jax

      # support for torch, tf, jax etc.
      if config.TORCH_AVAILABLE and "torch" in sys.modules:
           import torch

           if isinstance(A_     ,						torch.Tensor       ):
                return self._tensorize(data_struct.detach().cpu().numpy()[()]       )
      if hasattr(A_     ,						"__array__"       ) and not isinstance(A_     ,						jax.Array       ):
           _lowerCAmelCase   :					str        =    data_struct.__array__()
      # support for nested types like struct of list of struct
      if isinstance(A_     ,						np.ndarray       ):
           if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
                return self._consolidate([self.recursive_tensorize(A_       ) for substruct in data_struct]       )
      elif isinstance(A_     ,						(list, tuple)       ):
           return self._consolidate([self.recursive_tensorize(A_       ) for substruct in data_struct]       )
      return self._tensorize(A_       )

 def    __magic_name__       (		self :	Optional[Any]     ,						A_ :	dict       ):






      '''simple docstring'''




      return map_nested(self._recursive_tensorize     ,						A_     ,						map_list=A_       )

 def    __magic_name__       (		self :	Any     ,						A_ :	pa.Table       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[Any]        =    self.numpy_arrow_extractor().extract_row(A_       )
      _lowerCAmelCase   :					Tuple        =    self.python_features_decoder.decode_row(A_       )
      return self.recursive_tensorize(A_       )

 def    __magic_name__       (		self :	List[Any]     ,						A_ :	pa.Table       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[int]        =    self.numpy_arrow_extractor().extract_column(A_       )
      _lowerCAmelCase   :					str        =    self.python_features_decoder.decode_column(A_     ,						pa_table.column_names[0]       )
      _lowerCAmelCase   :					List[str]        =    self.recursive_tensorize(A_       )
      _lowerCAmelCase   :					List[str]        =    self._consolidate(A_       )
      return column


 def    __magic_name__       (		self :	Any     ,						A_ :	pa.Table       ):






      '''simple docstring'''




      _lowerCAmelCase   :					str        =    self.numpy_arrow_extractor().extract_batch(A_       )
      _lowerCAmelCase   :					List[str]        =    self.python_features_decoder.decode_batch(A_       )
      _lowerCAmelCase   :					Dict        =    self.recursive_tensorize(A_       )
      for column_name in batch:
           _lowerCAmelCase   :					Tuple        =    self._consolidate(batch[column_name]       )
      return batch









from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import center_crop, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
  import PIL


__UpperCAmelCase    				=     logging.get_logger(__name__)





class A__    ( A ):



 """simple docstring"""







 _lowercase     :					Union[str, Any]      =						['''pixel_values''']

 def __init__(		self :	List[str]     ,						A_ :	bool = True     ,						A_ :	Dict[str, int] = None     ,						A_ :	PILImageResampling = PIL.Image.BICUBIC     ,						A_ :	bool = True     ,						A_ :	Dict[str, int] = None     ,						A_ :	Union[int, float] = 1 / 2_5_5     ,						A_ :	bool = True     ,						A_ :	bool = True     ,						A_ :	Optional[Union[float, List[float]]] = None     ,						A_ :	Optional[Union[float, List[float]]] = None     ,						**A_ :	List[str]     ,						):






      '''simple docstring'''




      super().__init__(**A_       )
      _lowerCAmelCase   :					str        =    size if size is not None else {"height": 2_5_6, "width": 2_5_6}
      _lowerCAmelCase   :					Any        =    get_size_dict(A_       )
      _lowerCAmelCase   :					Tuple        =    crop_size if crop_size is not None else {"height": 2_2_4, "width": 2_2_4}
      _lowerCAmelCase   :					Any        =    get_size_dict(A_     ,						param_name="crop_size"       )

      _lowerCAmelCase   :					str        =    do_resize
      _lowerCAmelCase   :					str        =    size
      _lowerCAmelCase   :					Dict        =    resample
      _lowerCAmelCase   :					Union[str, Any]        =    do_center_crop
      _lowerCAmelCase   :					Union[str, Any]        =    crop_size
      _lowerCAmelCase   :					List[Any]        =    do_rescale
      _lowerCAmelCase   :					Tuple        =    rescale_factor
      _lowerCAmelCase   :					Any        =    do_normalize
      _lowerCAmelCase   :					List[Any]        =    image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
      _lowerCAmelCase   :					List[str]        =    image_std if image_std is not None else IMAGENET_STANDARD_STD

 def    __magic_name__       (		self :	Optional[int]     ,						A_ :	np.ndarray     ,						A_ :	Dict[str, int]     ,						A_ :	PILImageResampling = PIL.Image.BICUBIC     ,						A_ :	Optional[Union[str, ChannelDimension]] = None     ,						**A_ :	str     ,						):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[Any]        =    get_size_dict(A_       )
      if "height" not in size or "width" not in size:
           raise ValueError(F'''The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}'''       )
      return resize(
          A_     ,						size=(size["height"], size["width"])     ,						resample=A_     ,						data_format=A_     ,						**A_       )

 def    __magic_name__       (		self :	Any     ,						A_ :	np.ndarray     ,						A_ :	Dict[str, int]     ,						A_ :	Optional[Union[str, ChannelDimension]] = None     ,						**A_ :	str     ,						):






      '''simple docstring'''




      _lowerCAmelCase   :					Any        =    get_size_dict(A_       )
      if "height" not in size or "width" not in size:
           raise ValueError(F'''The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}'''       )
      return center_crop(A_     ,						size=(size["height"], size["width"])     ,						data_format=A_     ,						**A_       )

 def    __magic_name__       (		self :	List[Any]     ,						A_ :	np.ndarray     ,						A_ :	Union[int, float]     ,						A_ :	Optional[Union[str, ChannelDimension]] = None     ,						**A_ :	List[Any]     ,						):






      '''simple docstring'''




      return rescale(A_     ,						scale=A_     ,						data_format=A_     ,						**A_       )

 def    __magic_name__       (		self :	Optional[Any]     ,						A_ :	np.ndarray     ,						A_ :	Union[float, List[float]]     ,						A_ :	Union[float, List[float]]     ,						A_ :	Optional[Union[str, ChannelDimension]] = None     ,						**A_ :	Optional[int]     ,						):






      '''simple docstring'''




      return normalize(A_     ,						mean=A_     ,						std=A_     ,						data_format=A_     ,						**A_       )


 def    __magic_name__       (		self :	Union[str, Any]     ,						A_ :	ImageInput     ,						A_ :	bool = None     ,						A_ :	Dict[str, int] = None     ,						A_ :	Union[str, Any]=None     ,						A_ :	bool = None     ,						A_ :	Dict[str, int] = None     ,						A_ :	bool = None     ,						A_ :	float = None     ,						A_ :	bool = None     ,						A_ :	Optional[Union[float, List[float]]] = None     ,						A_ :	Optional[Union[float, List[float]]] = None     ,						A_ :	Optional[Union[str, TensorType]] = None     ,						A_ :	ChannelDimension = ChannelDimension.FIRST     ,						**A_ :	Dict     ,						):






      '''simple docstring'''




      _lowerCAmelCase   :					List[str]        =    do_resize if do_resize is not None else self.do_resize
      _lowerCAmelCase   :					str        =    resample if resample is not None else self.resample
      _lowerCAmelCase   :					int        =    do_center_crop if do_center_crop is not None else self.do_center_crop
      _lowerCAmelCase   :					str        =    do_rescale if do_rescale is not None else self.do_rescale
      _lowerCAmelCase   :					Optional[Any]        =    rescale_factor if rescale_factor is not None else self.rescale_factor
      _lowerCAmelCase   :					int        =    do_normalize if do_normalize is not None else self.do_normalize
      _lowerCAmelCase   :					str        =    image_mean if image_mean is not None else self.image_mean
      _lowerCAmelCase   :					List[Any]        =    image_std if image_std is not None else self.image_std

      _lowerCAmelCase   :					Optional[int]        =    size if size is not None else self.size
      _lowerCAmelCase   :					List[str]        =    get_size_dict(A_       )
      _lowerCAmelCase   :					int        =    crop_size if crop_size is not None else self.crop_size
      _lowerCAmelCase   :					List[str]        =    get_size_dict(A_     ,						param_name="crop_size"       )

      _lowerCAmelCase   :					Tuple        =    make_list_of_images(A_       )

      if not valid_images(A_       ):
           raise ValueError(
               "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
               "torch.Tensor, tf.Tensor or jax.ndarray."       )

      if do_resize and size is None or resample is None:
           raise ValueError("Size and resample must be specified if do_resize is True."       )

      if do_center_crop and crop_size is None:
           raise ValueError("Crop size must be specified if do_center_crop is True."       )

      if do_rescale and rescale_factor is None:
           raise ValueError("Rescale factor must be specified if do_rescale is True."       )

      if do_normalize and (image_mean is None or image_std is None):
           raise ValueError("Image mean and std must be specified if do_normalize is True."       )

      # All transformations expect numpy arrays.
      _lowerCAmelCase   :					Union[str, Any]        =    [to_numpy_array(A_       ) for image in images]

      if do_resize:
           _lowerCAmelCase   :					List[str]        =    [self.resize(image=A_     ,						size=A_     ,						resample=A_       ) for image in images]

      if do_center_crop:
           _lowerCAmelCase   :					Any        =    [self.center_crop(image=A_     ,						size=A_       ) for image in images]

      if do_rescale:
           _lowerCAmelCase   :					str        =    [self.rescale(image=A_     ,						scale=A_       ) for image in images]

      if do_normalize:
           _lowerCAmelCase   :					Dict        =    [self.normalize(image=A_     ,						mean=A_     ,						std=A_       ) for image in images]

      _lowerCAmelCase   :					Tuple        =    [to_channel_dimension_format(A_     ,						A_       ) for image in images]

      _lowerCAmelCase   :					int        =    {"pixel_values": images}
      return BatchFeature(data=A_     ,						tensor_type=A_       )






