






from sklearn.metrics import matthews_corrcoef

import datasets


_UpperCamelCase			:	str								='\nCompute the Matthews correlation coefficient (MCC)\n\nThe Matthews correlation coefficient is used in machine learning as a\nmeasure of the quality of binary and multiclass classifications. It takes\ninto account true and false positives and negatives and is generally\nregarded as a balanced measure which can be used even if the classes are of\nvery different sizes. The MCC is in essence a correlation coefficient value\nbetween -1 and +1. A coefficient of +1 represents a perfect prediction, 0\nan average random prediction and -1 an inverse prediction.  The statistic\nis also known as the phi coefficient. [source: Wikipedia]\n'

_UpperCamelCase			:	List[Any]								='\nArgs:\n    predictions (list of int): Predicted labels, as returned by a model.\n    references (list of int): Ground truth labels.\n    sample_weight (list of int, float, or bool): Sample weights. Defaults to `None`.\nReturns:\n    matthews_correlation (dict containing float): Matthews correlation.\nExamples:\n    Example 1, a basic example with only predictions and references as inputs:\n        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")\n        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],\n        ...                                     predictions=[1, 2, 2, 0, 3, 3])\n        >>> print(round(results[\'matthews_correlation\'], 2))\n        0.54\n\n    Example 2, the same example as above, but also including sample weights:\n        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")\n        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],\n        ...                                     predictions=[1, 2, 2, 0, 3, 3],\n        ...                                     sample_weight=[0.5, 3, 1, 1, 1, 2])\n        >>> print(round(results[\'matthews_correlation\'], 2))\n        0.1\n\n    Example 3, the same example as above, but with sample weights that cause a negative correlation:\n        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")\n        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],\n        ...                                     predictions=[1, 2, 2, 0, 3, 3],\n        ...                                     sample_weight=[0.5, 1, 0, 0, 0, 1])\n        >>> print(round(results[\'matthews_correlation\'], 2))\n        -0.25\n'

_UpperCamelCase			:	List[str]								='\\n@article{scikit-learn,\n  title={Scikit-learn: Machine Learning in {P}ython},\n  author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.\n         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.\n         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and\n         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},\n  journal={Journal of Machine Learning Research},\n  volume={12},\n  pages={2825--2830},\n  year={2011}\n}\n'






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION     ,				_KWARGS_DESCRIPTION							)
class 						UpperCAmelCase__ (							datasets.Metric							):







		def  A__      (					self		):
						return datasets.MetricInfo(
						    description=_DESCRIPTION       ,citation=_CITATION       ,inputs_description=_KWARGS_DESCRIPTION       ,features=datasets.Features(
						        {
						            '''predictions''': datasets.Value('''int32'''		),
						            '''references''': datasets.Value('''int32'''		),
						        }		)       ,reference_urls=[
						        '''https://scikit-learn.org/stable/modules/generated/sklearn.metrics.matthews_corrcoef.html'''
						    ]       ,)







		def  A__      (					self       ,A__       ,A__       ,A__=None		):
						return {
						    "matthews_correlation": float(matthews_corrcoef(__UpperCAmelCase       ,__UpperCAmelCase       ,sample_weight=__UpperCAmelCase		)		),
						}








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


a_								=       {
    """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__      (		_UpperCamelCase      :			Optional[Any]				):
   __lowerCamelCase    				=     {}
   state_dict.pop('''pixel_mean'''			,_UpperCamelCase				)
   state_dict.pop('''pixel_std'''			,_UpperCamelCase				)

   __lowerCamelCase    				=     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    				=     key.replace(_UpperCamelCase			,_UpperCamelCase				)

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

      __lowerCamelCase    				=     value

   __lowerCamelCase    				=     model_state_dict[
       '''prompt_encoder.shared_embedding.positional_embedding'''
   ]

   return model_state_dict





def 		a__      (		_UpperCamelCase      :			Tuple			,_UpperCamelCase      :			Dict			,_UpperCamelCase      :			Any			,_UpperCamelCase      :			Optional[int]="ybelkada/segment-anything"				):
   __lowerCamelCase    				=     hf_hub_download(_UpperCamelCase			,F"""checkpoints/{model_name}.pth"""				)

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

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

      __lowerCamelCase    				=     SamConfig(
          vision_config=_UpperCamelCase			,)

   __lowerCamelCase    				=     torch.load(_UpperCamelCase			,map_location='''cpu'''				)
   __lowerCamelCase    				=     replace_keys(_UpperCamelCase				)

   __lowerCamelCase    				=     SamImageProcessor()

   __lowerCamelCase    				=     SamProcessor(image_processor=_UpperCamelCase				)
   __lowerCamelCase    				=     SamModel(_UpperCamelCase				)

   hf_model.load_state_dict(_UpperCamelCase				)
   __lowerCamelCase    				=     hf_model.to('''cuda'''				)

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

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

   __lowerCamelCase    				=     processor(images=np.array(_UpperCamelCase				)			,return_tensors='''pt'''				).to('''cuda'''				)

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

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

      __lowerCamelCase    				=     processor(
          images=np.array(_UpperCamelCase				)			,input_points=_UpperCamelCase			,input_labels=_UpperCamelCase			,return_tensors='''pt'''				).to('''cuda'''				)

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

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

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

      __lowerCamelCase    				=     processor(images=np.array(_UpperCamelCase				)			,input_boxes=_UpperCamelCase			,return_tensors='''pt'''				).to('''cuda'''				)

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

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

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

      __lowerCamelCase    				=     processor(
          images=np.array(_UpperCamelCase				)			,input_points=_UpperCamelCase			,input_labels=_UpperCamelCase			,return_tensors='''pt'''				).to('''cuda'''				)

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

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


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

  a_								=       parser.parse_args()

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






