



"""simple docstring"""





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


lowerCAmelCase_     =	[
    '''openmmlab/upernet-convnext-tiny''',
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
lowerCAmelCase_     =	'''UperNetConfig'''




class  _snake_case		(    nn.Module ):







				"""simple docstring"""






				def __init__(      self							:							str      ,    _A							:							int      ,    _A							:							int      ,    _A							:							Union[int, Tuple[int, int]]      ,    _A							:							Union[int, Tuple[int, int], str] = 0      ,    _A							:							bool = False      ,    _A							:							Union[int, Tuple[int, int]] = 1      ,    ):



											"""simple docstring"""







											super().__init__()
											_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    nn.Convad(
											    in_channels=a_      ,    out_channels=a_      ,    kernel_size=a_      ,    padding=a_      ,    bias=a_      ,    dilation=a_      ,    )
											_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    nn.BatchNormad(a_)
											_SCREAMING_SNAKE_CASE		:						Tuple			  =    nn.ReLU()





				def 		_lowerCAmelCase		(      self							:							Optional[Any]      ,    _A							:							torch.Tensor):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    self.conv(a_)
											_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    self.batch_norm(a_)
											_SCREAMING_SNAKE_CASE		:						Dict			  =    self.activation(a_)

											return output




class  _snake_case		(    nn.Module ):







				"""simple docstring"""






				def __init__(      self							:							List[str]      ,    _A							:							int      ,    _A							:							int      ,    _A							:							int):



											"""simple docstring"""







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





				def 		_lowerCAmelCase		(      self							:							List[Any]      ,    _A							:							torch.Tensor):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						List[Any]			  =    input
											for layer in self.layers:
																		_SCREAMING_SNAKE_CASE		:						str			  =    layer(a_)
											return hidden_state




class  _snake_case		(    nn.Module ):







				"""simple docstring"""






				def __init__(      self							:							List[str]      ,    _A							:							Tuple[int, ...]      ,    _A							:							int      ,    _A							:							int      ,    _A							:							bool):



											"""simple docstring"""







											super().__init__()
											_SCREAMING_SNAKE_CASE		:						List[str]			  =    pool_scales
											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    align_corners
											_SCREAMING_SNAKE_CASE		:						List[Any]			  =    in_channels
											_SCREAMING_SNAKE_CASE		:						Dict			  =    channels
											_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    []
											for i, pool_scale in enumerate(a_):
																		_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    UperNetPyramidPoolingBlock(pool_scale=a_      ,    in_channels=a_      ,    channels=a_)
																		self.blocks.append(a_)
																		self.add_module(str(a_)      ,    a_)





				def 		_lowerCAmelCase		(      self							:							Union[str, Any]      ,    _A							:							torch.Tensor):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						str			  =    []
											for ppm in self.blocks:
																		_SCREAMING_SNAKE_CASE		:						Any			  =    ppm(a_)
																		_SCREAMING_SNAKE_CASE		:						Dict			  =    nn.functional.interpolate(
																		    a_      ,    size=x.size()[2:]      ,    mode="""bilinear"""      ,    align_corners=self.align_corners)
																		ppm_outs.append(a_)
											return ppm_outs




class  _snake_case		(    nn.Module ):







				"""simple docstring"""






				def __init__(      self							:							Any      ,    _A							:							Any      ,    _A							:							Dict):



											"""simple docstring"""







											super().__init__()

											_SCREAMING_SNAKE_CASE		:						Tuple			  =    config
											_SCREAMING_SNAKE_CASE		:						Tuple			  =    config.pool_scales  # e.g. (1, 2, 3, 6)
											_SCREAMING_SNAKE_CASE		:						str			  =    in_channels
											_SCREAMING_SNAKE_CASE		:						Dict			  =    config.hidden_size
											_SCREAMING_SNAKE_CASE		:						Any			  =    False
											_SCREAMING_SNAKE_CASE		:						List[Any]			  =    nn.Convad(self.channels      ,    config.num_labels      ,    kernel_size=1)

											# PSP Module
											_SCREAMING_SNAKE_CASE		:						List[Any]			  =    UperNetPyramidPoolingModule(
											    self.pool_scales      ,    self.in_channels[-1]      ,    self.channels      ,    align_corners=self.align_corners      ,    )
											_SCREAMING_SNAKE_CASE		:						int			  =    UperNetConvModule(
											    self.in_channels[-1] + len(self.pool_scales) * self.channels      ,    self.channels      ,    kernel_size=3      ,    padding=1      ,    )
											# FPN Module
											_SCREAMING_SNAKE_CASE		:						str			  =    nn.ModuleList()
											_SCREAMING_SNAKE_CASE		:						Tuple			  =    nn.ModuleList()
											for in_channels in self.in_channels[:-1]:  # skip the top layer
																		_SCREAMING_SNAKE_CASE		:						List[Any]			  =    UperNetConvModule(a_      ,    self.channels      ,    kernel_size=1)
																		_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    UperNetConvModule(self.channels      ,    self.channels      ,    kernel_size=3      ,    padding=1)
																		self.lateral_convs.append(a_)
																		self.fpn_convs.append(a_)

											_SCREAMING_SNAKE_CASE		:						str			  =    UperNetConvModule(
											    len(self.in_channels) * self.channels      ,    self.channels      ,    kernel_size=3      ,    padding=1      ,    )





				def 		_lowerCAmelCase		(      self							:							List[Any]):



											"""simple docstring"""







											self.apply(self._init_weights)





				def 		_lowerCAmelCase		(      self							:							Optional[Any]      ,    _A							:							List[Any]):



											"""simple docstring"""







											if isinstance(a_      ,    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 		_lowerCAmelCase		(      self							:							List[Any]      ,    _A							:							Any):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    inputs[-1]
											_SCREAMING_SNAKE_CASE		:						Tuple			  =    [x]
											psp_outs.extend(self.psp_modules(a_))
											_SCREAMING_SNAKE_CASE		:						Tuple			  =    torch.cat(a_      ,    dim=1)
											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    self.bottleneck(a_)

											return output





				def 		_lowerCAmelCase		(      self							:							Optional[Any]      ,    _A							:							torch.Tensor):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						int			  =    [lateral_conv(encoder_hidden_states[i]) for i, lateral_conv in enumerate(self.lateral_convs)]

											laterals.append(self.psp_forward(a_))

											# build top-down path
											_SCREAMING_SNAKE_CASE		:						List[str]			  =    len(a_)
											for i in range(used_backbone_levels - 1      ,    0      ,    -1):
																		_SCREAMING_SNAKE_CASE		:						Dict			  =    laterals[i - 1].shape[2:]
																		_SCREAMING_SNAKE_CASE		:						Tuple			  =    laterals[i - 1] + nn.functional.interpolate(
																		    laterals[i]      ,    size=a_      ,    mode="""bilinear"""      ,    align_corners=self.align_corners)

											# build outputs
											_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    [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):
																		_SCREAMING_SNAKE_CASE		:						int			  =    nn.functional.interpolate(
																		    fpn_outs[i]      ,    size=fpn_outs[0].shape[2:]      ,    mode="""bilinear"""      ,    align_corners=self.align_corners)
											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    torch.cat(a_      ,    dim=1)
											_SCREAMING_SNAKE_CASE		:						Any			  =    self.fpn_bottleneck(a_)
											_SCREAMING_SNAKE_CASE		:						Dict			  =    self.classifier(a_)

											return output




class  _snake_case		(    nn.Module ):







				"""simple docstring"""






				def __init__(      self							:							Dict      ,    _A							:							int      ,    _A							:							int = 2      ,    _A							:							int = 3      ,    _A							:							Union[int, Tuple[int, int]] = 1):



											"""simple docstring"""







											super().__init__()

											_SCREAMING_SNAKE_CASE		:						List[Any]			  =    config
											_SCREAMING_SNAKE_CASE		:						Dict			  =    config.auxiliary_in_channels
											_SCREAMING_SNAKE_CASE		:						Dict			  =    config.auxiliary_channels
											_SCREAMING_SNAKE_CASE		:						int			  =    config.auxiliary_num_convs
											_SCREAMING_SNAKE_CASE		:						Any			  =    config.auxiliary_concat_input
											_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    in_index

											_SCREAMING_SNAKE_CASE		:						int			  =    (kernel_size // 2) * dilation
											_SCREAMING_SNAKE_CASE		:						Tuple			  =    []
											convs.append(
											    UperNetConvModule(
											        self.in_channels      ,    self.channels      ,    kernel_size=a_      ,    padding=a_      ,    dilation=a_))
											for i in range(self.num_convs - 1):
																		convs.append(
																		    UperNetConvModule(
																		        self.channels      ,    self.channels      ,    kernel_size=a_      ,    padding=a_      ,    dilation=a_))
											if self.num_convs == 0:
																		_SCREAMING_SNAKE_CASE		:						str			  =    nn.Identity()
											else:
																		_SCREAMING_SNAKE_CASE		:						List[Any]			  =    nn.Sequential(*a_)
											if self.concat_input:
																		_SCREAMING_SNAKE_CASE		:						Dict			  =    UperNetConvModule(
																		    self.in_channels + self.channels      ,    self.channels      ,    kernel_size=a_      ,    padding=kernel_size // 2)

											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    nn.Convad(self.channels      ,    config.num_labels      ,    kernel_size=1)





				def 		_lowerCAmelCase		(      self							:							int):



											"""simple docstring"""







											self.apply(self._init_weights)





				def 		_lowerCAmelCase		(      self							:							Dict      ,    _A							:							List[str]):



											"""simple docstring"""







											if isinstance(a_      ,    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 		_lowerCAmelCase		(      self							:							List[Any]      ,    _A							:							torch.Tensor):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						Any			  =    encoder_hidden_states[self.in_index]
											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    self.convs(a_)
											if self.concat_input:
																		_SCREAMING_SNAKE_CASE		:						Any			  =    self.conv_cat(torch.cat([hidden_states, output]      ,    dim=1))
											_SCREAMING_SNAKE_CASE		:						str			  =    self.classifier(a_)
											return output




class  _snake_case		(    __a ):







				"""simple docstring"""






				a			=						UperNetConfig
				a			=						"pixel_values"
				a			=						True





				def 		_lowerCAmelCase		(      self							:							Dict      ,    _A							:							Optional[int]):



											"""simple docstring"""







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





				def 		_lowerCAmelCase		(      self							:							str):



											"""simple docstring"""







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





				def 		_lowerCAmelCase		(      self							:							int      ,    _A							:							Union[str, Any]      ,    _A							:							str=False):



											"""simple docstring"""







											if isinstance(a_      ,    a_):
																		_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    value


lowerCAmelCase_     =	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.
'''

lowerCAmelCase_     =	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."   ,       __a   ,       )
class  _snake_case		(    __a ):







				"""simple docstring"""






				def __init__(      self							:							Union[str, Any]      ,    _A							:							Any):



											"""simple docstring"""







											super().__init__(a_)

											_SCREAMING_SNAKE_CASE		:						Tuple			  =    AutoBackbone.from_config(config.backbone_config)

											# Semantic segmentation head(s)
											_SCREAMING_SNAKE_CASE		:						int			  =    UperNetHead(a_      ,    in_channels=self.backbone.channels)
											_SCREAMING_SNAKE_CASE		:						Tuple			  =    UperNetFCNHead(a_) 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=a_      ,    config_class=_CONFIG_FOR_DOC)
				def 		_lowerCAmelCase		(      self							:							Union[str, Any]      ,    _A							:							Optional[torch.Tensor] = None      ,    _A							:							Optional[bool] = None      ,    _A							:							Optional[bool] = None      ,    _A							:							Optional[torch.Tensor] = None      ,    _A							:							Optional[bool] = None      ,    ):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						Dict			  =    return_dict if return_dict is not None else self.config.use_return_dict
											_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											_SCREAMING_SNAKE_CASE		:						Tuple			  =    output_attentions if output_attentions is not None else self.config.output_attentions

											_SCREAMING_SNAKE_CASE		:						str			  =    self.backbone.forward_with_filtered_kwargs(
											    a_      ,    output_hidden_states=a_      ,    output_attentions=a_)
											_SCREAMING_SNAKE_CASE		:						Dict			  =    outputs.feature_maps

											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    self.decode_head(a_)
											_SCREAMING_SNAKE_CASE		:						int			  =    nn.functional.interpolate(a_      ,    size=pixel_values.shape[2:]      ,    mode="""bilinear"""      ,    align_corners=a_)

											_SCREAMING_SNAKE_CASE		:						Dict			  =    None
											if self.auxiliary_head is not None:
																		_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    self.auxiliary_head(a_)
																		_SCREAMING_SNAKE_CASE		:						List[str]			  =    nn.functional.interpolate(
																		    a_      ,    size=pixel_values.shape[2:]      ,    mode="""bilinear"""      ,    align_corners=a_)

											_SCREAMING_SNAKE_CASE		:						Any			  =    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
																									_SCREAMING_SNAKE_CASE		:						str			  =    CrossEntropyLoss(ignore_index=self.config.loss_ignore_index)
																									_SCREAMING_SNAKE_CASE		:						Any			  =    loss_fct(a_      ,    a_)
																									_SCREAMING_SNAKE_CASE		:						List[str]			  =    loss_fct(a_      ,    a_)
																									_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

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

											return SemanticSegmenterOutput(
											    loss=a_      ,    logits=a_      ,    hidden_states=outputs.hidden_states      ,    attentions=outputs.attentions      ,    )



"""simple docstring"""





def      lowerCamelCase_(__SCREAMING_SNAKE_CASE  )->    int:

							if not isinstance(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE  ):
														raise TypeError("""only integers accepted as input"""  )
							else:
														_SCREAMING_SNAKE_CASE		:						List[Any]			  =    str(abs(__SCREAMING_SNAKE_CASE  )  )
														_SCREAMING_SNAKE_CASE		:						List[str]			  =    [list(__SCREAMING_SNAKE_CASE  ) for char in range(len(__SCREAMING_SNAKE_CASE  )  )]
														for index in range(len(__SCREAMING_SNAKE_CASE  )  ):
																					num_transpositions[index].pop(__SCREAMING_SNAKE_CASE  )
														return max(
														    int("""""".join(list(__SCREAMING_SNAKE_CASE  )  )  ) for transposition in num_transpositions  )


if __name__ == "__main__":
			__import__('''doctest''').testmod()


