



'''simple docstring'''




import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation



def 	a			( lowerCamelCase__    ):
			'''simple docstring'''





			A_	:						Optional[int]	=     3_84
			if "tiny" in model_name:
						A_	:						Dict	=     [3, 3, 9, 3]
						A_	:						Tuple	=     [96, 1_92, 3_84, 7_68]
			if "small" in model_name:
						A_	:						int	=     [3, 3, 27, 3]
						A_	:						List[str]	=     [96, 1_92, 3_84, 7_68]
			if "base" in model_name:
						A_	:						Optional[int]	=     [3, 3, 27, 3]
						A_	:						Tuple	=     [1_28, 2_56, 5_12, 10_24]
						A_	:						Tuple	=     5_12
			if "large" in model_name:
						A_	:						List[str]	=     [3, 3, 27, 3]
						A_	:						Tuple	=     [1_92, 3_84, 7_68, 15_36]
						A_	:						Union[str, Any]	=     7_68
			if "xlarge" in model_name:
						A_	:						Dict	=     [3, 3, 27, 3]
						A_	:						Tuple	=     [2_56, 5_12, 10_24, 20_48]
						A_	:						str	=     10_24

			# set label information
			A_	:						List[str]	=     1_50
			A_	:						Any	=     """huggingface/label-files"""
			A_	:						Tuple	=     """ade20k-id2label.json"""
			A_	:						List[Any]	=     json.load(open(hf_hub_download(lowerCamelCase__						,      lowerCamelCase__						,      repo_type="""dataset"""    )						,      """r"""    )    )
			A_	:						Tuple	=     {int(lowerCamelCase__    ): v for k, v in idalabel.items()}
			A_	:						Optional[int]	=     {v: k for k, v in idalabel.items()}

			A_	:						List[str]	=     ConvNextConfig(
			    depths=lowerCamelCase__						,      hidden_sizes=lowerCamelCase__						,      out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]    )
			A_	:						List[Any]	=     UperNetConfig(
			    backbone_config=lowerCamelCase__						,      auxiliary_in_channels=lowerCamelCase__						,      num_labels=lowerCamelCase__						,      idalabel=lowerCamelCase__						,      labelaid=lowerCamelCase__						,      )

			return config



def 	a			( lowerCamelCase__    ):
			'''simple docstring'''





			A_	:						Tuple	=     []

			# fmt: off
			# stem
			rename_keys.append(("""backbone.downsample_layers.0.0.weight""", """backbone.embeddings.patch_embeddings.weight""")    )
			rename_keys.append(("""backbone.downsample_layers.0.0.bias""", """backbone.embeddings.patch_embeddings.bias""")    )
			rename_keys.append(("""backbone.downsample_layers.0.1.weight""", """backbone.embeddings.layernorm.weight""")    )
			rename_keys.append(("""backbone.downsample_layers.0.1.bias""", """backbone.embeddings.layernorm.bias""")    )
			# stages
			for i in range(len(config.backbone_config.depths    )    ):
						for j in range(config.backbone_config.depths[i]    ):
									rename_keys.append((f'backbone.stages.{i}.{j}.gamma', f'backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter')    )
									rename_keys.append((f'backbone.stages.{i}.{j}.depthwise_conv.weight', f'backbone.encoder.stages.{i}.layers.{j}.dwconv.weight')    )
									rename_keys.append((f'backbone.stages.{i}.{j}.depthwise_conv.bias', f'backbone.encoder.stages.{i}.layers.{j}.dwconv.bias')    )
									rename_keys.append((f'backbone.stages.{i}.{j}.norm.weight', f'backbone.encoder.stages.{i}.layers.{j}.layernorm.weight')    )
									rename_keys.append((f'backbone.stages.{i}.{j}.norm.bias', f'backbone.encoder.stages.{i}.layers.{j}.layernorm.bias')    )
									rename_keys.append((f'backbone.stages.{i}.{j}.pointwise_conv1.weight', f'backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight')    )
									rename_keys.append((f'backbone.stages.{i}.{j}.pointwise_conv1.bias', f'backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias')    )
									rename_keys.append((f'backbone.stages.{i}.{j}.pointwise_conv2.weight', f'backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight')    )
									rename_keys.append((f'backbone.stages.{i}.{j}.pointwise_conv2.bias', f'backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias')    )
						if i > 0:
									rename_keys.append((f'backbone.downsample_layers.{i}.0.weight', f'backbone.encoder.stages.{i}.downsampling_layer.0.weight')    )
									rename_keys.append((f'backbone.downsample_layers.{i}.0.bias', f'backbone.encoder.stages.{i}.downsampling_layer.0.bias')    )
									rename_keys.append((f'backbone.downsample_layers.{i}.1.weight', f'backbone.encoder.stages.{i}.downsampling_layer.1.weight')    )
									rename_keys.append((f'backbone.downsample_layers.{i}.1.bias', f'backbone.encoder.stages.{i}.downsampling_layer.1.bias')    )

						rename_keys.append((f'backbone.norm{i}.weight', f'backbone.hidden_states_norms.stage{i+1}.weight')    )
						rename_keys.append((f'backbone.norm{i}.bias', f'backbone.hidden_states_norms.stage{i+1}.bias')    )

			# decode head
			rename_keys.extend(
			    [
			        ("""decode_head.conv_seg.weight""", """decode_head.classifier.weight"""),
			        ("""decode_head.conv_seg.bias""", """decode_head.classifier.bias"""),
			        ("""auxiliary_head.conv_seg.weight""", """auxiliary_head.classifier.weight"""),
			        ("""auxiliary_head.conv_seg.bias""", """auxiliary_head.classifier.bias"""),
			    ]    )
			# fmt: on

			return rename_keys



def 	a			( lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__    ):
			'''simple docstring'''





			A_	:						Tuple	=     dct.pop(lowerCamelCase__    )
			A_	:						int	=     val



def 	a			( lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__    ):
			'''simple docstring'''





			A_	:						Optional[Any]	=     {
			    """upernet-convnext-tiny""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth""",
			    """upernet-convnext-small""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth""",
			    """upernet-convnext-base""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth""",
			    """upernet-convnext-large""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth""",
			    """upernet-convnext-xlarge""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth""",
			}
			A_	:						Union[str, Any]	=     model_name_to_url[model_name]
			A_	:						List[str]	=     torch.hub.load_state_dict_from_url(lowerCamelCase__						,      map_location="""cpu"""    )["""state_dict"""]

			A_	:						str	=     get_upernet_config(lowerCamelCase__    )
			A_	:						Any	=     UperNetForSemanticSegmentation(lowerCamelCase__    )
			model.eval()

			# replace "bn" => "batch_norm"
			for key in state_dict.copy().keys():
						A_	:						Any	=     state_dict.pop(lowerCamelCase__    )
						if "bn" in key:
									A_	:						List[str]	=     key.replace("""bn"""						,      """batch_norm"""    )
						A_	:						Optional[Any]	=     val

			# rename keys
			A_	:						List[str]	=     create_rename_keys(lowerCamelCase__    )
			for src, dest in rename_keys:
						rename_key(lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__    )

			model.load_state_dict(lowerCamelCase__    )

			# verify on image
			A_	:						Any	=     """https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"""
			A_	:						List[Any]	=     Image.open(requests.get(lowerCamelCase__						,      stream=lowerCamelCase__    ).raw    ).convert("""RGB"""    )

			A_	:						List[Any]	=     SegformerImageProcessor()
			A_	:						Optional[int]	=     processor(lowerCamelCase__						,      return_tensors="""pt"""    ).pixel_values

			with torch.no_grad():
						A_	:						Any	=     model(lowerCamelCase__    )

			if model_name == "upernet-convnext-tiny":
						A_	:						Optional[Any]	=     torch.tensor(
						    [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]    )
			elif model_name == "upernet-convnext-small":
						A_	:						Union[str, Any]	=     torch.tensor(
						    [[-8.8_236, -8.8_236, -8.6_771], [-8.8_236, -8.8_236, -8.6_771], [-8.7_638, -8.7_638, -8.6_240]]    )
			elif model_name == "upernet-convnext-base":
						A_	:						int	=     torch.tensor(
						    [[-8.8_558, -8.8_558, -8.6_905], [-8.8_558, -8.8_558, -8.6_905], [-8.7_669, -8.7_669, -8.6_021]]    )
			elif model_name == "upernet-convnext-large":
						A_	:						List[Any]	=     torch.tensor(
						    [[-8.6_660, -8.6_660, -8.6_210], [-8.6_660, -8.6_660, -8.6_210], [-8.6_310, -8.6_310, -8.5_964]]    )
			elif model_name == "upernet-convnext-xlarge":
						A_	:						List[Any]	=     torch.tensor(
						    [[-8.4_980, -8.4_980, -8.3_977], [-8.4_980, -8.4_980, -8.3_977], [-8.4_379, -8.4_379, -8.3_412]]    )
			print("""Logits:"""						,      outputs.logits[0, 0, :3, :3]    )
			assert torch.allclose(outputs.logits[0, 0, :3, :3]						,      lowerCamelCase__						,      atol=1E-4    )
			print("""Looks ok!"""    )

			if pytorch_dump_folder_path is not None:
						print(f'Saving model {model_name} to {pytorch_dump_folder_path}'    )
						model.save_pretrained(lowerCamelCase__    )
						print(f'Saving processor to {pytorch_dump_folder_path}'    )
						processor.save_pretrained(lowerCamelCase__    )

			if push_to_hub:
						print(f'Pushing model and processor for {model_name} to hub'    )
						model.push_to_hub(f'openmmlab/{model_name}'    )
						processor.push_to_hub(f'openmmlab/{model_name}'    )


if __name__ == "__main__":
						lowerCamelCase							:List[str]				    =					argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '''--model_name''',
						    default='''upernet-convnext-tiny''',
						    type=str,
						    choices=[F"upernet-convnext-{size}" for size in ['''tiny''', '''small''', '''base''', '''large''', '''xlarge''']],
						    help='''Name of the ConvNext UperNet model you\'d like to convert.''',
						)
						parser.add_argument(
						    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
						)
						parser.add_argument(
						    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
						)

						lowerCamelCase							:Any				    =					parser.parse_args()
						convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




'''simple docstring'''




# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from .config import config_command_parser
from .config_args import default_config_file, load_config_from_file  # noqa: F401
from .default import default_command_parser
from .update import update_command_parser



def 	a			( lowerCamelCase__=None    ):
			'''simple docstring'''





			A_	:						int	=     argparse.ArgumentParser(add_help=lowerCamelCase__						,      allow_abbrev=lowerCamelCase__    )
			# The main config parser
			A_	:						int	=     config_command_parser(lowerCamelCase__    )
			# The subparser to add commands to
			A_	:						int	=     config_parser.add_subparsers(title="""subcommands"""						,      dest="""subcommand"""    )

			# Then add other parsers with the parent parser
			default_command_parser(lowerCamelCase__						,      parents=[parent_parser]    )
			update_command_parser(lowerCamelCase__						,      parents=[parent_parser]    )

			return config_parser



def 	a			( ):
			'''simple docstring'''





			A_	:						Optional[int]	=     get_config_parser()
			A_	:						List[str]	=     config_parser.parse_args()

			if not hasattr(lowerCamelCase__						,      """func"""    ):
						config_parser.print_help()
						exit(1    )

			# Run
			args.func(lowerCamelCase__    )


if __name__ == "__main__":
						main()