

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 	_snake_case		(    SCREAMING_SNAKE_CASE				)       ->      Optional[int]:

     """simple docstring"""







     _lowerCAmelCase   :					int        =    {}
     state_dict.pop("pixel_mean"		,							SCREAMING_SNAKE_CASE				)
     state_dict.pop("pixel_std"		,							SCREAMING_SNAKE_CASE				)

     _lowerCAmelCase   :					Dict        =    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   :					Any        =    key.replace(SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE				)

          if re.match(SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE				):
               _lowerCAmelCase   :					Tuple        =    int(re.match(SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE				).group(2				)				)
               if layer_nb == 0:
                    _lowerCAmelCase   :					str        =    key.replace("layers.0"		,							"proj_in"				)
               elif layer_nb == 1:
                    _lowerCAmelCase   :					Optional[int]        =    key.replace("layers.1"		,							"layers.0"				)
               elif layer_nb == 2:
                    _lowerCAmelCase   :					str        =    key.replace("layers.2"		,							"proj_out"				)

          _lowerCAmelCase   :					int        =    value

     _lowerCAmelCase   :					int        =    model_state_dict[
         "prompt_encoder.shared_embedding.positional_embedding"
     ]

     return model_state_dict







def 	_snake_case		(    SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE="ybelkada/segment-anything"				)       ->      Optional[int]:

     """simple docstring"""







     _lowerCAmelCase   :					Tuple        =    hf_hub_download(SCREAMING_SNAKE_CASE		,							f'''checkpoints/{model_name}.pth'''				)

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

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

          _lowerCAmelCase   :					List[str]        =    SamConfig(
              vision_config=SCREAMING_SNAKE_CASE		,							)

     _lowerCAmelCase   :					str        =    torch.load(SCREAMING_SNAKE_CASE		,							map_location="cpu"				)
     _lowerCAmelCase   :					Optional[int]        =    replace_keys(SCREAMING_SNAKE_CASE				)

     _lowerCAmelCase   :					Union[str, Any]        =    SamImageProcessor()

     _lowerCAmelCase   :					List[str]        =    SamProcessor(image_processor=SCREAMING_SNAKE_CASE				)
     _lowerCAmelCase   :					Union[str, Any]        =    SamModel(SCREAMING_SNAKE_CASE				)

     hf_model.load_state_dict(SCREAMING_SNAKE_CASE				)
     _lowerCAmelCase   :					Tuple        =    hf_model.to("cuda"				)

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

     _lowerCAmelCase   :					Dict        =    [[[400, 650]]]
     _lowerCAmelCase   :					int        =    [[1]]

     _lowerCAmelCase   :					Union[str, Any]        =    processor(images=np.array(SCREAMING_SNAKE_CASE				)		,							return_tensors="pt"				).to("cuda"				)

     with torch.no_grad():
          _lowerCAmelCase   :					Dict        =    hf_model(**SCREAMING_SNAKE_CASE				)
     _lowerCAmelCase   :					Union[str, Any]        =    output.iou_scores.squeeze()

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

          _lowerCAmelCase   :					Optional[int]        =    processor(
              images=np.array(SCREAMING_SNAKE_CASE				)		,							input_points=SCREAMING_SNAKE_CASE		,							input_labels=SCREAMING_SNAKE_CASE		,							return_tensors="pt"				).to("cuda"				)

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

          assert scores[-1].item() == 0.9_712_603_092_193_604

          _lowerCAmelCase   :					Dict        =    ((75, 275, 1725, 850),)

          _lowerCAmelCase   :					List[str]        =    processor(images=np.array(SCREAMING_SNAKE_CASE				)		,							input_boxes=SCREAMING_SNAKE_CASE		,							return_tensors="pt"				).to("cuda"				)

          with torch.no_grad():
               _lowerCAmelCase   :					Tuple        =    hf_model(**SCREAMING_SNAKE_CASE				)
          _lowerCAmelCase   :					List[str]        =    output.iou_scores.squeeze()

          assert scores[-1].item() == 0.8_686_015_605_926_514

          # Test with 2 points and 1 image.
          _lowerCAmelCase   :					Tuple        =    [[[400, 650], [800, 650]]]
          _lowerCAmelCase   :					Optional[int]        =    [[1, 1]]

          _lowerCAmelCase   :					Tuple        =    processor(
              images=np.array(SCREAMING_SNAKE_CASE				)		,							input_points=SCREAMING_SNAKE_CASE		,							input_labels=SCREAMING_SNAKE_CASE		,							return_tensors="pt"				).to("cuda"				)

          with torch.no_grad():
               _lowerCAmelCase   :					Tuple        =    hf_model(**SCREAMING_SNAKE_CASE				)
          _lowerCAmelCase   :					List[str]        =    output.iou_scores.squeeze()

          assert scores[-1].item() == 0.9_936_047_792_434_692


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)









from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_tf_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_tf_available():
  import tensorflow as tf


__UpperCAmelCase    				=     logging.get_logger(__name__)





@dataclass
class A__    ( A ):



 """simple docstring"""







 _lowercase     :					List[Any]      =						[
     '''no_inference''',
     '''no_cuda''',
     '''no_tpu''',
     '''no_speed''',
     '''no_memory''',
     '''no_env_print''',
     '''no_multi_process''',
 ]

 def __init__(		self :	List[Any]     ,						**A_ :	str       ):






      '''simple docstring'''




      for deprecated_arg in self.deprecated_args:
           if deprecated_arg in kwargs:
                _lowerCAmelCase   :					str        =    deprecated_arg[3:]
                _lowerCAmelCase   :					List[Any]        =    not kwargs.pop(A_       )
                logger.warning(
                    F'''{deprecated_arg} is depreciated. Please use --no-{positive_arg} or'''
                    F''' {positive_arg}={kwargs[positive_arg]}'''       )
      _lowerCAmelCase   :					Any        =    kwargs.pop("tpu_name"     ,						self.tpu_name       )
      _lowerCAmelCase   :					str        =    kwargs.pop("device_idx"     ,						self.device_idx       )
      _lowerCAmelCase   :					str        =    kwargs.pop("eager_mode"     ,						self.eager_mode       )
      _lowerCAmelCase   :					Tuple        =    kwargs.pop("use_xla"     ,						self.use_xla       )
      super().__init__(**A_       )

 _lowercase     :					str      =						field(
     default=A  ,		metadata={'''help''': '''Name of TPU'''}  ,		)
 _lowercase     :					int      =						field(
     default=0  ,		metadata={'''help''': '''CPU / GPU device index. Defaults to 0.'''}  ,		)
 _lowercase     :					bool      =						field(default=A  ,		metadata={'''help''': '''Benchmark models in eager model.'''} )
 _lowercase     :					bool      =						field(
     default=A  ,		metadata={
         '''help''': '''Benchmark models using XLA JIT compilation. Note that `eager_model` has to be set to `False`.'''
     }  ,		)

 @cached_property
 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      _lowerCAmelCase   :					Any        =    None
      if self.tpu:
           try:
                if self.tpu_name:
                     _lowerCAmelCase   :					Union[str, Any]        =    tf.distribute.cluster_resolver.TPUClusterResolver(self.tpu_name       )
                else:
                     _lowerCAmelCase   :					List[str]        =    tf.distribute.cluster_resolver.TPUClusterResolver()
           except ValueError:
                _lowerCAmelCase   :					Optional[int]        =    None
      return tpu

 @cached_property
 def    __magic_name__       (		self :	Optional[int]       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      if self.is_tpu:
           tf.config.experimental_connect_to_cluster(self._setup_tpu       )
           tf.tpu.experimental.initialize_tpu_system(self._setup_tpu       )

           _lowerCAmelCase   :					List[Any]        =    tf.distribute.TPUStrategy(self._setup_tpu       )
      else:
           # currently no multi gpu is allowed
           if self.is_gpu:
                # TODO: Currently only single GPU is supported
                tf.config.set_visible_devices(self.gpu_list[self.device_idx]     ,						"GPU"       )
                _lowerCAmelCase   :					Dict        =    tf.distribute.OneDeviceStrategy(device=F'''/gpu:{self.device_idx}'''       )
           else:
                tf.config.set_visible_devices([]     ,						"GPU"       )  # disable GPU
                _lowerCAmelCase   :					Optional[Any]        =    tf.distribute.OneDeviceStrategy(device=F'''/cpu:{self.device_idx}'''       )

      return strategy

 @property
 def    __magic_name__       (		self :	int       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      return self._setup_tpu is not None

 @property
 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      return self._setup_strategy

 @property
 def    __magic_name__       (		self :	List[str]       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      return tf.config.list_physical_devices("GPU"       )

 @property
 def    __magic_name__       (		self :	Optional[Any]       ):






      '''simple docstring'''




      requires_backends(self     ,						["tf"]       )
      if self.cuda:
           return len(self.gpu_list       )
      return 0


 @property
 def    __magic_name__       (		self :	Optional[Any]       ):






      '''simple docstring'''




      return self.n_gpu > 0






