

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


UpperCAmelCase :       List[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    a__       (  a__						):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      {}
     state_dict.pop("""pixel_mean"""      ,    a__						)
     state_dict.pop("""pixel_std"""      ,    a__						)

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

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

          __SCREAMING_SNAKE_CASE     =      value

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

     return model_state_dict







def    a__       (  a__      ,    a__      ,    a__      ,    a__="ybelkada/segment-anything"						):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      hf_hub_download(a__      ,    F'checkpoints/{model_name}.pth'						)

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

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

          __SCREAMING_SNAKE_CASE     =      SamConfig(
              vision_config=a__      ,    )

     __SCREAMING_SNAKE_CASE     =      torch.load(a__      ,    map_location="""cpu"""						)
     __SCREAMING_SNAKE_CASE     =      replace_keys(a__						)

     __SCREAMING_SNAKE_CASE     =      SamImageProcessor()

     __SCREAMING_SNAKE_CASE     =      SamProcessor(image_processor=a__						)
     __SCREAMING_SNAKE_CASE     =      SamModel(a__						)

     hf_model.load_state_dict(a__						)
     __SCREAMING_SNAKE_CASE     =      hf_model.to("""cuda"""						)

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

     __SCREAMING_SNAKE_CASE     =      [[[4_00, 6_50]]]
     __SCREAMING_SNAKE_CASE     =      [[1]]

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

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

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

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

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

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

          __SCREAMING_SNAKE_CASE     =      ((75, 2_75, 17_25, 8_50),)

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

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

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

          # Test with 2 points and 1 image.
          __SCREAMING_SNAKE_CASE     =      [[[4_00, 6_50], [8_00, 6_50]]]
          __SCREAMING_SNAKE_CASE     =      [[1, 1]]

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

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

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


if __name__ == "__main__":
  UpperCAmelCase :       Tuple			   =     argparse.ArgumentParser()
  UpperCAmelCase :       Optional[int]			   =     ['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 :       str			   =     parser.parse_args()

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









'''simple docstring'''



def    a__       (  a__						):




     """simple docstring"""


     if isinstance(a__      ,    a__						):
          raise TypeError("""'float' object cannot be interpreted as an integer"""						)
     if isinstance(a__      ,    a__						):
          raise TypeError("""'str' object cannot be interpreted as an integer"""						)

     if num == 0:
          return "0b0"

     __SCREAMING_SNAKE_CASE     =      False

     if num < 0:
          __SCREAMING_SNAKE_CASE     =      True
          __SCREAMING_SNAKE_CASE     =      -num

     __SCREAMING_SNAKE_CASE     =      []
     while num > 0:
          binary.insert(0      ,    num % 2						)
          num >>= 1

     if negative:
          return "-0b" + "".join(str(a__						) for e in binary						)

     return "0b" + "".join(str(a__						) for e in binary						)


if __name__ == "__main__":
  import doctest

  doctest.testmod()






