



"""simple docstring"""






import argparse
import re

import numpy as np
import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SamConfig,
    SamImageProcessor,
    SamModel,
    SamProcessor,
    SamVisionConfig,
)


_lowercase				:				Union[str, Any]   		=		{
    "iou_prediction_head.layers.0": "iou_prediction_head.proj_in",
    "iou_prediction_head.layers.1": "iou_prediction_head.layers.0",
    "iou_prediction_head.layers.2": "iou_prediction_head.proj_out",
    "mask_decoder.output_upscaling.0": "mask_decoder.upscale_conv1",
    "mask_decoder.output_upscaling.1": "mask_decoder.upscale_layer_norm",
    "mask_decoder.output_upscaling.3": "mask_decoder.upscale_conv2",
    "mask_downscaling.0": "mask_embed.conv1",
    "mask_downscaling.1": "mask_embed.layer_norm1",
    "mask_downscaling.3": "mask_embed.conv2",
    "mask_downscaling.4": "mask_embed.layer_norm2",
    "mask_downscaling.6": "mask_embed.conv3",
    "point_embeddings": "point_embed",
    "pe_layer.positional_encoding_gaussian_matrix": "shared_embedding.positional_embedding",
    "image_encoder": "vision_encoder",
    "neck.0": "neck.conv1",
    "neck.1": "neck.layer_norm1",
    "neck.2": "neck.conv2",
    "neck.3": "neck.layer_norm2",
    "patch_embed.proj": "patch_embed.projection",
    ".norm": ".layer_norm",
    "blocks": "layers",
}






def   snake_case__     (	__lowerCamelCase       :					Dict			):
			"""simple docstring"""
			lowerCamelCase__	:	Union[str, Any]				={}
			state_dict.pop('''pixel_mean'''							,	lowerCamelCase__			)
			state_dict.pop('''pixel_std'''							,	lowerCamelCase__			)

			lowerCamelCase__	:	Union[str, Any]				=R".*.output_hypernetworks_mlps.(\d+).layers.(\d+).*"

			for key, value in state_dict.items():
						for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
									if key_to_modify in key:
												lowerCamelCase__	:	str				=key.replace(lowerCamelCase__							,	lowerCamelCase__			)

						if re.match(lowerCamelCase__							,	lowerCamelCase__			):
									lowerCamelCase__	:	Tuple				=int(re.match(lowerCamelCase__							,	lowerCamelCase__			).group(2			)			)
									if layer_nb == 0:
												lowerCamelCase__	:	List[str]				=key.replace('''layers.0'''							,	'''proj_in'''			)
									elif layer_nb == 1:
												lowerCamelCase__	:	int				=key.replace('''layers.1'''							,	'''layers.0'''			)
									elif layer_nb == 2:
												lowerCamelCase__	:	List[str]				=key.replace('''layers.2'''							,	'''proj_out'''			)

						lowerCamelCase__	:	Optional[Any]				=value

			lowerCamelCase__	:	str				=model_state_dict[
			    "prompt_encoder.shared_embedding.positional_embedding"
			]

			return model_state_dict






def   snake_case__     (	__lowerCamelCase       :					Any							,	__lowerCamelCase       :					List[Any]							,	__lowerCamelCase       :					Tuple							,	__lowerCamelCase       :					str="ybelkada/segment-anything"			):
			"""simple docstring"""
			lowerCamelCase__	:	Dict				=hf_hub_download(lowerCamelCase__							,	f'''checkpoints/{model_name}.pth'''			)

			if "sam_vit_b" in model_name:
						lowerCamelCase__	:	Optional[int]				=SamConfig()
			elif "sam_vit_l" in model_name:
						lowerCamelCase__	:	Dict				=SamVisionConfig(
						    hidden_size=1024							,	num_hidden_layers=24							,	num_attention_heads=16							,	global_attn_indexes=[5, 11, 17, 23]							,	)

						lowerCamelCase__	:	Optional[int]				=SamConfig(
						    vision_config=lowerCamelCase__							,	)
			elif "sam_vit_h" in model_name:
						lowerCamelCase__	:	List[str]				=SamVisionConfig(
						    hidden_size=1280							,	num_hidden_layers=32							,	num_attention_heads=16							,	global_attn_indexes=[7, 15, 23, 31]							,	)

						lowerCamelCase__	:	Any				=SamConfig(
						    vision_config=lowerCamelCase__							,	)

			lowerCamelCase__	:	Union[str, Any]				=torch.load(lowerCamelCase__							,	map_location='''cpu'''			)
			lowerCamelCase__	:	int				=replace_keys(lowerCamelCase__			)

			lowerCamelCase__	:	int				=SamImageProcessor()

			lowerCamelCase__	:	Any				=SamProcessor(image_processor=lowerCamelCase__			)
			lowerCamelCase__	:	str				=SamModel(lowerCamelCase__			)

			hf_model.load_state_dict(lowerCamelCase__			)
			lowerCamelCase__	:	List[str]				=hf_model.to('''cuda'''			)

			lowerCamelCase__	:	List[Any]				="https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
			lowerCamelCase__	:	Tuple				=Image.open(requests.get(lowerCamelCase__							,	stream=lowerCamelCase__			).raw			).convert('''RGB'''			)

			lowerCamelCase__	:	List[str]				=[[[400, 650]]]
			lowerCamelCase__	:	Any				=[[1]]

			lowerCamelCase__	:	Any				=processor(images=np.array(lowerCamelCase__			)							,	return_tensors='''pt'''			).to('''cuda'''			)

			with torch.no_grad():
						lowerCamelCase__	:	str				=hf_model(**lowerCamelCase__			)
			lowerCamelCase__	:	int				=output.iou_scores.squeeze()

			if model_name == "sam_vit_h_4b8939":
						assert scores[-1].item() == 0.5_79_89_02_51_15_96_68

						lowerCamelCase__	:	Tuple				=processor(
						    images=np.array(lowerCamelCase__			)							,	input_points=lowerCamelCase__							,	input_labels=lowerCamelCase__							,	return_tensors='''pt'''			).to('''cuda'''			)

						with torch.no_grad():
									lowerCamelCase__	:	Union[str, Any]				=hf_model(**lowerCamelCase__			)
						lowerCamelCase__	:	Optional[Any]				=output.iou_scores.squeeze()

						assert scores[-1].item() == 0.97_12_60_30_92_19_36_04

						lowerCamelCase__	:	Any				=((75, 275, 1725, 850),)

						lowerCamelCase__	:	Any				=processor(images=np.array(lowerCamelCase__			)							,	input_boxes=lowerCamelCase__							,	return_tensors='''pt'''			).to('''cuda'''			)

						with torch.no_grad():
									lowerCamelCase__	:	Dict				=hf_model(**lowerCamelCase__			)
						lowerCamelCase__	:	Tuple				=output.iou_scores.squeeze()

						assert scores[-1].item() == 0.86_86_01_56_05_92_65_14

						# Test with 2 points and 1 image.
						lowerCamelCase__	:	str				=[[[400, 650], [800, 650]]]
						lowerCamelCase__	:	Tuple				=[[1, 1]]

						lowerCamelCase__	:	str				=processor(
						    images=np.array(lowerCamelCase__			)							,	input_points=lowerCamelCase__							,	input_labels=lowerCamelCase__							,	return_tensors='''pt'''			).to('''cuda'''			)

						with torch.no_grad():
									lowerCamelCase__	:	List[Any]				=hf_model(**lowerCamelCase__			)
						lowerCamelCase__	:	str				=output.iou_scores.squeeze()

						assert scores[-1].item() == 0.99_36_04_77_92_43_46_92


if __name__ == "__main__":
		_lowercase				:				List[Any]   		=		argparse.ArgumentParser()
		_lowercase				:				List[str]   		=		["sam_vit_b_01ec64", "sam_vit_h_4b8939", "sam_vit_l_0b3195"]
		parser.add_argument(
		    "--model_name",
		    default="sam_vit_h_4b8939",
		    choices=choices,
		    type=str,
		    help="Path to hf config.json of model to convert",
		)
		parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
		parser.add_argument(
		    "--push_to_hub",
		    action="store_true",
		    help="Whether to push the model and processor to the hub after converting",
		)
		parser.add_argument(
		    "--model_hub_id",
		    default="ybelkada/segment-anything",
		    choices=choices,
		    type=str,
		    help="Path to hf config.json of model to convert",
		)

		_lowercase				:				Dict   		=		parser.parse_args()

		convert_sam_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.model_hub_id)









"""simple docstring"""






import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import UperNetForSemanticSegmentation
		from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor

class  __SCREAMING_SNAKE_CASE			:


					'''simple docstring'''





					def __init__( self							:						Optional[Any],						lowerCamelCase							:						Tuple,						lowerCamelCase							:						List[str]=13,						lowerCamelCase							:						List[Any]=32,						lowerCamelCase							:						Dict=3,						lowerCamelCase							:						int=4,						lowerCamelCase							:						str=[10, 20, 30, 40],						lowerCamelCase							:						Any=[2, 2, 3, 2],						lowerCamelCase							:						int=True,						lowerCamelCase							:						int=True,						lowerCamelCase							:						str=37,						lowerCamelCase							:						Optional[int]="gelu",						lowerCamelCase							:						Optional[int]=10,						lowerCamelCase							:						Any=0.02,						lowerCamelCase							:						Union[str, Any]=["stage2", "stage3", "stage4"],						lowerCamelCase							:						Optional[int]=3,						lowerCamelCase							:						Tuple=None,						)->      List[str]:
								lowerCamelCase__	:	List[str]				=parent
								lowerCamelCase__	:	Tuple				=batch_size
								lowerCamelCase__	:	str				=image_size
								lowerCamelCase__	:	Any				=num_channels
								lowerCamelCase__	:	Tuple				=num_stages
								lowerCamelCase__	:	List[str]				=hidden_sizes
								lowerCamelCase__	:	Any				=depths
								lowerCamelCase__	:	Union[str, Any]				=is_training
								lowerCamelCase__	:	Tuple				=use_labels
								lowerCamelCase__	:	int				=intermediate_size
								lowerCamelCase__	:	Optional[int]				=hidden_act
								lowerCamelCase__	:	Dict				=type_sequence_label_size
								lowerCamelCase__	:	Tuple				=initializer_range
								lowerCamelCase__	:	Any				=out_features
								lowerCamelCase__	:	Tuple				=num_labels
								lowerCamelCase__	:	Optional[int]				=scope
								lowerCamelCase__	:	Optional[int]				=num_stages



					def 	snake_case ( self							:						str  )->      Optional[int]:
								lowerCamelCase__	:	str				=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

								lowerCamelCase__	:	Tuple				=None
								if self.use_labels:
											lowerCamelCase__	:	Union[str, Any]				=ids_tensor([self.batch_size],						self.type_sequence_label_size  )

								lowerCamelCase__	:	int				=self.get_config()

								return config, pixel_values, labels



					def 	snake_case ( self							:						Union[str, Any]  )->      Any:
								return ConvNextConfig(
								    num_channels=self.num_channels,						num_stages=self.num_stages,						hidden_sizes=self.hidden_sizes,						depths=self.depths,						is_training=self.is_training,						intermediate_size=self.intermediate_size,						hidden_act=self.hidden_act,						out_features=self.out_features,						)



					def 	snake_case ( self							:						Union[str, Any]  )->      Any:
								return UperNetConfig(
								    backbone_config=self.get_backbone_config(),						hidden_size=512,						pool_scales=[1, 2, 3, 6],						use_auxiliary_head=lowerCamelCase,						auxiliary_loss_weight=0.4,						auxiliary_in_channels=40,						auxiliary_channels=256,						auxiliary_num_convs=1,						auxiliary_concat_input=lowerCamelCase,						loss_ignore_index=255,						num_labels=self.num_labels,						)



					def 	snake_case ( self							:						int,						lowerCamelCase							:						str,						lowerCamelCase							:						List[str],						lowerCamelCase							:						List[Any]  )->      Tuple:
								lowerCamelCase__	:	List[str]				=UperNetForSemanticSegmentation(config=lowerCamelCase  )
								model.to(lowerCamelCase  )
								model.eval()
								lowerCamelCase__	:	int				=model(lowerCamelCase  )
								self.parent.assertEqual(
								    result.logits.shape,						(self.batch_size, self.num_labels, self.image_size, self.image_size)  )




					def 	snake_case ( self							:						Any  )->      Tuple:
								lowerCamelCase__	:	Dict				=self.prepare_config_and_inputs()
								(
								    (
								    lowerCamelCase__
								)			,     (
								    lowerCamelCase__
								)			,     (
								    lowerCamelCase__
								)			,     
								)	:	Any				=config_and_inputs
								lowerCamelCase__	:	Optional[int]				={'''pixel_values''': pixel_values}
								return config, inputs_dict





@require_torch
class  __SCREAMING_SNAKE_CASE			(   lowerCAmelCase_	,     lowerCAmelCase_	,     unittest.TestCase     ):


					'''simple docstring'''





					_a         =							(UperNetForSemanticSegmentation,) if is_torch_available() else ()
					_a         =							{'image-segmentation': UperNetForSemanticSegmentation} if is_torch_available() else {}
					_a         =							False
					_a         =							False
					_a         =							False
					_a         =							False
					_a         =							False
					_a         =							False



					def 	snake_case ( self							:						Optional[int]  )->      Optional[int]:
								lowerCamelCase__	:	Optional[Any]				=UperNetModelTester(self  )
								lowerCamelCase__	:	Union[str, Any]				=ConfigTester(self,						config_class=lowerCamelCase,						has_text_modality=lowerCamelCase,						hidden_size=37  )



					def 	snake_case ( self							:						Optional[int]  )->      Optional[int]:
								self.create_and_test_config_common_properties()
								self.config_tester.create_and_test_config_to_json_string()
								self.config_tester.create_and_test_config_to_json_file()
								self.config_tester.create_and_test_config_from_and_save_pretrained()
								self.config_tester.create_and_test_config_with_num_labels()
								self.config_tester.check_config_can_be_init_without_params()
								self.config_tester.check_config_arguments_init()



					def 	snake_case ( self							:						List[str]  )->      Dict:
								return



					def 	snake_case ( self							:						Optional[int]  )->      List[str]:
								lowerCamelCase__			,     lowerCamelCase__	:	str				=self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											lowerCamelCase__	:	Union[str, Any]				=model_class(lowerCamelCase  )
											lowerCamelCase__	:	Tuple				=inspect.signature(model.forward  )
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											lowerCamelCase__	:	Tuple				=[*signature.parameters.keys()]

											lowerCamelCase__	:	List[Any]				=['''pixel_values''']
											self.assertListEqual(arg_names[:1],						lowerCamelCase  )



					def 	snake_case ( self							:						Any  )->      Union[str, Any]:
								lowerCamelCase__	:	Dict				=self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_semantic_segmentation(*lowerCamelCase  )



					@unittest.skip(reason='''UperNet does not use inputs_embeds'''  )
					def 	snake_case ( self							:						Optional[Any]  )->      List[Any]:
								pass



					@unittest.skip(reason='''UperNet does not support input and output embeddings'''  )
					def 	snake_case ( self							:						Any  )->      List[str]:
								pass



					@unittest.skip(reason='''UperNet does not have a base model'''  )
					def 	snake_case ( self							:						int  )->      Any:
								pass



					@unittest.skip(reason='''UperNet does not have a base model'''  )
					def 	snake_case ( self							:						Dict  )->      str:
								pass



					@require_torch_multi_gpu
					@unittest.skip(reason='''UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''  )
					def 	snake_case ( self							:						List[Any]  )->      List[str]:
								pass



					@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''  )
					def 	snake_case ( self							:						Tuple  )->      str:
								pass



					def 	snake_case ( self							:						Optional[int]  )->      List[str]:
								def check_hidden_states_output(lowerCamelCase							:						Dict,						lowerCamelCase							:						int,						lowerCamelCase							:						List[str]  ):
											lowerCamelCase__	:	Union[str, Any]				=model_class(lowerCamelCase  )
											model.to(lowerCamelCase  )
											model.eval()

											with torch.no_grad():
														lowerCamelCase__	:	Optional[Any]				=model(**self._prepare_for_class(lowerCamelCase,						lowerCamelCase  )  )

											lowerCamelCase__	:	Optional[Any]				=outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

											lowerCamelCase__	:	List[str]				=self.model_tester.num_stages
											self.assertEqual(len(lowerCamelCase  ),						expected_num_stages + 1  )

											# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
											self.assertListEqual(
											    list(hidden_states[0].shape[-2:]  ),						[self.model_tester.image_size // 4, self.model_tester.image_size // 4],						)

								lowerCamelCase__			,     lowerCamelCase__	:	Dict				=self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											lowerCamelCase__	:	Optional[int]				=True
											check_hidden_states_output(lowerCamelCase,						lowerCamelCase,						lowerCamelCase  )

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											lowerCamelCase__	:	Optional[Any]				=True

											check_hidden_states_output(lowerCamelCase,						lowerCamelCase,						lowerCamelCase  )



					def 	snake_case ( self							:						Any  )->      List[Any]:
								lowerCamelCase__			,     lowerCamelCase__	:	Optional[int]				=self.model_tester.prepare_config_and_inputs_for_common()

								lowerCamelCase__	:	str				=_config_zero_init(lowerCamelCase  )
								lowerCamelCase__	:	Union[str, Any]				=_config_zero_init(configs_no_init.backbone_config  )
								for model_class in self.all_model_classes:
											lowerCamelCase__	:	Optional[int]				=model_class(config=lowerCamelCase  )
											for name, param in model.named_parameters():
														if param.requires_grad:
																	self.assertIn(
																	    ((param.data.mean() * 1E9).round() / 1E9).item(),						[0.0, 1.0],						msg=F'''Parameter {name} of model {model_class} seems not properly initialized''',						)



					@unittest.skip(reason='''UperNet does not have tied weights'''  )
					def 	snake_case ( self							:						Any  )->      str:
								pass




					@slow
					def 	snake_case ( self							:						int  )->      Union[str, Any]:
								for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											lowerCamelCase__	:	str				=UperNetForSemanticSegmentation.from_pretrained(lowerCamelCase  )
											self.assertIsNotNone(lowerCamelCase  )






def   snake_case__     (	):
			"""simple docstring"""
			lowerCamelCase__	:	Optional[int]				=hf_hub_download(
			    repo_id='''hf-internal-testing/fixtures_ade20k'''							,	repo_type='''dataset'''							,	filename='''ADE_val_00000001.jpg'''			)
			lowerCamelCase__	:	List[str]				=Image.open(__lowerCamelCase			).convert('''RGB'''			)
			return image

@require_torch
@require_vision
@slow
class  __SCREAMING_SNAKE_CASE			(   unittest.TestCase     ):


					'''simple docstring'''





					def 	snake_case ( self							:						str  )->      Union[str, Any]:
								lowerCamelCase__	:	List[Any]				=AutoImageProcessor.from_pretrained('''openmmlab/upernet-swin-tiny'''  )
								lowerCamelCase__	:	List[Any]				=UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-swin-tiny'''  ).to(lowerCamelCase  )

								lowerCamelCase__	:	List[Any]				=prepare_img()
								lowerCamelCase__	:	List[Any]				=processor(images=lowerCamelCase,						return_tensors='''pt'''  ).to(lowerCamelCase  )

								with torch.no_grad():
											lowerCamelCase__	:	List[Any]				=model(**lowerCamelCase  )

								lowerCamelCase__	:	Optional[int]				=torch.Size((1, model.config.num_labels, 512, 512)  )
								self.assertEqual(outputs.logits.shape,						lowerCamelCase  )

								lowerCamelCase__	:	Dict				=torch.tensor(
								    [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]]  ).to(lowerCamelCase  )
								self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3],						lowerCamelCase,						atol=1E-4  )  )




					def 	snake_case ( self							:						Optional[int]  )->      Optional[Any]:
								lowerCamelCase__	:	str				=AutoImageProcessor.from_pretrained('''openmmlab/upernet-convnext-tiny'''  )
								lowerCamelCase__	:	Tuple				=UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-convnext-tiny'''  ).to(lowerCamelCase  )

								lowerCamelCase__	:	Dict				=prepare_img()
								lowerCamelCase__	:	Any				=processor(images=lowerCamelCase,						return_tensors='''pt'''  ).to(lowerCamelCase  )

								with torch.no_grad():
											lowerCamelCase__	:	Any				=model(**lowerCamelCase  )

								lowerCamelCase__	:	Dict				=torch.Size((1, model.config.num_labels, 512, 512)  )
								self.assertEqual(outputs.logits.shape,						lowerCamelCase  )

								lowerCamelCase__	:	List[str]				=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]]  ).to(lowerCamelCase  )
								self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3],						lowerCamelCase,						atol=1E-4  )  )




