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


__magic_name__								={
    "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  __UpperCamelCase    (						A							):
     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  __UpperCamelCase    (						A ,							A ,							A ,							A="ybelkada/segment-anything"							):
     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__":
     __magic_name__								=argparse.ArgumentParser()
     __magic_name__								=["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''',
     )

     __magic_name__								=parser.parse_args()

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




import os
from typing import Optional

import fsspec
from fsspec.archive import AbstractArchiveFileSystem
from fsspec.utils import DEFAULT_BLOCK_SIZE



class     _A   (    __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	List[str]						=""
  SCREAMING_SNAKE_CASE_	:	str						=(
      None  # protocol passed in prefix to the url. ex: "gzip", for gzip://file.txt::http://foo.bar/file.txt.gz
  )
  SCREAMING_SNAKE_CASE_	:	str						=None  # compression type in fsspec. ex: "gzip"
  SCREAMING_SNAKE_CASE_	:	str						=None  # extension of the filename to strip. ex: "".gz" to get file.txt from file.txt.gz






  def __init__(self   ,				SCREAMING_SNAKE_CASE_ = ""   ,				SCREAMING_SNAKE_CASE_ = None   ,				SCREAMING_SNAKE_CASE_ = None   ,				**SCREAMING_SNAKE_CASE_	)		->      List[str]:





       '''simple docstring'''





       super().__init__(self   ,				**SCREAMING_SNAKE_CASE_	)
       # always open as "rb" since fsspec can then use the TextIOWrapper to make it work for "r" mode
       UpperCamelCase__										=      fsspec.open(
           SCREAMING_SNAKE_CASE_   ,				mode='''rb'''   ,				protocol=SCREAMING_SNAKE_CASE_   ,				compression=self.compression   ,				client_kwargs={
               '''requote_redirect_url''': False,  # see https://github.com/huggingface/datasets/pull/5459
               '''trust_env''': True,  # Enable reading proxy env variables.
               **(target_options or {}).pop('''client_kwargs'''   ,				{}	),  # To avoid issues if it was already passed.
           }   ,				**(target_options or {})   ,				)
       UpperCamelCase__										=      os.path.basename(self.file.path.split('''::'''	)[0]	)
       UpperCamelCase__										=      (
           self.compressed_name[: self.compressed_name.rindex('''.'''	)]
           if '''.''' in self.compressed_name
           else self.compressed_name
       )
       UpperCamelCase__										=      None






  @classmethod
  def 		_a					(cls   ,				SCREAMING_SNAKE_CASE_	)		->      Optional[Any]:





       '''simple docstring'''





       return super()._strip_protocol(SCREAMING_SNAKE_CASE_	).lstrip('''/'''	)






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





       '''simple docstring'''





       if self.dir_cache is None:
            UpperCamelCase__										=      {**self.file.fs.info(self.file.path	), '''name''': self.uncompressed_name}
            UpperCamelCase__										=      {f['''name''']: f}






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_	)		->      List[Any]:





       '''simple docstring'''





       return self.file.open().read()






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = "rb"   ,				SCREAMING_SNAKE_CASE_=None   ,				SCREAMING_SNAKE_CASE_=True   ,				SCREAMING_SNAKE_CASE_=None   ,				**SCREAMING_SNAKE_CASE_   ,				)		->      List[Any]:





       '''simple docstring'''





       UpperCamelCase__										=      self._strip_protocol(SCREAMING_SNAKE_CASE_	)
       if mode != "rb":
            raise ValueError(F"Tried to read with mode {mode} on file {self.file.path} opened with mode 'rb'"	)
       return self.file.open()



class     _A   (    __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	Tuple						="bz2"
  SCREAMING_SNAKE_CASE_	:	List[str]						="bz2"
  SCREAMING_SNAKE_CASE_	:	List[Any]						=".bz2"



class     _A   (    __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	Tuple						="gzip"
  SCREAMING_SNAKE_CASE_	:	Dict						="gzip"
  SCREAMING_SNAKE_CASE_	:	List[Any]						=".gz"



class     _A   (    __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	List[Any]						="lz4"
  SCREAMING_SNAKE_CASE_	:	str						="lz4"
  SCREAMING_SNAKE_CASE_	:	Any						=".lz4"



class     _A   (    __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	Dict						="xz"
  SCREAMING_SNAKE_CASE_	:	List[str]						="xz"
  SCREAMING_SNAKE_CASE_	:	Optional[int]						=".xz"



class     _A   (    __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	Optional[Any]						="zstd"
  SCREAMING_SNAKE_CASE_	:	Optional[int]						="zstd"
  SCREAMING_SNAKE_CASE_	:	Optional[int]						=".zst"






  def __init__(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = "rb"   ,				SCREAMING_SNAKE_CASE_ = None   ,				SCREAMING_SNAKE_CASE_ = None   ,				SCREAMING_SNAKE_CASE_ = DEFAULT_BLOCK_SIZE   ,				**SCREAMING_SNAKE_CASE_   ,				)		->      List[Any]:





       '''simple docstring'''





       super().__init__(
           fo=SCREAMING_SNAKE_CASE_   ,				mode=SCREAMING_SNAKE_CASE_   ,				target_protocol=SCREAMING_SNAKE_CASE_   ,				target_options=SCREAMING_SNAKE_CASE_   ,				block_size=SCREAMING_SNAKE_CASE_   ,				**SCREAMING_SNAKE_CASE_   ,				)
       # We need to wrap the zstd decompressor to avoid this error in fsspec==2021.7.0 and zstandard==0.15.2:
       #
       # File "/Users/user/.virtualenvs/hf-datasets/lib/python3.7/site-packages/fsspec/core.py", line 145, in open
       #     out.close = close
       # AttributeError: 'zstd.ZstdDecompressionReader' object attribute 'close' is read-only
       #
       # see https://github.com/intake/filesystem_spec/issues/725
       UpperCamelCase__										=      self.file.__enter__



       class     _A   :






         def __init__(self   ,				SCREAMING_SNAKE_CASE_	)		->      int:





              '''simple docstring'''





              UpperCamelCase__										=      file_






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





              '''simple docstring'''





              self._file.__enter__()
              return self






         def __exit__(self   ,				*SCREAMING_SNAKE_CASE_   ,				**SCREAMING_SNAKE_CASE_	)		->      List[Any]:





              '''simple docstring'''





              self._file.__exit__(*SCREAMING_SNAKE_CASE_   ,				**SCREAMING_SNAKE_CASE_	)






         def __iter__(self	)		->      Dict:





              '''simple docstring'''





              return iter(self._file	)






         def 		_a					(self	)		->      Dict:





              '''simple docstring'''





              return next(self._file	)






         def __getattr__(self   ,				SCREAMING_SNAKE_CASE_	)		->      str:





              '''simple docstring'''





              return getattr(self._file   ,				SCREAMING_SNAKE_CASE_	)

       def fixed_enter(*SCREAMING_SNAKE_CASE_   ,				**SCREAMING_SNAKE_CASE_	):
            return WrappedFile(_enter(*SCREAMING_SNAKE_CASE_   ,				**SCREAMING_SNAKE_CASE_	)	)

       UpperCamelCase__										=      fixed_enter



