





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,
)


_UpperCAmelCase        =				{
    '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  lowerCAmelCase_ (							UpperCamelCase_  )     ->		List[str]:
       UpperCamelCase_          =  {}
       state_dict.pop("pixel_mean"					,		UpperCamelCase_  )
       state_dict.pop("pixel_std"					,		UpperCamelCase_  )

       UpperCamelCase_          =  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:
                            UpperCamelCase_          =  key.replace(UpperCamelCase_					,		UpperCamelCase_  )

              if re.match(UpperCamelCase_					,		UpperCamelCase_  ):
                     UpperCamelCase_          =  int(re.match(UpperCamelCase_					,		UpperCamelCase_  ).group(2  )  )
                     if layer_nb == 0:
                            UpperCamelCase_          =  key.replace("layers.0"					,		"proj_in"  )
                     elif layer_nb == 1:
                            UpperCamelCase_          =  key.replace("layers.1"					,		"layers.0"  )
                     elif layer_nb == 2:
                            UpperCamelCase_          =  key.replace("layers.2"					,		"proj_out"  )

              UpperCamelCase_          =  value

       UpperCamelCase_          =  model_state_dict[
           "prompt_encoder.shared_embedding.positional_embedding"
       ]

       return model_state_dict




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_="ybelkada/segment-anything"  )     ->		Tuple:
       UpperCamelCase_          =  hf_hub_download(UpperCamelCase_					,		F'''checkpoints/{model_name}.pth'''  )

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

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

              UpperCamelCase_          =  SamConfig(
                  vision_config=UpperCamelCase_					,		)

       UpperCamelCase_          =  torch.load(UpperCamelCase_					,		map_location="cpu"  )
       UpperCamelCase_          =  replace_keys(UpperCamelCase_  )

       UpperCamelCase_          =  SamImageProcessor()

       UpperCamelCase_          =  SamProcessor(image_processor=UpperCamelCase_  )
       UpperCamelCase_          =  SamModel(UpperCamelCase_  )

       hf_model.load_state_dict(UpperCamelCase_  )
       UpperCamelCase_          =  hf_model.to("cuda"  )

       UpperCamelCase_          =  "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
       UpperCamelCase_          =  Image.open(requests.get(UpperCamelCase_					,		stream=UpperCamelCase_  ).raw  ).convert("RGB"  )

       UpperCamelCase_          =  [[[400, 650]]]
       UpperCamelCase_          =  [[1]]

       UpperCamelCase_          =  processor(images=np.array(UpperCamelCase_  )					,		return_tensors="pt"  ).to("cuda"  )

       with torch.no_grad():
              UpperCamelCase_          =  hf_model(**UpperCamelCase_  )
       UpperCamelCase_          =  output.iou_scores.squeeze()

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

              UpperCamelCase_          =  processor(
                  images=np.array(UpperCamelCase_  )					,		input_points=UpperCamelCase_					,		input_labels=UpperCamelCase_					,		return_tensors="pt"  ).to("cuda"  )

              with torch.no_grad():
                     UpperCamelCase_          =  hf_model(**UpperCamelCase_  )
              UpperCamelCase_          =  output.iou_scores.squeeze()

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

              UpperCamelCase_          =  ((75, 275, 1725, 850),)

              UpperCamelCase_          =  processor(images=np.array(UpperCamelCase_  )					,		input_boxes=UpperCamelCase_					,		return_tensors="pt"  ).to("cuda"  )

              with torch.no_grad():
                     UpperCamelCase_          =  hf_model(**UpperCamelCase_  )
              UpperCamelCase_          =  output.iou_scores.squeeze()

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

              # Test with 2 points and 1 image.
              UpperCamelCase_          =  [[[400, 650], [800, 650]]]
              UpperCamelCase_          =  [[1, 1]]

              UpperCamelCase_          =  processor(
                  images=np.array(UpperCamelCase_  )					,		input_points=UpperCamelCase_					,		input_labels=UpperCamelCase_					,		return_tensors="pt"  ).to("cuda"  )

              with torch.no_grad():
                     UpperCamelCase_          =  hf_model(**UpperCamelCase_  )
              UpperCamelCase_          =  output.iou_scores.squeeze()

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


if __name__ == "__main__":
     _UpperCAmelCase        =				argparse.ArgumentParser()
     _UpperCAmelCase        =				['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',
     )

     _UpperCAmelCase        =				parser.parse_args()

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








import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate




def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		Optional[int]:
       return x + 2







class 							_UpperCamelCase					(	unittest.TestCase     ):

       def lowercase (		self: Optional[int]		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  "x = 3"
              UpperCamelCase_          =  {}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {}					,   state=_SCREAMING_SNAKE_CASE		)
              assert result == 3
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3}		)

              UpperCamelCase_          =  "x = y"
              UpperCamelCase_          =  {"y": 5}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {}					,   state=_SCREAMING_SNAKE_CASE		)
              # evaluate returns the value of the last assignment.
              assert result == 5
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 5, "y": 5}		)

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




              """simple docstring"""


              UpperCamelCase_          =  "y = add_two(x)"
              UpperCamelCase_          =  {"x": 3}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {"add_two": add_two}					,   state=_SCREAMING_SNAKE_CASE		)
              assert result == 5
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "y": 5}		)

              # Won't work without the tool
              with CaptureStdout() as out:
                     UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {}					,   state=_SCREAMING_SNAKE_CASE		)
              assert result is None
              assert "tried to execute add_two" in out.out

       def lowercase (		self: Union[str, Any]		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  "x = 3"
              UpperCamelCase_          =  {}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {}					,   state=_SCREAMING_SNAKE_CASE		)
              assert result == 3
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3}		)

       def lowercase (		self: Dict		)       -> Any:




              """simple docstring"""


              UpperCamelCase_          =  "test_dict = {'x': x, 'y': add_two(x)}"
              UpperCamelCase_          =  {"x": 3}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {"add_two": add_two}					,   state=_SCREAMING_SNAKE_CASE		)
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "y": 5}		)
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "test_dict": {"x": 3, "y": 5}}		)

       def lowercase (		self: List[str]		)       -> Union[str, Any]:




              """simple docstring"""


              UpperCamelCase_          =  "x = 3\ny = 5"
              UpperCamelCase_          =  {}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {}					,   state=_SCREAMING_SNAKE_CASE		)
              # evaluate returns the value of the last assignment.
              assert result == 5
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "y": 5}		)

       def lowercase (		self: List[Any]		)       -> int:




              """simple docstring"""


              UpperCamelCase_          =  "text = f'This is x: {x}.'"
              UpperCamelCase_          =  {"x": 3}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {}					,   state=_SCREAMING_SNAKE_CASE		)
              # evaluate returns the value of the last assignment.
              assert result == "This is x: 3."
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "text": "This is x: 3."}		)

       def lowercase (		self: int		)       -> Union[str, Any]:




              """simple docstring"""


              UpperCamelCase_          =  "if x <= 3:\n    y = 2\nelse:\n    y = 5"
              UpperCamelCase_          =  {"x": 3}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {}					,   state=_SCREAMING_SNAKE_CASE		)
              # evaluate returns the value of the last assignment.
              assert result == 2
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "y": 2}		)

              UpperCamelCase_          =  {"x": 8}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {}					,   state=_SCREAMING_SNAKE_CASE		)
              # evaluate returns the value of the last assignment.
              assert result == 5
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 8, "y": 5}		)

       def lowercase (		self: str		)       -> Optional[Any]:




              """simple docstring"""


              UpperCamelCase_          =  "test_list = [x, add_two(x)]"
              UpperCamelCase_          =  {"x": 3}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {"add_two": add_two}					,   state=_SCREAMING_SNAKE_CASE		)
              self.assertListEqual(_SCREAMING_SNAKE_CASE					,   [3, 5]		)
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "test_list": [3, 5]}		)

       def lowercase (		self: List[str]		)       -> int:




              """simple docstring"""


              UpperCamelCase_          =  "y = x"
              UpperCamelCase_          =  {"x": 3}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {}					,   state=_SCREAMING_SNAKE_CASE		)
              assert result == 3
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "y": 3}		)

       def lowercase (		self: str		)       -> int:




              """simple docstring"""


              UpperCamelCase_          =  "test_list = [x, add_two(x)]\ntest_list[1]"
              UpperCamelCase_          =  {"x": 3}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {"add_two": add_two}					,   state=_SCREAMING_SNAKE_CASE		)
              assert result == 5
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "test_list": [3, 5]}		)

              UpperCamelCase_          =  "test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"
              UpperCamelCase_          =  {"x": 3}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {"add_two": add_two}					,   state=_SCREAMING_SNAKE_CASE		)
              assert result == 5
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 3, "test_dict": {"x": 3, "y": 5}}		)



       def lowercase (		self: Dict		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_          =  "x = 0\nfor i in range(3):\n    x = i"
              UpperCamelCase_          =  {}
              UpperCamelCase_          =  evaluate(_SCREAMING_SNAKE_CASE					,   {"range": range}					,   state=_SCREAMING_SNAKE_CASE		)
              assert result == 2
              self.assertDictEqual(_SCREAMING_SNAKE_CASE					,   {"x": 2, "i": 2}		)

