
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,
    flip_channel_order,
    get_resize_output_image_size,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
				import PIL

if is_torch_available():
				import torch


__snake_case								=  logging.get_logger(__name__)




class 						lowercase__  (						_UpperCAmelCase  ):
	A__    :   Optional[Any]				       =["""pixel_values"""]
	def __init__(			self							:   Optional[int]       ,   UpperCAmelCase_							:   bool = True       ,   UpperCAmelCase_							:   Dict[str, int] = None       ,   UpperCAmelCase_							:   PILImageResampling = PILImageResampling.BILINEAR       ,   UpperCAmelCase_							:   bool = True       ,   UpperCAmelCase_							:   Union[int, float] = 1 / 255       ,   UpperCAmelCase_							:   bool = True       ,   UpperCAmelCase_							:   Dict[str, int] = None       ,   UpperCAmelCase_							:   bool = True       ,   **UpperCAmelCase_							:   Union[str, Any]       ,   ):
						super().__init__(**UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= size if size is not None else {'shortest_edge': 224}
						SCREAMING_SNAKE_CASE__													= get_size_dict(UpperCAmelCase_       ,   default_to_square=UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= crop_size if crop_size is not None else {'height': 256, 'width': 256}
						SCREAMING_SNAKE_CASE__													= get_size_dict(UpperCAmelCase_       ,   param_name='crop_size'							)

						SCREAMING_SNAKE_CASE__													= do_resize
						SCREAMING_SNAKE_CASE__													= size
						SCREAMING_SNAKE_CASE__													= resample
						SCREAMING_SNAKE_CASE__													= do_rescale
						SCREAMING_SNAKE_CASE__													= rescale_factor
						SCREAMING_SNAKE_CASE__													= do_center_crop
						SCREAMING_SNAKE_CASE__													= crop_size
						SCREAMING_SNAKE_CASE__													= do_flip_channel_order
	def      A_				(			self							:   Optional[int]       ,   UpperCAmelCase_							:   np.ndarray       ,   UpperCAmelCase_							:   Dict[str, int]       ,   UpperCAmelCase_							:   PILImageResampling = PIL.Image.BILINEAR       ,   UpperCAmelCase_							:   Optional[Union[str, ChannelDimension]] = None       ,   **UpperCAmelCase_							:   Union[str, Any]       ,   ):
						SCREAMING_SNAKE_CASE__													= get_size_dict(UpperCAmelCase_       ,   default_to_square=UpperCAmelCase_							)
						if "shortest_edge" not in size:
											raise ValueError(F'The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}'							)
						SCREAMING_SNAKE_CASE__													= get_resize_output_image_size(UpperCAmelCase_       ,   size=size['shortest_edge']       ,   default_to_square=UpperCAmelCase_							)
						return resize(UpperCAmelCase_       ,   size=UpperCAmelCase_       ,   resample=UpperCAmelCase_       ,   data_format=UpperCAmelCase_       ,   **UpperCAmelCase_							)
	def      A_				(			self							:   Tuple       ,   UpperCAmelCase_							:   np.ndarray       ,   UpperCAmelCase_							:   Dict[str, int]       ,   UpperCAmelCase_							:   Optional[Union[str, ChannelDimension]] = None       ,   **UpperCAmelCase_							:   List[str]       ,   ):
						SCREAMING_SNAKE_CASE__													= get_size_dict(UpperCAmelCase_							)
						if "height" not in size or "width" not in size:
											raise ValueError(F'The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}'							)
						return center_crop(UpperCAmelCase_       ,   size=(size['height'], size['width'])       ,   data_format=UpperCAmelCase_       ,   **UpperCAmelCase_							)
	def      A_				(			self							:   Dict       ,   UpperCAmelCase_							:   np.ndarray       ,   UpperCAmelCase_							:   Union[int, float]       ,   UpperCAmelCase_							:   Optional[Union[str, ChannelDimension]] = None       ,   **UpperCAmelCase_							:   List[str]       ,   ):
						return rescale(UpperCAmelCase_       ,   scale=UpperCAmelCase_       ,   data_format=UpperCAmelCase_       ,   **UpperCAmelCase_							)
	def      A_				(			self							:   Union[str, Any]       ,   UpperCAmelCase_							:   np.ndarray       ,   UpperCAmelCase_							:   Optional[Union[str, ChannelDimension]] = None							):
						return flip_channel_order(UpperCAmelCase_       ,   data_format=UpperCAmelCase_							)
	def      A_				(			self							:   Optional[Any]       ,   UpperCAmelCase_							:   ImageInput       ,   UpperCAmelCase_							:   bool = None       ,   UpperCAmelCase_							:   Dict[str, int] = None       ,   UpperCAmelCase_							:   PILImageResampling = None       ,   UpperCAmelCase_							:   bool = None       ,   UpperCAmelCase_							:   float = None       ,   UpperCAmelCase_							:   bool = None       ,   UpperCAmelCase_							:   Dict[str, int] = None       ,   UpperCAmelCase_							:   bool = None       ,   UpperCAmelCase_							:   Optional[Union[str, TensorType]] = None       ,   UpperCAmelCase_							:   ChannelDimension = ChannelDimension.FIRST       ,   **UpperCAmelCase_							:   Union[str, Any]       ,   ):
						SCREAMING_SNAKE_CASE__													= do_resize if do_resize is not None else self.do_resize
						SCREAMING_SNAKE_CASE__													= resample if resample is not None else self.resample
						SCREAMING_SNAKE_CASE__													= do_rescale if do_rescale is not None else self.do_rescale
						SCREAMING_SNAKE_CASE__													= rescale_factor if rescale_factor is not None else self.rescale_factor
						SCREAMING_SNAKE_CASE__													= do_center_crop if do_center_crop is not None else self.do_center_crop
						SCREAMING_SNAKE_CASE__													= (
						    do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
						)

						SCREAMING_SNAKE_CASE__													= size if size is not None else self.size
						SCREAMING_SNAKE_CASE__													= get_size_dict(UpperCAmelCase_       ,   default_to_square=UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= crop_size if crop_size is not None else self.crop_size
						SCREAMING_SNAKE_CASE__													= get_size_dict(UpperCAmelCase_       ,   param_name='crop_size'							)

						SCREAMING_SNAKE_CASE__													= make_list_of_images(UpperCAmelCase_							)

						if not valid_images(UpperCAmelCase_							):
											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_rescale and rescale_factor is None:
											raise ValueError('Rescale factor must be specified if do_rescale is True.'							)

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

						# All transformations expect numpy arrays.
						SCREAMING_SNAKE_CASE__													= [to_numpy_array(UpperCAmelCase_							) for image in images]

						if do_resize:
											SCREAMING_SNAKE_CASE__													= [self.resize(image=UpperCAmelCase_       ,   size=UpperCAmelCase_       ,   resample=UpperCAmelCase_							) for image in images]

						if do_center_crop:
											SCREAMING_SNAKE_CASE__													= [self.center_crop(image=UpperCAmelCase_       ,   size=UpperCAmelCase_							) for image in images]

						if do_rescale:
											SCREAMING_SNAKE_CASE__													= [self.rescale(image=UpperCAmelCase_       ,   scale=UpperCAmelCase_							) for image in images]

						# the pretrained checkpoints assume images are BGR, not RGB
						if do_flip_channel_order:
											SCREAMING_SNAKE_CASE__													= [self.flip_channel_order(image=UpperCAmelCase_							) for image in images]

						SCREAMING_SNAKE_CASE__													= [to_channel_dimension_format(UpperCAmelCase_       ,   UpperCAmelCase_							) for image in images]

						SCREAMING_SNAKE_CASE__													= {'pixel_values': images}
						return BatchFeature(data=UpperCAmelCase_       ,   tensor_type=UpperCAmelCase_							)





	def      A_				(			self							:   int       ,   UpperCAmelCase_							:   List[str]       ,   UpperCAmelCase_							:   List[Tuple] = None							):
						SCREAMING_SNAKE_CASE__													= outputs.logits

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

											if is_torch_tensor(UpperCAmelCase_							):
																SCREAMING_SNAKE_CASE__													= target_sizes.numpy()

											SCREAMING_SNAKE_CASE__													= []

											for idx in range(len(UpperCAmelCase_							)							):
																SCREAMING_SNAKE_CASE__													= torch.nn.functional.interpolate(
																    logits[idx].unsqueeze(dim=0							)       ,   size=target_sizes[idx]       ,   mode='bilinear'       ,   align_corners=UpperCAmelCase_							)
																SCREAMING_SNAKE_CASE__													= resized_logits[0].argmax(dim=0							)
																semantic_segmentation.append(UpperCAmelCase_							)
						else:
											SCREAMING_SNAKE_CASE__													= logits.argmax(dim=1							)
											SCREAMING_SNAKE_CASE__													= [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]							)]

						return semantic_segmentation







from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


__snake_case								=  logging.get_logger(__name__)


__snake_case								=  TypeVar("""DatasetType""", Dataset, IterableDataset)






def 							_lowercase							(     UpperCamelCase_	,	UpperCamelCase_ = None	,	UpperCamelCase_ = None	,	UpperCamelCase_ = None	,	UpperCamelCase_ = None	,	UpperCamelCase_ = "first_exhausted"	,	)		->		DatasetType:
					'''simple docstring'''


					from .arrow_dataset import Dataset
					from .iterable_dataset import IterableDataset

					if not datasets:
										raise ValueError('Unable to interleave an empty list of datasets.'		)
					for i, dataset in enumerate(UpperCamelCase_		):
										if not isinstance(UpperCamelCase_	,	(Dataset, IterableDataset)		):
															if isinstance(UpperCamelCase_	,	(DatasetDict, IterableDatasetDict)		):
																				if not dataset:
																									raise ValueError(
																									    F'Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} '
																									    'is an empty dataset dictionary.'		)
																				raise ValueError(
																				    F'Dataset at position {i} has at least one split: {list(UpperCamelCase_		)}\n'
																				    F'Please pick one to interleave with the other datasets, for example: dataset[\'{next(iter(UpperCamelCase_		)		)}\']'		)
															raise ValueError(
															    F'Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(UpperCamelCase_		).__name__}.'		)
										if i == 0:
															SCREAMING_SNAKE_CASE__						, SCREAMING_SNAKE_CASE__													= (
															    (Dataset, IterableDataset) if isinstance(UpperCamelCase_	,	UpperCamelCase_		) else (IterableDataset, Dataset)
															)
										elif not isinstance(UpperCamelCase_	,	UpperCamelCase_		):
															raise ValueError(
															    F'Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects.'		)
					if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
										raise ValueError(F'{stopping_strategy} is not supported. Please enter a valid stopping_strategy.'		)
					if dataset_type is Dataset:
										return _interleave_map_style_datasets(
										    UpperCamelCase_	,	UpperCamelCase_	,	UpperCamelCase_	,	info=UpperCamelCase_	,	split=UpperCamelCase_	,	stopping_strategy=UpperCamelCase_		)
					else:
										return _interleave_iterable_datasets(
										    UpperCamelCase_	,	UpperCamelCase_	,	UpperCamelCase_	,	info=UpperCamelCase_	,	split=UpperCamelCase_	,	stopping_strategy=UpperCamelCase_		)






def 							_lowercase							(     UpperCamelCase_	,	UpperCamelCase_ = None	,	UpperCamelCase_ = None	,	UpperCamelCase_ = 0	,	)		->		DatasetType:
					'''simple docstring'''


					if not dsets:
										raise ValueError('Unable to concatenate an empty list of datasets.'		)
					for i, dataset in enumerate(UpperCamelCase_		):
										if not isinstance(UpperCamelCase_	,	(Dataset, IterableDataset)		):
															if isinstance(UpperCamelCase_	,	(DatasetDict, IterableDatasetDict)		):
																				if not dataset:
																									raise ValueError(
																									    F'Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} '
																									    'is an empty dataset dictionary.'		)
																				raise ValueError(
																				    F'Dataset at position {i} has at least one split: {list(UpperCamelCase_		)}\n'
																				    F'Please pick one to interleave with the other datasets, for example: dataset[\'{next(iter(UpperCamelCase_		)		)}\']'		)
															raise ValueError(
															    F'Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(UpperCamelCase_		).__name__}.'		)
										if i == 0:
															SCREAMING_SNAKE_CASE__						, SCREAMING_SNAKE_CASE__													= (
															    (Dataset, IterableDataset) if isinstance(UpperCamelCase_	,	UpperCamelCase_		) else (IterableDataset, Dataset)
															)
										elif not isinstance(UpperCamelCase_	,	UpperCamelCase_		):
															raise ValueError(
															    F'Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects.'		)
					if dataset_type is Dataset:
										return _concatenate_map_style_datasets(UpperCamelCase_	,	info=UpperCamelCase_	,	split=UpperCamelCase_	,	axis=UpperCamelCase_		)
					else:
										return _concatenate_iterable_datasets(UpperCamelCase_	,	info=UpperCamelCase_	,	split=UpperCamelCase_	,	axis=UpperCamelCase_		)





