




"""simple docstring"""

import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase	:      Union[str, Any]   =							logging.get_logger(__name__)

def 		lowercase_   (		_UpperCAmelCase						):




	"""simple docstring"""






	A_	:					Tuple						 =		ConvNextConfig()

	if "tiny" in checkpoint_url:
		A_	:					Optional[Any]						 =		[3, 3, 9, 3]
		A_	:					Union[str, Any]						 =		[96, 192, 384, 768]
	if "small" in checkpoint_url:
		A_	:					Union[str, Any]						 =		[3, 3, 27, 3]
		A_	:					Union[str, Any]						 =		[96, 192, 384, 768]
	if "base" in checkpoint_url:
		A_	:					Optional[Any]						 =		[3, 3, 27, 3]
		A_	:					Optional[Any]						 =		[128, 256, 512, 1024]
	if "large" in checkpoint_url:
		A_	:					Any						 =		[3, 3, 27, 3]
		A_	:					int						 =		[192, 384, 768, 1536]
	if "xlarge" in checkpoint_url:
		A_	:					Dict						 =		[3, 3, 27, 3]
		A_	:					int						 =		[256, 512, 1024, 2048]

	if "1k" in checkpoint_url:
		A_	:					str						 =		1000
		A_	:					str						 =		'''imagenet-1k-id2label.json'''
		A_	:					str						 =		(1, 1000)
	else:
		A_	:					Union[str, Any]						 =		2_1841
		A_	:					List[Any]						 =		'''imagenet-22k-id2label.json'''
		A_	:					Optional[Any]						 =		(1, 2_1841)

	A_	:					int						 =		'''huggingface/label-files'''
	A_	:					Tuple						 =		num_labels
	A_	:					List[str]						 =		json.load(open(hf_hub_download(_UpperCAmelCase	,   _UpperCAmelCase	,   repo_type='''dataset'''						)	,   '''r'''						)						)
	A_	:					List[str]						 =		{int(_UpperCAmelCase						): v for k, v in idalabel.items()}
	if "1k" not in checkpoint_url:
		# this dataset contains 21843 labels but the model only has 21841
		# we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
		del idalabel[9205]
		del idalabel[1_5027]
	A_	:					Union[str, Any]						 =		idalabel
	A_	:					int						 =		{v: k for k, v in idalabel.items()}
	A_	:					int						 =		hidden_sizes
	A_	:					Dict						 =		depths

	return config, expected_shape

def 		lowercase_   (		_UpperCAmelCase						):




	"""simple docstring"""






	if "downsample_layers.0.0" in name:
		A_	:					Any						 =		name.replace('''downsample_layers.0.0'''	,   '''embeddings.patch_embeddings'''						)
	if "downsample_layers.0.1" in name:
		A_	:					int						 =		name.replace('''downsample_layers.0.1'''	,   '''embeddings.norm'''						)  # we rename to layernorm later on
	if "downsample_layers.1.0" in name:
		A_	:					Any						 =		name.replace('''downsample_layers.1.0'''	,   '''stages.1.downsampling_layer.0'''						)
	if "downsample_layers.1.1" in name:
		A_	:					Dict						 =		name.replace('''downsample_layers.1.1'''	,   '''stages.1.downsampling_layer.1'''						)
	if "downsample_layers.2.0" in name:
		A_	:					Optional[Any]						 =		name.replace('''downsample_layers.2.0'''	,   '''stages.2.downsampling_layer.0'''						)
	if "downsample_layers.2.1" in name:
		A_	:					Any						 =		name.replace('''downsample_layers.2.1'''	,   '''stages.2.downsampling_layer.1'''						)
	if "downsample_layers.3.0" in name:
		A_	:					str						 =		name.replace('''downsample_layers.3.0'''	,   '''stages.3.downsampling_layer.0'''						)
	if "downsample_layers.3.1" in name:
		A_	:					List[str]						 =		name.replace('''downsample_layers.3.1'''	,   '''stages.3.downsampling_layer.1'''						)
	if "stages" in name and "downsampling_layer" not in name:
		# stages.0.0. for instance should be renamed to stages.0.layers.0.
		A_	:					Optional[Any]						 =		name[: len('''stages.0'''						)] + '''.layers''' + name[len('''stages.0'''						) :]
	if "stages" in name:
		A_	:					Dict						 =		name.replace('''stages'''	,   '''encoder.stages'''						)
	if "norm" in name:
		A_	:					int						 =		name.replace('''norm'''	,   '''layernorm'''						)
	if "gamma" in name:
		A_	:					Optional[Any]						 =		name.replace('''gamma'''	,   '''layer_scale_parameter'''						)
	if "head" in name:
		A_	:					Optional[int]						 =		name.replace('''head'''	,   '''classifier'''						)

	return name

def 		lowercase_   (		):




	"""simple docstring"""






	A_	:					Dict						 =		'''http://images.cocodataset.org/val2017/000000039769.jpg'''
	A_	:					int						 =		Image.open(requests.get(_UpperCAmelCase	,   stream=_UpperCAmelCase						).raw						)
	return im

@torch.no_grad()
def 		lowercase_   (		_UpperCAmelCase	,   _UpperCAmelCase						):




	"""simple docstring"""






	A_ ,			A_	:					Optional[Any]						 =		get_convnext_config(_UpperCAmelCase						)
	# load original state_dict from URL
	A_	:					int						 =		torch.hub.load_state_dict_from_url(_UpperCAmelCase						)['''model''']
	# rename keys
	for key in state_dict.copy().keys():
		A_	:					Union[str, Any]						 =		state_dict.pop(_UpperCAmelCase						)
		A_	:					Optional[int]						 =		val
	# add prefix to all keys expect classifier head
	for key in state_dict.copy().keys():
		A_	:					List[str]						 =		state_dict.pop(_UpperCAmelCase						)
		if not key.startswith('''classifier'''						):
			A_	:					str						 =		'''convnext.''' + key
		A_	:					List[Any]						 =		val

	# load HuggingFace model
	A_	:					Union[str, Any]						 =		ConvNextForImageClassification(_UpperCAmelCase						)
	model.load_state_dict(_UpperCAmelCase						)
	model.eval()

	# Check outputs on an image, prepared by ConvNextImageProcessor
	A_	:					int						 =		224 if '''224''' in checkpoint_url else 384
	A_	:					int						 =		ConvNextImageProcessor(size=_UpperCAmelCase						)
	A_	:					List[Any]						 =		image_processor(images=prepare_img()	,   return_tensors='''pt'''						).pixel_values

	A_	:					Dict						 =		model(_UpperCAmelCase						).logits

	# note: the logits below were obtained without center cropping
	if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
		A_	:					int						 =		torch.tensor([-0.1210, -0.6605, 0.1918]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
		A_	:					List[Any]						 =		torch.tensor([-0.4473, -0.1847, -0.6365]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
		A_	:					List[Any]						 =		torch.tensor([0.4525, 0.7539, 0.0308]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
		A_	:					Dict						 =		torch.tensor([0.3561, 0.6350, -0.0384]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
		A_	:					Tuple						 =		torch.tensor([0.4174, -0.0989, 0.1489]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
		A_	:					List[str]						 =		torch.tensor([0.2513, -0.1349, -0.1613]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
		A_	:					int						 =		torch.tensor([1.2980, 0.3631, -0.1198]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
		A_	:					Dict						 =		torch.tensor([1.2963, 0.1227, 0.1723]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
		A_	:					Dict						 =		torch.tensor([1.7956, 0.8390, 0.2820]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
		A_	:					List[Any]						 =		torch.tensor([-0.2822, -0.0502, -0.0878]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
		A_	:					Dict						 =		torch.tensor([-0.5672, -0.0730, -0.4348]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
		A_	:					List[Any]						 =		torch.tensor([0.2681, 0.2365, 0.6246]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
		A_	:					List[Any]						 =		torch.tensor([-0.2642, 0.3931, 0.5116]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
		A_	:					Optional[int]						 =		torch.tensor([-0.6677, -0.1873, -0.8379]						)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
		A_	:					Any						 =		torch.tensor([-0.7749, -0.2967, -0.6444]						)
	else:
		raise ValueError(f"""Unknown URL: {checkpoint_url}"""						)

	assert torch.allclose(logits[0, :3]	,   _UpperCAmelCase	,   atol=1E-3						)
	assert logits.shape == expected_shape

	Path(_UpperCAmelCase						).mkdir(exist_ok=_UpperCAmelCase						)
	print(f"""Saving model to {pytorch_dump_folder_path}"""						)
	model.save_pretrained(_UpperCAmelCase						)
	print(f"""Saving image processor to {pytorch_dump_folder_path}"""						)
	image_processor.save_pretrained(_UpperCAmelCase						)

	print('''Pushing model to the hub...'''						)
	A_	:					Any						 =		'''convnext'''
	if "tiny" in checkpoint_url:
		model_name += "-tiny"
	elif "small" in checkpoint_url:
		model_name += "-small"
	elif "base" in checkpoint_url:
		model_name += "-base"
	elif "xlarge" in checkpoint_url:
		model_name += "-xlarge"
	elif "large" in checkpoint_url:
		model_name += "-large"
	if "224" in checkpoint_url:
		model_name += "-224"
	elif "384" in checkpoint_url:
		model_name += "-384"
	if "22k" in checkpoint_url and "1k" not in checkpoint_url:
		model_name += "-22k"
	if "22k" in checkpoint_url and "1k" in checkpoint_url:
		model_name += "-22k-1k"

	model.push_to_hub(
	    repo_path_or_name=Path(_UpperCAmelCase	,   _UpperCAmelCase						)	,   organization='''nielsr'''	,   commit_message='''Add model'''	,   )


if __name__ == "__main__":
		_lowerCamelCase	:      Any   =							argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '--checkpoint_url',
		    default='https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth',
		    type=str,
		    help='URL of the original ConvNeXT checkpoint you\'d like to convert.',
		)
		parser.add_argument(
		    '--pytorch_dump_folder_path',
		    default=None,
		    type=str,
		    required=True,
		    help='Path to the output PyTorch model directory.',
		)

		_lowerCamelCase	:      List[str]   =							parser.parse_args()
		convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)











"""simple docstring"""

import unittest

from transformers import (
    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
    PreTrainedTokenizer,
    is_vision_available,
)
from transformers.pipelines import ImageClassificationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_or_tf,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
		from PIL import Image
else:

		class 							lowercase   :







						@staticmethod
						def        a_	(				*_lowerCamelCase :       Tuple   ,  **_lowerCamelCase :       str							):



							"""simple docstring"""

							pass

@is_pipeline_test
@require_torch_or_tf
@require_vision
class 							lowercase   (     unittest.TestCase):
				__lowerCAmelCase     :     str								= MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
				__lowerCAmelCase     :     Union[str, Any]								= TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING







				def        a_	(				self :       Tuple   ,  _lowerCamelCase :       Any   ,  _lowerCamelCase :       str   ,  _lowerCamelCase :       Optional[Any]							):



					"""simple docstring"""

					A_	:					List[str]						 =		ImageClassificationPipeline(model=_lowerCamelCase   ,  image_processor=_lowerCamelCase   ,  top_k=2							)
					A_	:					Tuple						 =		[
					    Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							),
					    '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					]
					return image_classifier, examples







				def        a_	(				self :       Optional[Any]   ,  _lowerCamelCase :       Optional[int]   ,  _lowerCamelCase :       Optional[Any]							):



					"""simple docstring"""

					A_	:					int						 =		image_classifier('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)

					self.assertEqual(
					    _lowerCamelCase   ,  [
					        {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					    ]   ,  )

					import datasets

					A_	:					Optional[Any]						 =		datasets.load_dataset('''hf-internal-testing/fixtures_image_utils'''   ,  '''image'''   ,  split='''test'''							)

					# Accepts URL + PIL.Image + lists
					A_	:					Union[str, Any]						 =		image_classifier(
					    [
					        Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							),
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					        # RGBA
					        dataset[0]['''file'''],
					        # LA
					        dataset[1]['''file'''],
					        # L
					        dataset[2]['''file'''],
					    ]							)
					self.assertEqual(
					    _lowerCamelCase   ,  [
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					    ]   ,  )







				@require_torch
				def        a_	(				self :       Dict							):



					"""simple docstring"""

					A_	:					Optional[int]						 =		'''hf-internal-testing/tiny-random-vit'''
					A_	:					Tuple						 =		pipeline('''image-classification'''   ,  model=_lowerCamelCase							)

					A_	:					Optional[int]						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}]   ,  )

					A_	:					str						 =		image_classifier(
					    [
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					    ]   ,  top_k=2   ,  )
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}],
					        [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}],
					    ]   ,  )







				@require_tf
				def        a_	(				self :       Tuple							):



					"""simple docstring"""

					A_	:					List[str]						 =		'''hf-internal-testing/tiny-random-vit'''
					A_	:					Optional[int]						 =		pipeline('''image-classification'''   ,  model=_lowerCamelCase   ,  framework='''tf'''							)

					A_	:					Tuple						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}]   ,  )

					A_	:					Optional[Any]						 =		image_classifier(
					    [
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					    ]   ,  top_k=2   ,  )
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}],
					        [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}],
					    ]   ,  )







				def        a_	(				self :       Dict							):



					"""simple docstring"""

					A_	:					Any						 =		PreTrainedTokenizer()

					# Assert that the pipeline can be initialized with a feature extractor that is not in any mapping
					A_	:					List[Any]						 =		pipeline(
					    '''image-classification'''   ,  model='''hf-internal-testing/tiny-random-vit'''   ,  tokenizer=_lowerCamelCase							)

					self.assertIs(image_classifier.tokenizer   ,  _lowerCamelCase							)







				@slow
				@require_torch
				def        a_	(				self :       List[str]							):



					"""simple docstring"""

					A_	:					List[Any]						 =		pipeline('''image-classification'''   ,  model='''deepmind/vision-perceiver-conv'''							)
					A_	:					Optional[Any]						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        {'''score''': 0.4385, '''label''': '''tabby, tabby cat'''},
					        {'''score''': 0.321, '''label''': '''tiger cat'''},
					        {'''score''': 0.0502, '''label''': '''Egyptian cat'''},
					        {'''score''': 0.0137, '''label''': '''crib, cot'''},
					        {'''score''': 0.007, '''label''': '''radiator'''},
					    ]   ,  )

					A_	:					Dict						 =		pipeline('''image-classification'''   ,  model='''deepmind/vision-perceiver-fourier'''							)
					A_	:					str						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        {'''score''': 0.5658, '''label''': '''tabby, tabby cat'''},
					        {'''score''': 0.1309, '''label''': '''tiger cat'''},
					        {'''score''': 0.0722, '''label''': '''Egyptian cat'''},
					        {'''score''': 0.0707, '''label''': '''remote control, remote'''},
					        {'''score''': 0.0082, '''label''': '''computer keyboard, keypad'''},
					    ]   ,  )

					A_	:					Optional[int]						 =		pipeline('''image-classification'''   ,  model='''deepmind/vision-perceiver-learned'''							)
					A_	:					int						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        {'''score''': 0.3022, '''label''': '''tabby, tabby cat'''},
					        {'''score''': 0.2362, '''label''': '''Egyptian cat'''},
					        {'''score''': 0.1856, '''label''': '''tiger cat'''},
					        {'''score''': 0.0324, '''label''': '''remote control, remote'''},
					        {'''score''': 0.0096, '''label''': '''quilt, comforter, comfort, puff'''},
					    ]   ,  )





