



"""simple docstring"""





from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    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_torch_available, is_torch_tensor, logging


if is_torch_available():
     import torch


__lowerCAmelCase       :					List[Any]      					=					logging.get_logger(__name__)



class    _lowerCAmelCase    (				lowerCamelCase__		):





  """simple docstring"""






  _lowerCamelCase			  =       ['''pixel_values''']





  def __init__(					self							,    _lowercase = True							,    _lowercase = None							,    _lowercase = PILImageResampling.BILINEAR							,    _lowercase = True							,    _lowercase = None							,    _lowercase = True							,    _lowercase = 1 / 2_5_5							,    _lowercase = True							,    _lowercase = None							,    _lowercase = None							,    **_lowercase							,    ) ->						None:







        '''simple docstring'''

        super().__init__(**__lowerCamelCase   )
        snake_case_   : List[Any]						=			size if size is not None else {'''shortest_edge''': 2_5_6}
        snake_case_   : Any						=			get_size_dict(__lowerCamelCase							,    default_to_square=__lowerCamelCase   )
        snake_case_   : Optional[Any]						=			crop_size if crop_size is not None else {'''height''': 2_2_4, '''width''': 2_2_4}
        snake_case_   : Any						=			get_size_dict(__lowerCamelCase							,    param_name="""crop_size"""   )
        snake_case_   : Dict						=			do_resize
        snake_case_   : str						=			size
        snake_case_   : Tuple						=			resample
        snake_case_   : List[Any]						=			do_center_crop
        snake_case_   : Dict						=			crop_size
        snake_case_   : Optional[int]						=			do_rescale
        snake_case_   : List[str]						=			rescale_factor
        snake_case_   : List[Any]						=			do_normalize
        snake_case_   : Optional[Any]						=			image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
        snake_case_   : str						=			image_std if image_std is not None else IMAGENET_STANDARD_STD





  def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase							,    _lowercase = PILImageResampling.BICUBIC							,    _lowercase = None							,    **_lowercase							,    ) ->						np.ndarray:







        '''simple docstring'''

        snake_case_   : Dict						=			get_size_dict(__lowerCamelCase							,    default_to_square=__lowerCamelCase   )
        if "shortest_edge" not in size:
              raise ValueError(f'The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'   )
        snake_case_   : str						=			get_resize_output_image_size(__lowerCamelCase							,    size=size["""shortest_edge"""]							,    default_to_square=__lowerCamelCase   )
        return resize(__lowerCamelCase							,    size=__lowerCamelCase							,    resample=__lowerCamelCase							,    data_format=__lowerCamelCase							,    **__lowerCamelCase   )





  def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase							,    _lowercase = None							,    **_lowercase							,    ) ->						np.ndarray:







        '''simple docstring'''

        snake_case_   : Union[str, Any]						=			get_size_dict(__lowerCamelCase   )
        if "height" not in size or "width" not in size:
              raise ValueError(f'The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}'   )
        return center_crop(__lowerCamelCase							,    size=(size["""height"""], size["""width"""])							,    data_format=__lowerCamelCase							,    **__lowerCamelCase   )





  def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase							,    _lowercase = None							,    **_lowercase   ) ->						np.ndarray:







        '''simple docstring'''

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





  def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase							,    _lowercase							,    _lowercase = None							,    **_lowercase							,    ) ->						np.ndarray:







        '''simple docstring'''

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





  def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase = None							,    _lowercase = None							,    _lowercase = None							,    _lowercase = None							,    _lowercase = None							,    _lowercase = None							,    _lowercase = None							,    _lowercase = None							,    _lowercase = None							,    _lowercase = None							,    _lowercase = None							,    _lowercase = ChannelDimension.FIRST							,    **_lowercase							,    ) ->						Optional[Any]:







        '''simple docstring'''

        snake_case_   : Optional[int]						=			do_resize if do_resize is not None else self.do_resize
        snake_case_   : Optional[int]						=			size if size is not None else self.size
        snake_case_   : Union[str, Any]						=			get_size_dict(__lowerCamelCase							,    default_to_square=__lowerCamelCase   )
        snake_case_   : Dict						=			resample if resample is not None else self.resample
        snake_case_   : Union[str, Any]						=			do_center_crop if do_center_crop is not None else self.do_center_crop
        snake_case_   : Any						=			crop_size if crop_size is not None else self.crop_size
        snake_case_   : List[str]						=			get_size_dict(__lowerCamelCase							,    param_name="""crop_size"""   )
        snake_case_   : List[str]						=			do_rescale if do_rescale is not None else self.do_rescale
        snake_case_   : str						=			rescale_factor if rescale_factor is not None else self.rescale_factor
        snake_case_   : int						=			do_normalize if do_normalize is not None else self.do_normalize
        snake_case_   : Optional[Any]						=			image_mean if image_mean is not None else self.image_mean
        snake_case_   : List[Any]						=			image_std if image_std is not None else self.image_std

        snake_case_   : Any						=			make_list_of_images(__lowerCamelCase   )

        if not valid_images(__lowerCamelCase   ):
              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:
              raise ValueError("""Size 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.
        snake_case_   : Union[str, Any]						=			[to_numpy_array(__lowerCamelCase   ) for image in images]

        if do_resize:
              snake_case_   : str						=			[self.resize(image=__lowerCamelCase							,    size=__lowerCamelCase							,    resample=__lowerCamelCase   ) for image in images]

        if do_center_crop:
              snake_case_   : List[Any]						=			[self.center_crop(image=__lowerCamelCase							,    size=__lowerCamelCase   ) for image in images]

        if do_rescale:
              snake_case_   : Any						=			[self.rescale(image=__lowerCamelCase							,    scale=__lowerCamelCase   ) for image in images]

        if do_normalize:
              snake_case_   : Tuple						=			[self.normalize(image=__lowerCamelCase							,    mean=__lowerCamelCase							,    std=__lowerCamelCase   ) for image in images]

        snake_case_   : int						=			[to_channel_dimension_format(__lowerCamelCase							,    __lowerCamelCase   ) for image in images]

        snake_case_   : int						=			{'''pixel_values''': images}
        return BatchFeature(data=__lowerCamelCase							,    tensor_type=__lowerCamelCase   )





  def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase = None   ) ->						Union[str, Any]:







        '''simple docstring'''

        snake_case_   : Optional[int]						=			outputs.logits

        # Resize logits and compute semantic segmentation maps
        if target_sizes is not None:
              if len(__lowerCamelCase   ) != len(__lowerCamelCase   ):
                    raise ValueError(
                        """Make sure that you pass in as many target sizes as the batch dimension of the logits"""   )

              if is_torch_tensor(__lowerCamelCase   ):
                    snake_case_   : str						=			target_sizes.numpy()

              snake_case_   : Optional[Any]						=			[]

              for idx in range(len(__lowerCamelCase   )   ):
                    snake_case_   : Optional[int]						=			torch.nn.functional.interpolate(
                        logits[idx].unsqueeze(dim=0   )							,    size=target_sizes[idx]							,    mode="""bilinear"""							,    align_corners=__lowerCamelCase   )
                    snake_case_   : Optional[Any]						=			resized_logits[0].argmax(dim=0   )
                    semantic_segmentation.append(__lowerCamelCase   )
        else:
              snake_case_   : Optional[Any]						=			logits.argmax(dim=1   )
              snake_case_   : int						=			[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]   )]

        return semantic_segmentation






"""simple docstring"""





from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__lowerCAmelCase       :					str      					=					logging.get_logger(__name__)

__lowerCAmelCase       :					Tuple      					=					{
    '''shi-labs/nat-mini-in1k-224''': '''https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json''',
    # See all Nat models at https://huggingface.co/models?filter=nat
}



class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__  ,				SCREAMING_SNAKE_CASE__		):





		"""simple docstring"""






		_lowerCamelCase			  =       '''nat'''

		_lowerCamelCase			  =       {
		    '''num_attention_heads''': '''num_heads''',
		    '''num_hidden_layers''': '''num_layers''',
		}





		def __init__(					self							,    _lowercase=4							,    _lowercase=3							,    _lowercase=6_4							,    _lowercase=[3, 4, 6, 5]							,    _lowercase=[2, 4, 8, 1_6]							,    _lowercase=7							,    _lowercase=3.0							,    _lowercase=True							,    _lowercase=0.0							,    _lowercase=0.0							,    _lowercase=0.1							,    _lowercase="gelu"							,    _lowercase=0.02							,    _lowercase=1E-5							,    _lowercase=0.0							,    _lowercase=None							,    _lowercase=None							,    **_lowercase							,    ) ->						Optional[Any]:







								'''simple docstring'''

								super().__init__(**_lowercase   )

								snake_case_   : Any						=			patch_size
								snake_case_   : Optional[Any]						=			num_channels
								snake_case_   : Optional[Any]						=			embed_dim
								snake_case_   : Tuple						=			depths
								snake_case_   : int						=			len(_lowercase   )
								snake_case_   : Optional[int]						=			num_heads
								snake_case_   : List[str]						=			kernel_size
								snake_case_   : str						=			mlp_ratio
								snake_case_   : str						=			qkv_bias
								snake_case_   : str						=			hidden_dropout_prob
								snake_case_   : Tuple						=			attention_probs_dropout_prob
								snake_case_   : Tuple						=			drop_path_rate
								snake_case_   : Dict						=			hidden_act
								snake_case_   : Union[str, Any]						=			layer_norm_eps
								snake_case_   : Tuple						=			initializer_range
								# we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
								# this indicates the channel dimension after the last stage of the model
								snake_case_   : Union[str, Any]						=			int(embed_dim * 2 ** (len(_lowercase   ) - 1)   )
								snake_case_   : Union[str, Any]						=			layer_scale_init_value
								snake_case_   : Optional[Any]						=			["""stem"""] + [f'stage{idx}' for idx in range(1							,    len(_lowercase   ) + 1   )]
								snake_case_			,		snake_case_   : Any						=			get_aligned_output_features_output_indices(
								    out_features=_lowercase							,    out_indices=_lowercase							,    stage_names=self.stage_names   )

