





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


_SCREAMING_SNAKE_CASE		       =							{
    "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(SCREAMING_SNAKE_CASE_     :    Optional[int]			):



    '''simple docstring'''





    _lowerCAmelCase			=					{}
    state_dict.pop("pixel_mean"						,				__snake_case			)
    state_dict.pop("pixel_std"						,				__snake_case			)

    _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(__snake_case						,				__snake_case			)

        if re.match(__snake_case						,				__snake_case			):
            _lowerCAmelCase			=					int(re.match(__snake_case						,				__snake_case			).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(SCREAMING_SNAKE_CASE_     :    str						,				SCREAMING_SNAKE_CASE_     :    Any						,				SCREAMING_SNAKE_CASE_     :    Any						,				SCREAMING_SNAKE_CASE_     :    Tuple="ybelkada/segment-anything"			):



    '''simple docstring'''





    _lowerCAmelCase			=					hf_hub_download(__snake_case						,				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=1024						,				num_hidden_layers=24						,				num_attention_heads=16						,				global_attn_indexes=[5, 11, 17, 23]						,				)

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

        _lowerCAmelCase			=					SamConfig(
            vision_config=__snake_case						,				)

    _lowerCAmelCase			=					torch.load(__snake_case						,				map_location="cpu"			)
    _lowerCAmelCase			=					replace_keys(__snake_case			)

    _lowerCAmelCase			=					SamImageProcessor()

    _lowerCAmelCase			=					SamProcessor(image_processor=__snake_case			)
    _lowerCAmelCase			=					SamModel(__snake_case			)

    hf_model.load_state_dict(__snake_case			)
    _lowerCAmelCase			=					hf_model.to("cuda"			)

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

    _lowerCAmelCase			=					[[[400, 650]]]
    _lowerCAmelCase			=					[[1]]

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

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

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

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

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

        assert scores[-1].item() == 0.9712_6030_9219_3604

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

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

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

        assert scores[-1].item() == 0.8686_0156_0592_6514

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

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

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

        assert scores[-1].item() == 0.9936_0477_9243_4692


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

       _SCREAMING_SNAKE_CASE		       =							parser.parse_args()

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









import pytest

from datasets.parallel import ParallelBackendConfig, parallel_backend
from datasets.utils.py_utils import map_nested

from .utils import require_dill_gt_0_3_2, require_joblibspark, require_not_windows





def lowerCAmelCase__(__snake_case							)				->      int:  # picklable for multiprocessing







	'''simple docstring'''







	return i + 1





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
def lowerCAmelCase__()				->      Any:







	'''simple docstring'''







	with parallel_backend('''spark'''							):
		assert ParallelBackendConfig.backend_name == "spark"

	lowerCamelCase__								=		[1, 2, 3]
	with pytest.raises(__snake_case							):
		with parallel_backend('''unsupported backend'''							):
			map_nested(__snake_case       ,__snake_case       ,num_proc=2							)

	with pytest.raises(__snake_case							):
		with parallel_backend('''unsupported backend'''							):
			map_nested(__snake_case       ,__snake_case       ,num_proc=-1							)





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
@pytest.mark.parametrize('''num_proc'''       ,[2, -1]							)
def lowerCAmelCase__(__snake_case							)				->      Tuple:







	'''simple docstring'''







	lowerCamelCase__								=		[1, 2]
	lowerCamelCase__								=		{'''a''': 1, '''b''': 2}
	lowerCamelCase__								=		{'''a''': [1, 2], '''b''': [3, 4]}
	lowerCamelCase__								=		{'''a''': {'''1''': 1}, '''b''': 2}
	lowerCamelCase__								=		{'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
	lowerCamelCase__								=		[2, 3]
	lowerCamelCase__								=		{'''a''': 2, '''b''': 3}
	lowerCamelCase__								=		{'''a''': [2, 3], '''b''': [4, 5]}
	lowerCamelCase__								=		{'''a''': {'''1''': 2}, '''b''': 3}
	lowerCamelCase__								=		{'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}

	with parallel_backend('''spark'''							):
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa



