

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


__lowercase     : Dict       =						logging.get_logger(__name__)



class 						_A     (				snake_case				):





  '''simple docstring'''



  __lowerCamelCase			:      Optional[int]					=   ['''pixel_values''']





  def __init__(  self   ,SCREAMING_SNAKE_CASE_ = True   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = PILImageResampling.BILINEAR   ,SCREAMING_SNAKE_CASE_ = True   ,SCREAMING_SNAKE_CASE_ = 1 / 255   ,SCREAMING_SNAKE_CASE_ = True   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = True   ,**SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      super().__init__(**SCREAMING_SNAKE_CASE_    )
      snake_case							:		List[Any]        =		size if size is not None else {"""shortest_edge""": 224}
      snake_case							:		str        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,default_to_square=SCREAMING_SNAKE_CASE_    )
      snake_case							:		str        =		crop_size if crop_size is not None else {"""height""": 256, """width""": 256}
      snake_case							:		Optional[Any]        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,param_name="""crop_size"""    )

      snake_case							:		Optional[Any]        =		do_resize
      snake_case							:		Union[str, Any]        =		size
      snake_case							:		Dict        =		resample
      snake_case							:		Dict        =		do_rescale
      snake_case							:		Dict        =		rescale_factor
      snake_case							:		List[str]        =		do_center_crop
      snake_case							:		Dict        =		crop_size
      snake_case							:		Any        =		do_flip_channel_order





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = PIL.Image.BILINEAR   ,SCREAMING_SNAKE_CASE_ = None   ,**SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      snake_case							:		str        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,default_to_square=SCREAMING_SNAKE_CASE_    )
      if "shortest_edge" not in size:
          raise ValueError(F"""The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}"""    )
      snake_case							:		List[Any]        =		get_resize_output_image_size(SCREAMING_SNAKE_CASE_   ,size=size["""shortest_edge"""]   ,default_to_square=SCREAMING_SNAKE_CASE_    )
      return resize(SCREAMING_SNAKE_CASE_   ,size=SCREAMING_SNAKE_CASE_   ,resample=SCREAMING_SNAKE_CASE_   ,data_format=SCREAMING_SNAKE_CASE_   ,**SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = None   ,**SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      snake_case							:		Union[str, Any]        =		get_size_dict(SCREAMING_SNAKE_CASE_    )
      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(SCREAMING_SNAKE_CASE_   ,size=(size["""height"""], size["""width"""])   ,data_format=SCREAMING_SNAKE_CASE_   ,**SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = None   ,**SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      return rescale(SCREAMING_SNAKE_CASE_   ,scale=SCREAMING_SNAKE_CASE_   ,data_format=SCREAMING_SNAKE_CASE_   ,**SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = None    ):

      '''simple docstring'''
      return flip_channel_order(SCREAMING_SNAKE_CASE_   ,data_format=SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = ChannelDimension.FIRST   ,**SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      snake_case							:		List[Any]        =		do_resize if do_resize is not None else self.do_resize
      snake_case							:		List[str]        =		resample if resample is not None else self.resample
      snake_case							:		Union[str, Any]        =		do_rescale if do_rescale is not None else self.do_rescale
      snake_case							:		Union[str, Any]        =		rescale_factor if rescale_factor is not None else self.rescale_factor
      snake_case							:		str        =		do_center_crop if do_center_crop is not None else self.do_center_crop
      snake_case							:		Union[str, Any]        =		(
          do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
      )

      snake_case							:		Tuple        =		size if size is not None else self.size
      snake_case							:		Optional[Any]        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,default_to_square=SCREAMING_SNAKE_CASE_    )
      snake_case							:		str        =		crop_size if crop_size is not None else self.crop_size
      snake_case							:		Optional[Any]        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,param_name="""crop_size"""    )

      snake_case							:		List[Any]        =		make_list_of_images(SCREAMING_SNAKE_CASE_    )

      if not valid_images(SCREAMING_SNAKE_CASE_    ):
          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.
      snake_case							:		Dict        =		[to_numpy_array(SCREAMING_SNAKE_CASE_    ) for image in images]

      if do_resize:
          snake_case							:		Union[str, Any]        =		[self.resize(image=SCREAMING_SNAKE_CASE_   ,size=SCREAMING_SNAKE_CASE_   ,resample=SCREAMING_SNAKE_CASE_    ) for image in images]

      if do_center_crop:
          snake_case							:		Optional[Any]        =		[self.center_crop(image=SCREAMING_SNAKE_CASE_   ,size=SCREAMING_SNAKE_CASE_    ) for image in images]

      if do_rescale:
          snake_case							:		Dict        =		[self.rescale(image=SCREAMING_SNAKE_CASE_   ,scale=SCREAMING_SNAKE_CASE_    ) for image in images]

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

      snake_case							:		List[Any]        =		[to_channel_dimension_format(SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_    ) for image in images]

      snake_case							:		int        =		{"""pixel_values""": images}
      return BatchFeature(data=SCREAMING_SNAKE_CASE_   ,tensor_type=SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = None    ):

      '''simple docstring'''
      snake_case							:		Dict        =		outputs.logits

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

          if is_torch_tensor(SCREAMING_SNAKE_CASE_    ):
              snake_case							:		int        =		target_sizes.numpy()

          snake_case							:		Optional[Any]        =		[]

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

      return semantic_segmentation





import os
import sys
import tempfile

import torch

from .state import AcceleratorState
from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment







def 							A__       (    lowerCamelCase  ,     lowerCamelCase=()  ,     lowerCamelCase=None  ,     lowerCamelCase="no"  ,     lowerCamelCase="29500"    )	->						List[Any]:
				UpperCamelCase_:				Any						=    False
				UpperCamelCase_:				int						=    False
				if any(key.startswith("""KAGGLE"""    ) for key in os.environ.keys()    ):
								UpperCamelCase_:				str						=    True
				elif "IPython" in sys.modules:
								UpperCamelCase_:				Tuple						=    """google.colab""" in str(sys.modules["""IPython"""].get_ipython()    )

				try:
								UpperCamelCase_:				Union[str, Any]						=    PrecisionType(mixed_precision.lower()    )
				except ValueError:
								raise ValueError(
								    F'''Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.'''    )

				if (in_colab or in_kaggle) and (os.environ.get("""TPU_NAME"""  ,     lowerCamelCase    ) is not None):
								# TPU launch
								import torch_xla.distributed.xla_multiprocessing as xmp

								if len(AcceleratorState._shared_state    ) > 0:
												raise ValueError(
												    """To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside """
												    """your training function. Restart your notebook and make sure no cells initializes an """
												    """`Accelerator`."""    )
								if num_processes is None:
												UpperCamelCase_:				str						=    8

								UpperCamelCase_:				int						=    PrepareForLaunch(lowerCamelCase  ,     distributed_type="""TPU"""    )
								print(F'''Launching a training on {num_processes} TPU cores.'''    )
								xmp.spawn(lowerCamelCase  ,     args=lowerCamelCase  ,     nprocs=lowerCamelCase  ,     start_method="""fork"""    )
				elif in_colab:
								# No need for a distributed launch otherwise as it's either CPU or one GPU.
								if torch.cuda.is_available():
												print("""Launching training on one GPU."""    )
								else:
												print("""Launching training on one CPU."""    )
								function(*lowerCamelCase    )
				else:
								if num_processes is None:
												raise ValueError(
												    """You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call."""    )

								if num_processes > 1:
												# Multi-GPU launch
												from torch.multiprocessing import start_processes
												from torch.multiprocessing.spawn import ProcessRaisedException

												if len(AcceleratorState._shared_state    ) > 0:
																raise ValueError(
																    """To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized """
																    """inside your training function. Restart your notebook and make sure no cells initializes an """
																    """`Accelerator`."""    )

												if torch.cuda.is_initialized():
																raise ValueError(
																    """To launch a multi-GPU training from your notebook, you need to avoid running any instruction """
																    """using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA """
																    """function."""    )

												# torch.distributed will expect a few environment variable to be here. We set the ones common to each
												# process here (the other ones will be set be the launcher).
												with patch_environment(
												    world_size=lowerCamelCase  ,     master_addr="""127.0.01"""  ,     master_port=lowerCamelCase  ,     mixed_precision=lowerCamelCase    ):
																UpperCamelCase_:				Optional[int]						=    PrepareForLaunch(lowerCamelCase  ,     distributed_type="""MULTI_GPU"""    )
																print(F'''Launching training on {num_processes} GPUs.'''    )
																try:
																				start_processes(lowerCamelCase  ,     args=lowerCamelCase  ,     nprocs=lowerCamelCase  ,     start_method="""fork"""    )
																except ProcessRaisedException as e:
																				if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
																								raise RuntimeError(
																								    """CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. """
																								    """This likely stems from an outside import causing issues once the `notebook_launcher()` is called. """
																								    """Please review your imports and test them when running the `notebook_launcher()` to identify """
																								    """which one is problematic."""    ) from e

								else:
												# No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
												if is_mps_available():
																UpperCamelCase_:				List[Any]						=    """1"""
																print("""Launching training on MPS."""    )
												elif torch.cuda.is_available():
																print("""Launching training on one GPU."""    )
												else:
																print("""Launching training on CPU."""    )
												function(*lowerCamelCase    )







def 							A__       (    lowerCamelCase  ,     lowerCamelCase=()  ,     lowerCamelCase=2    )	->						List[str]:
				from torch.multiprocessing import start_processes

				with tempfile.NamedTemporaryFile() as tmp_file:
								# torch.distributed will expect a few environment variable to be here. We set the ones common to each
								# process here (the other ones will be set be the launcher).
								with patch_environment(
								    world_size=lowerCamelCase  ,     master_addr="""127.0.01"""  ,     master_port="""29500"""  ,     accelerate_mixed_precision="""no"""  ,     accelerate_debug_rdv_file=tmp_file.name  ,     accelerate_use_cpu="""yes"""  ,     ):
												UpperCamelCase_:				str						=    PrepareForLaunch(lowerCamelCase  ,     debug=lowerCamelCase    )
												start_processes(lowerCamelCase  ,     args=lowerCamelCase  ,     nprocs=lowerCamelCase  ,     start_method="""fork"""    )






