

"""simple docstring"""




import json
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
				import torch

				if is_vision_available():
								from transformers import OneFormerImageProcessor
								from transformers.models.oneformer.image_processing_oneformer import binary_mask_to_rle
								from transformers.models.oneformer.modeling_oneformer import OneFormerForUniversalSegmentationOutput

if is_vision_available():
				from PIL import Image





def 	lowercase__		(    snake_case_					:Union[str, Any]       ,						snake_case_					:Dict="shi-labs/oneformer_demo"  ):
	with open(hf_hub_download(_A       ,						_A       ,						repo_type='''dataset'''  )       ,						'''r'''  ) as f:
		__UpperCAmelCase							=      json.load(_A  )
	__UpperCAmelCase							=      {}
	__UpperCAmelCase							=      []
	__UpperCAmelCase							=      []
	for key, info in class_info.items():
		__UpperCAmelCase							=      info['''name''']
		class_names.append(info['''name''']  )
		if info["isthing"]:
			thing_ids.append(int(_A  )  )
	__UpperCAmelCase							=      thing_ids
	__UpperCAmelCase							=      class_names
	return metadata
class 			_UpperCAmelCase					(				unittest.TestCase    ):



	def __init__(   self			:    Tuple       ,  _lowercase			:    Optional[int]       ,  _lowercase			:    Any=7       ,  _lowercase			:    str=3       ,  _lowercase			:    Optional[int]=30       ,  _lowercase			:    List[str]=4_00       ,  _lowercase			:    Optional[int]=None       ,  _lowercase			:    Any=True       ,  _lowercase			:    int=True       ,  _lowercase			:    Any=[0.5, 0.5, 0.5]       ,  _lowercase			:    Optional[Any]=[0.5, 0.5, 0.5]       ,  _lowercase			:    Optional[int]=10       ,  _lowercase			:    Tuple=False       ,  _lowercase			:    Optional[int]=2_55       ,  _lowercase			:    Tuple="shi-labs/oneformer_demo"       ,  _lowercase			:    str="ade20k_panoptic.json"       ,  _lowercase			:    int=10       ,  ):
		__UpperCAmelCase							=      parent
		__UpperCAmelCase							=      batch_size
		__UpperCAmelCase							=      num_channels
		__UpperCAmelCase							=      min_resolution
		__UpperCAmelCase							=      max_resolution
		__UpperCAmelCase							=      do_resize
		__UpperCAmelCase							=      {'''shortest_edge''': 32, '''longest_edge''': 13_33} if size is None else size
		__UpperCAmelCase							=      do_normalize
		__UpperCAmelCase							=      image_mean
		__UpperCAmelCase							=      image_std
		__UpperCAmelCase							=      class_info_file
		__UpperCAmelCase							=      prepare_metadata(UpperCamelCase__       ,  UpperCamelCase__		)
		__UpperCAmelCase							=      num_text
		__UpperCAmelCase							=      repo_path

		# for the post_process_functions
		__UpperCAmelCase							=      2
		__UpperCAmelCase							=      10
		__UpperCAmelCase							=      10
		__UpperCAmelCase							=      3
		__UpperCAmelCase							=      4
		__UpperCAmelCase							=      num_labels
		__UpperCAmelCase							=      do_reduce_labels
		__UpperCAmelCase							=      ignore_index



	def     a   (   self			:    Tuple		):
		return {
		    "do_resize": self.do_resize,
		    "size": self.size,
		    "do_normalize": self.do_normalize,
		    "image_mean": self.image_mean,
		    "image_std": self.image_std,
		    "num_labels": self.num_labels,
		    "do_reduce_labels": self.do_reduce_labels,
		    "ignore_index": self.ignore_index,
		    "class_info_file": self.class_info_file,
		    "metadata": self.metadata,
		    "num_text": self.num_text,
		}



	def     a   (   self			:    Union[str, Any]       ,  _lowercase			:    str       ,  _lowercase			:    int=False		):



		if not batched:
			__UpperCAmelCase							=      image_inputs[0]
			if isinstance(UpperCamelCase__       ,  Image.Image		):
				__UpperCAmelCase						,				__UpperCAmelCase							=      image.size
			else:
				__UpperCAmelCase						,				__UpperCAmelCase							=      image.shape[1], image.shape[2]
			if w < h:
				__UpperCAmelCase							=      int(self.size['''shortest_edge'''] * h / w		)
				__UpperCAmelCase							=      self.size['''shortest_edge''']
			elif w > h:
				__UpperCAmelCase							=      self.size['''shortest_edge''']
				__UpperCAmelCase							=      int(self.size['''shortest_edge'''] * w / h		)
			else:
				__UpperCAmelCase							=      self.size['''shortest_edge''']
				__UpperCAmelCase							=      self.size['''shortest_edge''']

		else:
			__UpperCAmelCase							=      []
			for image in image_inputs:
				__UpperCAmelCase						,				__UpperCAmelCase							=      self.get_expected_values([image]		)
				expected_values.append((expected_height, expected_width)		)
			__UpperCAmelCase							=      max(UpperCamelCase__       ,  key=lambda _lowercase		: item[0]		)[0]
			__UpperCAmelCase							=      max(UpperCamelCase__       ,  key=lambda _lowercase		: item[1]		)[1]

		return expected_height, expected_width



	def     a   (   self			:    Tuple		):
		return OneFormerForUniversalSegmentationOutput(
		    # +1 for null class
		    class_queries_logits=torch.randn((self.batch_size, self.num_queries, self.num_classes + 1)		)       ,  masks_queries_logits=torch.randn((self.batch_size, self.num_queries, self.height, self.width)		)       ,  )






@require_torch
@require_vision
class 			_UpperCAmelCase					(				__UpperCAmelCase						,  unittest.TestCase    ):
	a__						:      Union[str, Any]						       = OneFormerImageProcessor if (is_vision_available() and is_torch_available()) else None
	# only for test_image_processing_common.test_image_proc_to_json_string
	a__						:      str						       = image_processing_class



	def     a   (   self			:    List[Any]		):
		__UpperCAmelCase							=      OneFormerImageProcessorTester(self		)



	@property
	def     a   (   self			:    str		):
		return self.image_processing_tester.prepare_image_processor_dict()



	def     a   (   self			:    Optional[int]		):
		__UpperCAmelCase							=      self.image_processing_class(**self.image_processor_dict		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''image_mean'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''image_std'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''do_normalize'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''do_resize'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''size'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''ignore_index'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''class_info_file'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''num_text'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''repo_path'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''metadata'''		)		)
		self.assertTrue(hasattr(UpperCamelCase__       ,  '''do_reduce_labels'''		)		)



	def     a   (   self			:    Any		):
		pass



	def     a   (   self			:    Union[str, Any]		):
		__UpperCAmelCase							=      self.image_processing_class(**self.image_processor_dict		)
		# create random PIL images
		__UpperCAmelCase							=      prepare_image_inputs(self.image_processing_tester       ,  equal_resolution=UpperCamelCase__		)
		for image in image_inputs:
			self.assertIsInstance(UpperCamelCase__       ,  Image.Image		)

		# Test not batched input
		__UpperCAmelCase							=      image_processor(image_inputs[0]       ,  ['''semantic''']       ,  return_tensors='''pt'''		).pixel_values

		__UpperCAmelCase						,				__UpperCAmelCase							=      self.image_processing_tester.get_expected_values(UpperCamelCase__		)

		self.assertEqual(
		    encoded_images.shape       ,  (1, self.image_processing_tester.num_channels, expected_height, expected_width)       ,  )

		# Test batched
		__UpperCAmelCase						,				__UpperCAmelCase							=      self.image_processing_tester.get_expected_values(UpperCamelCase__       ,  batched=UpperCamelCase__		)

		__UpperCAmelCase							=      image_processor(
		    UpperCamelCase__       ,  ['''semantic'''] * len(UpperCamelCase__		)       ,  return_tensors='''pt'''		).pixel_values
		self.assertEqual(
		    encoded_images.shape       ,  (
		        self.image_processing_tester.batch_size,
		        self.image_processing_tester.num_channels,
		        expected_height,
		        expected_width,
		    )       ,  )



	def     a   (   self			:    List[Any]		):
		__UpperCAmelCase							=      self.image_processing_class(**self.image_processor_dict		)
		# create random numpy tensors
		__UpperCAmelCase							=      prepare_image_inputs(self.image_processing_tester       ,  equal_resolution=UpperCamelCase__       ,  numpify=UpperCamelCase__		)
		for image in image_inputs:
			self.assertIsInstance(UpperCamelCase__       ,  np.ndarray		)

		# Test not batched input
		__UpperCAmelCase							=      image_processor(image_inputs[0]       ,  ['''semantic''']       ,  return_tensors='''pt'''		).pixel_values

		__UpperCAmelCase						,				__UpperCAmelCase							=      self.image_processing_tester.get_expected_values(UpperCamelCase__		)

		self.assertEqual(
		    encoded_images.shape       ,  (1, self.image_processing_tester.num_channels, expected_height, expected_width)       ,  )

		# Test batched
		__UpperCAmelCase						,				__UpperCAmelCase							=      self.image_processing_tester.get_expected_values(UpperCamelCase__       ,  batched=UpperCamelCase__		)

		__UpperCAmelCase							=      image_processor(
		    UpperCamelCase__       ,  ['''semantic'''] * len(UpperCamelCase__		)       ,  return_tensors='''pt'''		).pixel_values
		self.assertEqual(
		    encoded_images.shape       ,  (
		        self.image_processing_tester.batch_size,
		        self.image_processing_tester.num_channels,
		        expected_height,
		        expected_width,
		    )       ,  )



	def     a   (   self			:    List[str]		):
		__UpperCAmelCase							=      self.image_processing_class(**self.image_processor_dict		)
		# create random PyTorch tensors
		__UpperCAmelCase							=      prepare_image_inputs(self.image_processing_tester       ,  equal_resolution=UpperCamelCase__       ,  torchify=UpperCamelCase__		)
		for image in image_inputs:
			self.assertIsInstance(UpperCamelCase__       ,  torch.Tensor		)

		# Test not batched input
		__UpperCAmelCase							=      image_processor(image_inputs[0]       ,  ['''semantic''']       ,  return_tensors='''pt'''		).pixel_values

		__UpperCAmelCase						,				__UpperCAmelCase							=      self.image_processing_tester.get_expected_values(UpperCamelCase__		)

		self.assertEqual(
		    encoded_images.shape       ,  (1, self.image_processing_tester.num_channels, expected_height, expected_width)       ,  )

		# Test batched
		__UpperCAmelCase						,				__UpperCAmelCase							=      self.image_processing_tester.get_expected_values(UpperCamelCase__       ,  batched=UpperCamelCase__		)

		__UpperCAmelCase							=      image_processor(
		    UpperCamelCase__       ,  ['''semantic'''] * len(UpperCamelCase__		)       ,  return_tensors='''pt'''		).pixel_values
		self.assertEqual(
		    encoded_images.shape       ,  (
		        self.image_processing_tester.batch_size,
		        self.image_processing_tester.num_channels,
		        expected_height,
		        expected_width,
		    )       ,  )



	def     a   (   self			:    Optional[Any]       ,  _lowercase			:    Optional[int]=False       ,  _lowercase			:    List[str]=False       ,  _lowercase			:    Optional[int]="np"		):
		__UpperCAmelCase							=      self.image_processing_class(**self.image_processor_dict		)
		# prepare image and target
		__UpperCAmelCase							=      self.image_processing_tester.num_labels
		__UpperCAmelCase							=      None
		__UpperCAmelCase							=      None
		__UpperCAmelCase							=      prepare_image_inputs(self.image_processing_tester       ,  equal_resolution=UpperCamelCase__		)
		if with_segmentation_maps:
			__UpperCAmelCase							=      num_labels
			if is_instance_map:
				__UpperCAmelCase							=      list(range(UpperCamelCase__		)		) * 2
				__UpperCAmelCase							=      dict(enumerate(UpperCamelCase__		)		)
			__UpperCAmelCase							=      [
			    np.random.randint(0       ,  high * 2       ,  (img.size[1], img.size[0])		).astype(np.uinta		) for img in image_inputs
			]
			if segmentation_type == "pil":
				__UpperCAmelCase							=      [Image.fromarray(UpperCamelCase__		) for annotation in annotations]

		__UpperCAmelCase							=      image_processor(
		    UpperCamelCase__       ,  ['''semantic'''] * len(UpperCamelCase__		)       ,  UpperCamelCase__       ,  return_tensors='''pt'''       ,  instance_id_to_semantic_id=UpperCamelCase__       ,  pad_and_return_pixel_mask=UpperCamelCase__       ,  )

		return inputs



	def     a   (   self			:    Optional[Any]		):
		pass



	def     a   (   self			:    int		):



		def common(_lowercase			:    List[str]=False       ,  _lowercase			:    int=None		):
			__UpperCAmelCase							=      self.comm_get_image_processor_inputs(
			    with_segmentation_maps=UpperCamelCase__       ,  is_instance_map=UpperCamelCase__       ,  segmentation_type=UpperCamelCase__		)

			__UpperCAmelCase							=      inputs['''mask_labels''']
			__UpperCAmelCase							=      inputs['''class_labels''']
			__UpperCAmelCase							=      inputs['''pixel_values''']
			__UpperCAmelCase							=      inputs['''text_inputs''']

			# check the batch_size
			for mask_label, class_label, text_input in zip(UpperCamelCase__       ,  UpperCamelCase__       ,  UpperCamelCase__		):
				self.assertEqual(mask_label.shape[0]       ,  class_label.shape[0]		)
				# this ensure padding has happened
				self.assertEqual(mask_label.shape[1:]       ,  pixel_values.shape[2:]		)
				self.assertEqual(len(UpperCamelCase__		)       ,  self.image_processing_tester.num_text		)

		common()
		common(is_instance_map=UpperCamelCase__		)
		common(is_instance_map=UpperCamelCase__       ,  segmentation_type='''pil'''		)
		common(is_instance_map=UpperCamelCase__       ,  segmentation_type='''pil'''		)



	def     a   (   self			:    Optional[Any]		):
		__UpperCAmelCase							=      np.zeros((20, 50)		)
		__UpperCAmelCase							=      1
		__UpperCAmelCase							=      1
		__UpperCAmelCase							=      1

		__UpperCAmelCase							=      binary_mask_to_rle(UpperCamelCase__		)
		self.assertEqual(len(UpperCamelCase__		)       ,  4		)
		self.assertEqual(rle[0]       ,  21		)
		self.assertEqual(rle[1]       ,  45		)



	def     a   (   self			:    Any		):
		__UpperCAmelCase							=      self.image_processing_class(
		    num_labels=self.image_processing_tester.num_classes       ,  max_seq_length=77       ,  task_seq_length=77       ,  class_info_file='''ade20k_panoptic.json'''       ,  num_text=self.image_processing_tester.num_text       ,  repo_path='''shi-labs/oneformer_demo'''       ,  )
		__UpperCAmelCase							=      self.image_processing_tester.get_fake_oneformer_outputs()

		__UpperCAmelCase							=      fature_extractor.post_process_semantic_segmentation(UpperCamelCase__		)

		self.assertEqual(len(UpperCamelCase__		)       ,  self.image_processing_tester.batch_size		)
		self.assertEqual(
		    segmentation[0].shape       ,  (
		        self.image_processing_tester.height,
		        self.image_processing_tester.width,
		    )       ,  )

		__UpperCAmelCase							=      [(1, 4) for i in range(self.image_processing_tester.batch_size		)]
		__UpperCAmelCase							=      fature_extractor.post_process_semantic_segmentation(UpperCamelCase__       ,  target_sizes=UpperCamelCase__		)

		self.assertEqual(segmentation[0].shape       ,  target_sizes[0]		)



	def     a   (   self			:    Optional[Any]		):
		__UpperCAmelCase							=      self.image_processing_class(
		    num_labels=self.image_processing_tester.num_classes       ,  max_seq_length=77       ,  task_seq_length=77       ,  class_info_file='''ade20k_panoptic.json'''       ,  num_text=self.image_processing_tester.num_text       ,  repo_path='''shi-labs/oneformer_demo'''       ,  )
		__UpperCAmelCase							=      self.image_processing_tester.get_fake_oneformer_outputs()
		__UpperCAmelCase							=      image_processor.post_process_instance_segmentation(UpperCamelCase__       ,  threshold=0		)

		self.assertTrue(len(UpperCamelCase__		) == self.image_processing_tester.batch_size		)
		for el in segmentation:
			self.assertTrue('''segmentation''' in el		)
			self.assertTrue('''segments_info''' in el		)
			self.assertEqual(type(el['''segments_info''']		)       ,  UpperCamelCase__		)
			self.assertEqual(
			    el['''segmentation'''].shape       ,  (self.image_processing_tester.height, self.image_processing_tester.width)		)



	def     a   (   self			:    Optional[Any]		):
		__UpperCAmelCase							=      self.image_processing_class(
		    num_labels=self.image_processing_tester.num_classes       ,  max_seq_length=77       ,  task_seq_length=77       ,  class_info_file='''ade20k_panoptic.json'''       ,  num_text=self.image_processing_tester.num_text       ,  repo_path='''shi-labs/oneformer_demo'''       ,  )
		__UpperCAmelCase							=      self.image_processing_tester.get_fake_oneformer_outputs()
		__UpperCAmelCase							=      image_processor.post_process_panoptic_segmentation(UpperCamelCase__       ,  threshold=0		)

		self.assertTrue(len(UpperCamelCase__		) == self.image_processing_tester.batch_size		)
		for el in segmentation:
			self.assertTrue('''segmentation''' in el		)
			self.assertTrue('''segments_info''' in el		)
			self.assertEqual(type(el['''segments_info''']		)       ,  UpperCamelCase__		)
			self.assertEqual(
			    el['''segmentation'''].shape       ,  (self.image_processing_tester.height, self.image_processing_tester.width)		)






import argparse
import json

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

from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  [2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 18, 2]
						lowerCAmelCase_      			=  True if '''large''' in model_name or '''huge''' in model_name else False
						lowerCAmelCase_      			=  True if '''large''' in model_name or '''huge''' in model_name else False
						lowerCAmelCase_      			=  True if '''large''' in model_name or '''huge''' in model_name else False

						if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
												if "fl3" in model_name:
																		lowerCAmelCase_      			=  [3, 3, 3, 3]
																		lowerCAmelCase_      			=  [5, 5, 5, 5]
												elif "fl4" in model_name:
																		lowerCAmelCase_      			=  [4, 4, 4, 4]
																		lowerCAmelCase_      			=  [3, 3, 3, 3]

						if "tiny" in model_name or "small" in model_name or "base" in model_name:
												lowerCAmelCase_      			=  [3, 3, 3, 3]
												if "lrf" in model_name:
																		lowerCAmelCase_      			=  [3, 3, 3, 3]
												else:
																		lowerCAmelCase_      			=  [2, 2, 2, 2]

						if "tiny" in model_name:
												lowerCAmelCase_      			=  96
						elif "small" in model_name:
												lowerCAmelCase_      			=  96
						elif "base" in model_name:
												lowerCAmelCase_      			=  128
						elif "large" in model_name:
												lowerCAmelCase_      			=  192
						elif "xlarge" in model_name:
												lowerCAmelCase_      			=  256
						elif "huge" in model_name:
												lowerCAmelCase_      			=  352

						# set label information
						lowerCAmelCase_      			=  '''huggingface/label-files'''
						if "large" in model_name or "huge" in model_name:
												lowerCAmelCase_      			=  '''imagenet-22k-id2label.json'''
						else:
												lowerCAmelCase_      			=  '''imagenet-1k-id2label.json'''

						lowerCAmelCase_      			=  json.load(open(hf_hub_download(_A ,  _A ,  repo_type='''dataset'''  ) ,  '''r'''  )  )
						lowerCAmelCase_      			=  {int(_A  ): v for k, v in idalabel.items()}
						lowerCAmelCase_      			=  {v: k for k, v in idalabel.items()}

						lowerCAmelCase_      			=  FocalNetConfig(
						    embed_dim=_A ,  depths=_A ,  focal_levels=_A ,  focal_windows=_A ,  use_conv_embed=_A ,  idalabel=_A ,  labelaid=_A ,  use_post_layernorm=_A ,  use_layerscale=_A ,  )

						return config





def __UpperCamelCase			(       _A  ):
						if "patch_embed.proj" in name:
												lowerCAmelCase_      			=  name.replace('''patch_embed.proj''' ,  '''embeddings.patch_embeddings.projection'''  )
						if "patch_embed.norm" in name:
												lowerCAmelCase_      			=  name.replace('''patch_embed.norm''' ,  '''embeddings.norm'''  )
						if "layers" in name:
												lowerCAmelCase_      			=  '''encoder.''' + name
						if "encoder.layers" in name:
												lowerCAmelCase_      			=  name.replace('''encoder.layers''' ,  '''encoder.stages'''  )
						if "downsample.proj" in name:
												lowerCAmelCase_      			=  name.replace('''downsample.proj''' ,  '''downsample.projection'''  )
						if "blocks" in name:
												lowerCAmelCase_      			=  name.replace('''blocks''' ,  '''layers'''  )
						if "modulation.f.weight" in name or "modulation.f.bias" in name:
												lowerCAmelCase_      			=  name.replace('''modulation.f''' ,  '''modulation.projection_in'''  )
						if "modulation.h.weight" in name or "modulation.h.bias" in name:
												lowerCAmelCase_      			=  name.replace('''modulation.h''' ,  '''modulation.projection_context'''  )
						if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
												lowerCAmelCase_      			=  name.replace('''modulation.proj''' ,  '''modulation.projection_out'''  )

						if name == "norm.weight":
												lowerCAmelCase_      			=  '''layernorm.weight'''
						if name == "norm.bias":
												lowerCAmelCase_      			=  '''layernorm.bias'''

						if "head" in name:
												lowerCAmelCase_      			=  name.replace('''head''' ,  '''classifier'''  )
						else:
												lowerCAmelCase_      			=  '''focalnet.''' + name

						return name





def __UpperCamelCase			(       _A ,  _A ,  _A=False  ):
						# fmt: off
						lowerCAmelCase_      			=  {
						    '''focalnet-tiny''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth''',
						    '''focalnet-tiny-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth''',
						    '''focalnet-small''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth''',
						    '''focalnet-small-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth''',
						    '''focalnet-base''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth''',
						    '''focalnet-base-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth''',
						    '''focalnet-large-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth''',
						    '''focalnet-large-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth''',
						    '''focalnet-xlarge-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth''',
						    '''focalnet-xlarge-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth''',
						}
						# fmt: on

						lowerCAmelCase_      			=  model_name_to_url[model_name]
						print('''Checkpoint URL: ''' ,  _A  )
						lowerCAmelCase_      			=  torch.hub.load_state_dict_from_url(_A ,  map_location='''cpu'''  )['''model''']

						# rename keys
						for key in state_dict.copy().keys():
												lowerCAmelCase_      			=  state_dict.pop(_A  )
												lowerCAmelCase_      			=  val

						lowerCAmelCase_      			=  get_focalnet_config(_A  )
						lowerCAmelCase_      			=  FocalNetForImageClassification(_A  )
						model.eval()

						# load state dict
						model.load_state_dict(_A  )

						# verify conversion
						lowerCAmelCase_      			=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''

						lowerCAmelCase_      			=  BitImageProcessor(
						    do_resize=_A ,  size={'''shortest_edge''': 256} ,  resample=PILImageResampling.BILINEAR ,  do_center_crop=_A ,  crop_size=224 ,  do_normalize=_A ,  image_mean=_A ,  image_std=_A ,  )
						lowerCAmelCase_      			=  Image.open(requests.get(_A ,  stream=_A  ).raw  )
						lowerCAmelCase_      			=  processor(images=_A ,  return_tensors='''pt'''  )

						lowerCAmelCase_      			=  transforms.Compose(
						    [
						        transforms.Resize(256  ),
						        transforms.CenterCrop(224  ),
						        transforms.ToTensor(),
						        transforms.Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6] ,  std=[0.2_2_9, 0.2_2_4, 0.2_2_5]  ),
						    ]  )

						lowerCAmelCase_      			=  image_transforms(_A  ).unsqueeze(0  )

						# verify pixel_values
						assert torch.allclose(inputs.pixel_values ,  _A ,  atol=1E-4  )

						lowerCAmelCase_      			=  model(**_A  )

						lowerCAmelCase_      			=  outputs.logits.argmax(-1  ).item()
						print('''Predicted class:''' ,  model.config.idalabel[predicted_class_idx]  )

						print('''First values of logits:''' ,  outputs.logits[0, :3]  )

						if model_name == "focalnet-tiny":
												lowerCAmelCase_      			=  torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]  )
						elif model_name == "focalnet-tiny-lrf":
												lowerCAmelCase_      			=  torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]  )
						elif model_name == "focalnet-small":
												lowerCAmelCase_      			=  torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]  )
						elif model_name == "focalnet-small-lrf":
												lowerCAmelCase_      			=  torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]  )
						elif model_name == "focalnet-base":
												lowerCAmelCase_      			=  torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]  )
						elif model_name == "focalnet-base-lrf":
												lowerCAmelCase_      			=  torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]  )
						assert torch.allclose(outputs.logits[0, :3] ,  _A ,  atol=1E-4  )
						print('''Looks ok!'''  )

						if pytorch_dump_folder_path is not None:
												print(f"Saving model and processor of {model_name} to {pytorch_dump_folder_path}"  )
												model.save_pretrained(_A  )
												processor.save_pretrained(_A  )

						if push_to_hub:
												print(f"Pushing model and processor of {model_name} to the hub..."  )
												model.push_to_hub(f"{model_name}"  )
												processor.push_to_hub(f"{model_name}"  )


if __name__ == "__main__":
						_A     =			argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '''--model_name''',
						    default='''focalnet-tiny''',
						    type=str,
						    help='''Name of the FocalNet model you\'d like to convert.''',
						)
						parser.add_argument(
						    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
						)
						parser.add_argument(
						    '''--push_to_hub''',
						    action='''store_true''',
						    help='''Whether to push the model and processor to the hub.''',
						)

						_A     =			parser.parse_args()
						convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




