from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


__a    :	Optional[int]         =    [
    """openmmlab/upernet-convnext-tiny""",
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
__a    :	int         =    """UperNetConfig"""




class   _UpperCamelCase (		nn.Module					):


			"""simple docstring"""






			def __init__(  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ = 0 ,  lowerCAmelCase__ = False ,  lowerCAmelCase__ = 1 ,  )	->   None:

							'''simple docstring'''


							super().__init__()
							__lowercase					     =				nn.Convad(
							    in_channels=lowerCAmelCase__ ,  out_channels=lowerCAmelCase__ ,  kernel_size=lowerCAmelCase__ ,  padding=lowerCAmelCase__ ,  bias=lowerCAmelCase__ ,  dilation=lowerCAmelCase__ ,  )
							__lowercase					     =				nn.BatchNormad(lowerCAmelCase__		)
							__lowercase					     =				nn.ReLU()





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   torch.Tensor:

							'''simple docstring'''


							__lowercase					     =				self.conv(lowerCAmelCase__		)
							__lowercase					     =				self.batch_norm(lowerCAmelCase__		)
							__lowercase					     =				self.activation(lowerCAmelCase__		)

							return output




class   _UpperCamelCase (		nn.Module					):


			"""simple docstring"""






			def __init__(  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   None:

							'''simple docstring'''


							super().__init__()
							__lowercase					     =				[
							    nn.AdaptiveAvgPoolad(lowerCAmelCase__		),
							    UperNetConvModule(lowerCAmelCase__ ,  lowerCAmelCase__ ,  kernel_size=1		),
							]
							for i, layer in enumerate(self.layers		):
											self.add_module(str(lowerCAmelCase__		) ,  lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   torch.Tensor:

							'''simple docstring'''


							__lowercase					     =				input
							for layer in self.layers:
											__lowercase					     =				layer(lowerCAmelCase__		)
							return hidden_state




class   _UpperCamelCase (		nn.Module					):


			"""simple docstring"""






			def __init__(  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   None:

							'''simple docstring'''


							super().__init__()
							__lowercase					     =				pool_scales
							__lowercase					     =				align_corners
							__lowercase					     =				in_channels
							__lowercase					     =				channels
							__lowercase					     =				[]
							for i, pool_scale in enumerate(lowerCAmelCase__		):
											__lowercase					     =				UperNetPyramidPoolingBlock(pool_scale=lowerCAmelCase__ ,  in_channels=lowerCAmelCase__ ,  channels=lowerCAmelCase__		)
											self.blocks.append(lowerCAmelCase__		)
											self.add_module(str(lowerCAmelCase__		) ,  lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   List[torch.Tensor]:

							'''simple docstring'''


							__lowercase					     =				[]
							for ppm in self.blocks:
											__lowercase					     =				ppm(lowerCAmelCase__		)
											__lowercase					     =				nn.functional.interpolate(
											    lowerCAmelCase__ ,  size=x.size()[2:] ,  mode='''bilinear''' ,  align_corners=self.align_corners		)
											ppm_outs.append(lowerCAmelCase__		)
							return ppm_outs




class   _UpperCamelCase (		nn.Module					):


			"""simple docstring"""






			def __init__(  self ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   Any:

							'''simple docstring'''


							super().__init__()

							__lowercase					     =				config
							__lowercase					     =				config.pool_scales  # e.g. (1, 2, 3, 6)
							__lowercase					     =				in_channels
							__lowercase					     =				config.hidden_size
							__lowercase					     =				False
							__lowercase					     =				nn.Convad(self.channels ,  config.num_labels ,  kernel_size=1		)

							# PSP Module
							__lowercase					     =				UperNetPyramidPoolingModule(
							    self.pool_scales ,  self.in_channels[-1] ,  self.channels ,  align_corners=self.align_corners ,  )
							__lowercase					     =				UperNetConvModule(
							    self.in_channels[-1] + len(self.pool_scales		) * self.channels ,  self.channels ,  kernel_size=3 ,  padding=1 ,  )
							# FPN Module
							__lowercase					     =				nn.ModuleList()
							__lowercase					     =				nn.ModuleList()
							for in_channels in self.in_channels[:-1]:  # skip the top layer
											__lowercase					     =				UperNetConvModule(lowerCAmelCase__ ,  self.channels ,  kernel_size=1		)
											__lowercase					     =				UperNetConvModule(self.channels ,  self.channels ,  kernel_size=3 ,  padding=1		)
											self.lateral_convs.append(lowerCAmelCase__		)
											self.fpn_convs.append(lowerCAmelCase__		)

							__lowercase					     =				UperNetConvModule(
							    len(self.in_channels		) * self.channels ,  self.channels ,  kernel_size=3 ,  padding=1 ,  )





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Optional[Any]:

							'''simple docstring'''


							self.apply(self._init_weights		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   Optional[int]:

							'''simple docstring'''


							if isinstance(lowerCAmelCase__ ,  nn.Convad		):
											module.weight.data.normal_(mean=0.0 ,  std=self.config.initializer_range		)
											if module.bias is not None:
															module.bias.data.zero_()





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   Any:

							'''simple docstring'''


							__lowercase					     =				inputs[-1]
							__lowercase					     =				[x]
							psp_outs.extend(self.psp_modules(lowerCAmelCase__		)		)
							__lowercase					     =				torch.cat(lowerCAmelCase__ ,  dim=1		)
							__lowercase					     =				self.bottleneck(lowerCAmelCase__		)

							return output





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   torch.Tensor:

							'''simple docstring'''


							__lowercase					     =				[lateral_conv(encoder_hidden_states[i]		) for i, lateral_conv in enumerate(self.lateral_convs		)]

							laterals.append(self.psp_forward(lowerCAmelCase__		)		)

							# build top-down path
							__lowercase					     =				len(lowerCAmelCase__		)
							for i in range(used_backbone_levels - 1 ,  0 ,  -1		):
											__lowercase					     =				laterals[i - 1].shape[2:]
											__lowercase					     =				laterals[i - 1] + nn.functional.interpolate(
											    laterals[i] ,  size=lowerCAmelCase__ ,  mode='''bilinear''' ,  align_corners=self.align_corners		)

							# build outputs
							__lowercase					     =				[self.fpn_convs[i](laterals[i]		) for i in range(used_backbone_levels - 1		)]
							# append psp feature
							fpn_outs.append(laterals[-1]		)

							for i in range(used_backbone_levels - 1 ,  0 ,  -1		):
											__lowercase					     =				nn.functional.interpolate(
											    fpn_outs[i] ,  size=fpn_outs[0].shape[2:] ,  mode='''bilinear''' ,  align_corners=self.align_corners		)
							__lowercase					     =				torch.cat(lowerCAmelCase__ ,  dim=1		)
							__lowercase					     =				self.fpn_bottleneck(lowerCAmelCase__		)
							__lowercase					     =				self.classifier(lowerCAmelCase__		)

							return output




class   _UpperCamelCase (		nn.Module					):


			"""simple docstring"""






			def __init__(  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ = 2 ,  lowerCAmelCase__ = 3 ,  lowerCAmelCase__ = 1		)	->   None:

							'''simple docstring'''


							super().__init__()

							__lowercase					     =				config
							__lowercase					     =				config.auxiliary_in_channels
							__lowercase					     =				config.auxiliary_channels
							__lowercase					     =				config.auxiliary_num_convs
							__lowercase					     =				config.auxiliary_concat_input
							__lowercase					     =				in_index

							__lowercase					     =				(kernel_size // 2) * dilation
							__lowercase					     =				[]
							convs.append(
							    UperNetConvModule(
							        self.in_channels ,  self.channels ,  kernel_size=lowerCAmelCase__ ,  padding=lowerCAmelCase__ ,  dilation=lowerCAmelCase__		)		)
							for i in range(self.num_convs - 1		):
											convs.append(
											    UperNetConvModule(
											        self.channels ,  self.channels ,  kernel_size=lowerCAmelCase__ ,  padding=lowerCAmelCase__ ,  dilation=lowerCAmelCase__		)		)
							if self.num_convs == 0:
											__lowercase					     =				nn.Identity()
							else:
											__lowercase					     =				nn.Sequential(*lowerCAmelCase__		)
							if self.concat_input:
											__lowercase					     =				UperNetConvModule(
											    self.in_channels + self.channels ,  self.channels ,  kernel_size=lowerCAmelCase__ ,  padding=kernel_size // 2		)

							__lowercase					     =				nn.Convad(self.channels ,  config.num_labels ,  kernel_size=1		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   str:

							'''simple docstring'''


							self.apply(self._init_weights		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   Tuple:

							'''simple docstring'''


							if isinstance(lowerCAmelCase__ ,  nn.Convad		):
											module.weight.data.normal_(mean=0.0 ,  std=self.config.initializer_range		)
											if module.bias is not None:
															module.bias.data.zero_()





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   torch.Tensor:

							'''simple docstring'''


							__lowercase					     =				encoder_hidden_states[self.in_index]
							__lowercase					     =				self.convs(lowerCAmelCase__		)
							if self.concat_input:
											__lowercase					     =				self.conv_cat(torch.cat([hidden_states, output] ,  dim=1		)		)
							__lowercase					     =				self.classifier(lowerCAmelCase__		)
							return output




class   _UpperCamelCase (		_UpperCAmelCase					):


			"""simple docstring"""






			__a    :					List[Any]     					=		UperNetConfig
			__a    :					Optional[int]     					=		'''pixel_values'''
			__a    :					int     					=		True





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__		)	->   List[Any]:

							'''simple docstring'''


							if isinstance(lowerCAmelCase__ ,  lowerCAmelCase__		):
											module.backbone.init_weights()
											module.decode_head.init_weights()
											module.auxiliary_head.init_weights()





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Any:

							'''simple docstring'''


							self.backbone.init_weights()
							self.decode_head.init_weights()
							self.auxiliary_head.init_weights()





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__=False		)	->   Optional[int]:

							'''simple docstring'''


							if isinstance(lowerCAmelCase__ ,  lowerCAmelCase__		):
											__lowercase					     =				value


__a    :	Any         =    r"""
    Parameters:
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.
        config ([`UperNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

__a    :	List[str]         =    r"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See
            `attentions` under returned tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under
            returned tensors for more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""




@add_start_docstrings(
    '''UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes.'''		,_UpperCAmelCase		,)
class   _UpperCamelCase (		_UpperCAmelCase					):


			"""simple docstring"""






			def __init__(  self ,  lowerCAmelCase__		)	->   Optional[Any]:

							'''simple docstring'''


							super().__init__(lowerCAmelCase__		)

							__lowercase					     =				AutoBackbone.from_config(config.backbone_config		)

							# Semantic segmentation head(s)
							__lowercase					     =				UperNetHead(lowerCAmelCase__ ,  in_channels=self.backbone.channels		)
							__lowercase					     =				UperNetFCNHead(lowerCAmelCase__		) if config.use_auxiliary_head else None

							# Initialize weights and apply final processing
							self.post_init()





			@add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format('''batch_size, sequence_length'''		)		)
			@replace_return_docstrings(output_type=lowerCAmelCase__ ,  config_class=_CONFIG_FOR_DOC		)
			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ = None ,  lowerCAmelCase__ = None ,  lowerCAmelCase__ = None ,  lowerCAmelCase__ = None ,  lowerCAmelCase__ = None ,  )	->   Union[tuple, SemanticSegmenterOutput]:

							'''simple docstring'''


							__lowercase					     =				return_dict if return_dict is not None else self.config.use_return_dict
							__lowercase					     =				(
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							__lowercase					     =				output_attentions if output_attentions is not None else self.config.output_attentions

							__lowercase					     =				self.backbone.forward_with_filtered_kwargs(
							    lowerCAmelCase__ ,  output_hidden_states=lowerCAmelCase__ ,  output_attentions=lowerCAmelCase__		)
							__lowercase					     =				outputs.feature_maps

							__lowercase					     =				self.decode_head(lowerCAmelCase__		)
							__lowercase					     =				nn.functional.interpolate(lowerCAmelCase__ ,  size=pixel_values.shape[2:] ,  mode='''bilinear''' ,  align_corners=lowerCAmelCase__		)

							__lowercase					     =				None
							if self.auxiliary_head is not None:
											__lowercase					     =				self.auxiliary_head(lowerCAmelCase__		)
											__lowercase					     =				nn.functional.interpolate(
											    lowerCAmelCase__ ,  size=pixel_values.shape[2:] ,  mode='''bilinear''' ,  align_corners=lowerCAmelCase__		)

							__lowercase					     =				None
							if labels is not None:
											if self.config.num_labels == 1:
															raise ValueError('''The number of labels should be greater than one'''		)
											else:
															# compute weighted loss
															__lowercase					     =				CrossEntropyLoss(ignore_index=self.config.loss_ignore_index		)
															__lowercase					     =				loss_fct(lowerCAmelCase__ ,  lowerCAmelCase__		)
															__lowercase					     =				loss_fct(lowerCAmelCase__ ,  lowerCAmelCase__		)
															__lowercase					     =				main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

							if not return_dict:
											if output_hidden_states:
															__lowercase					     =				(logits,) + outputs[1:]
											else:
															__lowercase					     =				(logits,) + outputs[2:]
											return ((loss,) + output) if loss is not None else output

							return SemanticSegmenterOutput(
							    loss=lowerCAmelCase__ ,  logits=lowerCAmelCase__ ,  hidden_states=outputs.hidden_states ,  attentions=outputs.attentions ,  )
def 							UpperCAmelCase		(       lowercase				,     lowercase				):







				"""simple docstring"""



				__lowercase					     =				word.split()

				def justify(lowercase				,     lowercase				,     lowercase				) -> str:
								__lowercase					     =				max_width - width
								__lowercase					     =				len(lowercase				)
								if len(lowercase				) == 1:
												# if there is only word in line
												# just insert overall_spaces_count for the remainder of line
												return line[0] + " " * overall_spaces_count
								else:
												__lowercase					     =				words_count - 1
												# num_spaces_between_words_list[i] : tells you to insert
												# num_spaces_between_words_list[i] spaces
												# after word on line[i]
												__lowercase					     =				spaces_to_insert_between_words * [
												    overall_spaces_count // spaces_to_insert_between_words
												]
												__lowercase					     =				(
												    overall_spaces_count % spaces_to_insert_between_words
												)
												# distribute spaces via round robin to the left words
												for i in range(lowercase				):
																num_spaces_between_words_list[i] += 1
												__lowercase					     =				[]
												for i in range(lowercase				):
																# add the word
																aligned_words_list.append(line[i]				)
																# add the spaces to insert
																aligned_words_list.append(num_spaces_between_words_list[i] * ''' '''				)
												# just add the last word to the sentence
												aligned_words_list.append(line[-1]				)
												# join the aligned words list to form a justified line
												return "".join(lowercase				)

				__lowercase					     =				[]
				__lowercase					     =				[]
				__lowercase					     =				0
				for word in words:
								if width + len(lowercase				) + len(lowercase				) <= max_width:
												# keep adding words until we can fill out max_width
												# width = sum of length of all words (without overall_spaces_count)
												# len(word) = length of current word
												# len(line) = number of overall_spaces_count to insert between words
												line.append(lowercase				)
												width += len(lowercase				)
								else:
												# justify the line and add it to result
												answer.append(justify(lowercase				,     lowercase				,     lowercase				)				)
												# reset new line and new width
												__lowercase     ,	__lowercase					     =				[word], len(lowercase				)
				__lowercase					     =				max_width - width - len(lowercase				)
				answer.append(''' '''.join(lowercase				) + (remaining_spaces + 1) * ''' '''				)
				return answer


if __name__ == "__main__":
		from doctest import testmod

		testmod()