






"""simple docstring"""




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 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  _a  :       Optional[Any]             =	[2, 2, 6, 2] if """tiny""" in model_name else [2, 2, 1_8, 2]
  _a  :       Any             =	True if """large""" in model_name or """huge""" in model_name else False
  _a  :       Union[str, Any]             =	True if """large""" in model_name or """huge""" in model_name else False
  _a  :       Any             =	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:
      _a  :       Union[str, Any]             =	[3, 3, 3, 3]
      _a  :       Dict             =	[5, 5, 5, 5]
    elif "fl4" in model_name:
      _a  :       int             =	[4, 4, 4, 4]
      _a  :       Tuple             =	[3, 3, 3, 3]

  if "tiny" in model_name or "small" in model_name or "base" in model_name:
    _a  :       Tuple             =	[3, 3, 3, 3]
    if "lrf" in model_name:
      _a  :       List[str]             =	[3, 3, 3, 3]
    else:
      _a  :       List[Any]             =	[2, 2, 2, 2]

  if "tiny" in model_name:
    _a  :       int             =	9_6
  elif "small" in model_name:
    _a  :       List[str]             =	9_6
  elif "base" in model_name:
    _a  :       Dict             =	1_2_8
  elif "large" in model_name:
    _a  :       Optional[int]             =	1_9_2
  elif "xlarge" in model_name:
    _a  :       List[Any]             =	2_5_6
  elif "huge" in model_name:
    _a  :       Optional[int]             =	3_5_2

  # set label information
  _a  :       Any             =	"""huggingface/label-files"""
  if "large" in model_name or "huge" in model_name:
    _a  :       str             =	"""imagenet-22k-id2label.json"""
  else:
    _a  :       Optional[Any]             =	"""imagenet-1k-id2label.json"""

  _a  :       Tuple             =	json.load(open(hf_hub_download(UpperCamelCase__     ,   UpperCamelCase__     ,   repo_type="""dataset"""				)     ,   """r"""				)				)
  _a  :       Dict             =	{int(UpperCamelCase__				): v for k, v in idalabel.items()}
  _a  :       Optional[int]             =	{v: k for k, v in idalabel.items()}

  _a  :       List[str]             =	FocalNetConfig(
      embed_dim=UpperCamelCase__     ,   depths=UpperCamelCase__     ,   focal_levels=UpperCamelCase__     ,   focal_windows=UpperCamelCase__     ,   use_conv_embed=UpperCamelCase__     ,   idalabel=UpperCamelCase__     ,   labelaid=UpperCamelCase__     ,   use_post_layernorm=UpperCamelCase__     ,   use_layerscale=UpperCamelCase__     ,   )

  return config

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  if "patch_embed.proj" in name:
    _a  :       List[Any]             =	name.replace("""patch_embed.proj"""     ,   """embeddings.patch_embeddings.projection"""				)
  if "patch_embed.norm" in name:
    _a  :       Optional[int]             =	name.replace("""patch_embed.norm"""     ,   """embeddings.norm"""				)
  if "layers" in name:
    _a  :       str             =	"""encoder.""" + name
  if "encoder.layers" in name:
    _a  :       Any             =	name.replace("""encoder.layers"""     ,   """encoder.stages"""				)
  if "downsample.proj" in name:
    _a  :       Any             =	name.replace("""downsample.proj"""     ,   """downsample.projection"""				)
  if "blocks" in name:
    _a  :       List[str]             =	name.replace("""blocks"""     ,   """layers"""				)
  if "modulation.f.weight" in name or "modulation.f.bias" in name:
    _a  :       Tuple             =	name.replace("""modulation.f"""     ,   """modulation.projection_in"""				)
  if "modulation.h.weight" in name or "modulation.h.bias" in name:
    _a  :       List[str]             =	name.replace("""modulation.h"""     ,   """modulation.projection_context"""				)
  if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
    _a  :       Optional[Any]             =	name.replace("""modulation.proj"""     ,   """modulation.projection_out"""				)

  if name == "norm.weight":
    _a  :       Dict             =	"""layernorm.weight"""
  if name == "norm.bias":
    _a  :       int             =	"""layernorm.bias"""

  if "head" in name:
    _a  :       Union[str, Any]             =	name.replace("""head"""     ,   """classifier"""				)
  else:
    _a  :       Optional[Any]             =	"""focalnet.""" + name

  return name

def 			lowerCAmelCase__     (     UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__=False				):



  '''simple docstring'''
  # fmt: off
  _a  :       Dict             =	{
      """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

  _a  :       Optional[Any]             =	model_name_to_url[model_name]
  print("""Checkpoint URL: """     ,   UpperCamelCase__				)
  _a  :       List[Any]             =	torch.hub.load_state_dict_from_url(UpperCamelCase__     ,   map_location="""cpu"""				)["""model"""]

  # rename keys
  for key in state_dict.copy().keys():
    _a  :       List[str]             =	state_dict.pop(UpperCamelCase__				)
    _a  :       Any             =	val

  _a  :       Union[str, Any]             =	get_focalnet_config(UpperCamelCase__				)
  _a  :       List[Any]             =	FocalNetForImageClassification(UpperCamelCase__				)
  model.eval()

  # load state dict
  model.load_state_dict(UpperCamelCase__				)

  # verify conversion
  _a  :       Tuple             =	"""http://images.cocodataset.org/val2017/000000039769.jpg"""

  _a  :       Optional[Any]             =	BitImageProcessor(
      do_resize=UpperCamelCase__     ,   size={"""shortest_edge""": 2_5_6}     ,   resample=PILImageResampling.BILINEAR     ,   do_center_crop=UpperCamelCase__     ,   crop_size=2_2_4     ,   do_normalize=UpperCamelCase__     ,   image_mean=UpperCamelCase__     ,   image_std=UpperCamelCase__     ,   )
  _a  :       Optional[Any]             =	Image.open(requests.get(UpperCamelCase__     ,   stream=UpperCamelCase__				).raw				)
  _a  :       Union[str, Any]             =	processor(images=UpperCamelCase__     ,   return_tensors="""pt"""				)

  _a  :       Tuple             =	transforms.Compose(
      [
          transforms.Resize(2_5_6				),
          transforms.CenterCrop(2_2_4				),
          transforms.ToTensor(),
          transforms.Normalize(mean=[0.485, 0.456, 0.406]     ,   std=[0.229, 0.224, 0.225]				),
      ]				)

  _a  :       Dict             =	image_transforms(UpperCamelCase__				).unsqueeze(0				)

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

  _a  :       Any             =	model(**UpperCamelCase__				)

  _a  :       Dict             =	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":
    _a  :       str             =	torch.tensor([0.2_166, -0.4_368, 0.2_191]				)
  elif model_name == "focalnet-tiny-lrf":
    _a  :       int             =	torch.tensor([1.1_669, 0.0_125, -0.1_695]				)
  elif model_name == "focalnet-small":
    _a  :       Optional[Any]             =	torch.tensor([0.4_917, -0.0_430, 0.1_341]				)
  elif model_name == "focalnet-small-lrf":
    _a  :       Optional[Any]             =	torch.tensor([-0.2_588, -0.5_342, -0.2_331]				)
  elif model_name == "focalnet-base":
    _a  :       Dict             =	torch.tensor([-0.1_655, -0.4_090, -0.1_730]				)
  elif model_name == "focalnet-base-lrf":
    _a  :       int             =	torch.tensor([0.5_306, -0.0_483, -0.3_928]				)
  assert torch.allclose(outputs.logits[0, :3]     ,   UpperCamelCase__     ,   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(UpperCamelCase__				)
    processor.save_pretrained(UpperCamelCase__				)

  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__":
 _snake_case          =   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.',
 )

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









"""simple docstring"""




import shutil
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_tf_cross_test,
    require_tf,
    require_torch,
    require_torchvision,
    require_vision,
)
from transformers.utils import is_tf_available, is_torch_available, is_vision_available


if is_vision_available():
 from PIL import Image

 from transformers import AutoProcessor, SamImageProcessor, SamProcessor

if is_torch_available():
 import torch

if is_tf_available():
 import tensorflow as tf






@require_vision
@require_torchvision
class        UpperCamelCase			(	unittest.TestCase      ):


 def _lowercase	(     self		:						List[Any]							)					->       Dict:
   _a  :       Optional[int]             =	tempfile.mkdtemp()
   _a  :       Optional[Any]             =	SamImageProcessor()
   _a  :       int             =	SamProcessor(UpperCAmelCase__							)
   processor.save_pretrained(self.tmpdirname							)


 def _lowercase	(     self		:						Tuple							,						**UpperCAmelCase__		:						Any							)					->       Any:
   return AutoProcessor.from_pretrained(self.tmpdirname							,						**UpperCAmelCase__							).image_processor


 def _lowercase	(     self		:						str							)					->       int:
   shutil.rmtree(self.tmpdirname							)


 def _lowercase	(     self		:						Tuple							)					->       Dict:
   _a  :       List[Any]             =	[np.random.randint(255							,						size=(3, 30, 400)							,						dtype=np.uinta							)]

   _a  :       Optional[int]             =	[Image.fromarray(np.moveaxis(UpperCAmelCase__							,						0							,						-1							)							) for x in image_inputs]

   return image_inputs


 def _lowercase	(     self		:						Dict							)					->       Dict:
   _a  :       List[Any]             =	SamProcessor(image_processor=self.get_image_processor()							)
   processor.save_pretrained(self.tmpdirname							)

   _a  :       Tuple             =	self.get_image_processor(do_normalize=UpperCAmelCase__							,						padding_value=1.0							)

   _a  :       Tuple             =	SamProcessor.from_pretrained(self.tmpdirname							,						do_normalize=UpperCAmelCase__							,						padding_value=1.0							)

   self.assertEqual(processor.image_processor.to_json_string()							,						image_processor_add_kwargs.to_json_string()							)
   self.assertIsInstance(processor.image_processor							,						UpperCAmelCase__							)


 def _lowercase	(     self		:						Union[str, Any]							)					->       Tuple:
   _a  :       Optional[Any]             =	self.get_image_processor()

   _a  :       int             =	SamProcessor(image_processor=UpperCAmelCase__							)

   _a  :       Union[str, Any]             =	self.prepare_image_inputs()

   _a  :       List[str]             =	image_processor(UpperCAmelCase__							,						return_tensors="""np"""							)
   _a  :       List[str]             =	processor(images=UpperCAmelCase__							,						return_tensors="""np"""							)

   input_feat_extract.pop("""original_sizes"""							)  # pop original_sizes as it is popped in the processor
   input_feat_extract.pop("""reshaped_input_sizes"""							)  # pop original_sizes as it is popped in the processor

   for key in input_feat_extract.keys():
     self.assertAlmostEqual(input_feat_extract[key].sum()							,						input_processor[key].sum()							,						delta=1E-2							)







 @require_torch
 def _lowercase	(     self		:						Optional[Any]							)					->       Optional[Any]:
   _a  :       Optional[Any]             =	self.get_image_processor()

   _a  :       Dict             =	SamProcessor(image_processor=UpperCAmelCase__							)
   _a  :       Tuple             =	[torch.ones((1, 3, 5, 5)							)]

   _a  :       Tuple             =	[[1764, 2646]]

   _a  :       Optional[int]             =	[[683, 1024]]
   _a  :       List[Any]             =	processor.post_process_masks(UpperCAmelCase__							,						UpperCAmelCase__							,						UpperCAmelCase__							)
   self.assertEqual(masks[0].shape							,						(1, 3, 1764, 2646)							)

   _a  :       int             =	processor.post_process_masks(
       UpperCAmelCase__							,						torch.tensor(UpperCAmelCase__							)							,						torch.tensor(UpperCAmelCase__							)							)
   self.assertEqual(masks[0].shape							,						(1, 3, 1764, 2646)							)

   # should also work with np
   _a  :       Optional[Any]             =	[np.ones((1, 3, 5, 5)							)]
   _a  :       Tuple             =	processor.post_process_masks(UpperCAmelCase__							,						np.array(UpperCAmelCase__							)							,						np.array(UpperCAmelCase__							)							)

   self.assertEqual(masks[0].shape							,						(1, 3, 1764, 2646)							)

   _a  :       List[str]             =	[[1, 0], [0, 1]]
   with self.assertRaises(UpperCAmelCase__							):
     _a  :       str             =	processor.post_process_masks(UpperCAmelCase__							,						np.array(UpperCAmelCase__							)							,						np.array(UpperCAmelCase__							)							)







@require_vision
@require_tf
class        UpperCamelCase			(	unittest.TestCase      ):


 def _lowercase	(     self		:						Any							)					->       List[str]:
   _a  :       List[str]             =	tempfile.mkdtemp()
   _a  :       Any             =	SamImageProcessor()
   _a  :       Union[str, Any]             =	SamProcessor(UpperCAmelCase__							)
   processor.save_pretrained(self.tmpdirname							)


 def _lowercase	(     self		:						List[str]							,						**UpperCAmelCase__		:						Any							)					->       List[str]:
   return AutoProcessor.from_pretrained(self.tmpdirname							,						**UpperCAmelCase__							).image_processor


 def _lowercase	(     self		:						Optional[Any]							)					->       Union[str, Any]:
   shutil.rmtree(self.tmpdirname							)


 def _lowercase	(     self		:						Dict							)					->       List[str]:
   _a  :       List[Any]             =	[np.random.randint(255							,						size=(3, 30, 400)							,						dtype=np.uinta							)]

   _a  :       List[Any]             =	[Image.fromarray(np.moveaxis(UpperCAmelCase__							,						0							,						-1							)							) for x in image_inputs]

   return image_inputs


 def _lowercase	(     self		:						Union[str, Any]							)					->       Union[str, Any]:
   _a  :       Optional[int]             =	SamProcessor(image_processor=self.get_image_processor()							)
   processor.save_pretrained(self.tmpdirname							)

   _a  :       str             =	self.get_image_processor(do_normalize=UpperCAmelCase__							,						padding_value=1.0							)

   _a  :       Union[str, Any]             =	SamProcessor.from_pretrained(self.tmpdirname							,						do_normalize=UpperCAmelCase__							,						padding_value=1.0							)

   self.assertEqual(processor.image_processor.to_json_string()							,						image_processor_add_kwargs.to_json_string()							)
   self.assertIsInstance(processor.image_processor							,						UpperCAmelCase__							)


 def _lowercase	(     self		:						List[Any]							)					->       str:
   _a  :       Union[str, Any]             =	self.get_image_processor()

   _a  :       Dict             =	SamProcessor(image_processor=UpperCAmelCase__							)

   _a  :       int             =	self.prepare_image_inputs()

   _a  :       List[str]             =	image_processor(UpperCAmelCase__							,						return_tensors="""np"""							)
   _a  :       List[str]             =	processor(images=UpperCAmelCase__							,						return_tensors="""np"""							)

   input_feat_extract.pop("""original_sizes"""							)  # pop original_sizes as it is popped in the processor
   input_feat_extract.pop("""reshaped_input_sizes"""							)  # pop reshaped_input_sizes as it is popped in the processor

   for key in input_feat_extract.keys():
     self.assertAlmostEqual(input_feat_extract[key].sum()							,						input_processor[key].sum()							,						delta=1E-2							)







 @require_tf
 def _lowercase	(     self		:						Optional[Any]							)					->       int:
   _a  :       Optional[Any]             =	self.get_image_processor()

   _a  :       Dict             =	SamProcessor(image_processor=UpperCAmelCase__							)
   _a  :       Any             =	[tf.ones((1, 3, 5, 5)							)]

   _a  :       Tuple             =	[[1764, 2646]]

   _a  :       str             =	[[683, 1024]]
   _a  :       Union[str, Any]             =	processor.post_process_masks(UpperCAmelCase__							,						UpperCAmelCase__							,						UpperCAmelCase__							,						return_tensors="""tf"""							)
   self.assertEqual(masks[0].shape							,						(1, 3, 1764, 2646)							)

   _a  :       Union[str, Any]             =	processor.post_process_masks(
       UpperCAmelCase__							,						tf.convert_to_tensor(UpperCAmelCase__							)							,						tf.convert_to_tensor(UpperCAmelCase__							)							,						return_tensors="""tf"""							,						)
   self.assertEqual(masks[0].shape							,						(1, 3, 1764, 2646)							)

   # should also work with np
   _a  :       List[Any]             =	[np.ones((1, 3, 5, 5)							)]
   _a  :       Optional[int]             =	processor.post_process_masks(
       UpperCAmelCase__							,						np.array(UpperCAmelCase__							)							,						np.array(UpperCAmelCase__							)							,						return_tensors="""tf"""							)

   self.assertEqual(masks[0].shape							,						(1, 3, 1764, 2646)							)

   _a  :       Dict             =	[[1, 0], [0, 1]]
   with self.assertRaises(tf.errors.InvalidArgumentError							):
     _a  :       List[Any]             =	processor.post_process_masks(
         UpperCAmelCase__							,						np.array(UpperCAmelCase__							)							,						np.array(UpperCAmelCase__							)							,						return_tensors="""tf"""							)







@require_vision
@require_torchvision
class        UpperCamelCase			(	unittest.TestCase      ):


 def _lowercase	(     self		:						str							)					->       Optional[Any]:
   _a  :       Optional[Any]             =	tempfile.mkdtemp()
   _a  :       Dict             =	SamImageProcessor()
   _a  :       List[str]             =	SamProcessor(UpperCAmelCase__							)
   processor.save_pretrained(self.tmpdirname							)


 def _lowercase	(     self		:						Any							,						**UpperCAmelCase__		:						Dict							)					->       int:
   return AutoProcessor.from_pretrained(self.tmpdirname							,						**UpperCAmelCase__							).image_processor


 def _lowercase	(     self		:						Tuple							)					->       List[Any]:
   shutil.rmtree(self.tmpdirname							)


 def _lowercase	(     self		:						str							)					->       int:
   _a  :       str             =	[np.random.randint(255							,						size=(3, 30, 400)							,						dtype=np.uinta							)]

   _a  :       int             =	[Image.fromarray(np.moveaxis(UpperCAmelCase__							,						0							,						-1							)							) for x in image_inputs]

   return image_inputs


 @is_pt_tf_cross_test
 def _lowercase	(     self		:						int							)					->       List[Any]:
   _a  :       Optional[Any]             =	self.get_image_processor()

   _a  :       Optional[Any]             =	SamProcessor(image_processor=UpperCAmelCase__							)
   _a  :       Tuple             =	np.random.randint(0							,						2							,						size=(1, 3, 5, 5)							).astype(np.floataa							)
   _a  :       str             =	[tf.convert_to_tensor(UpperCAmelCase__							)]
   _a  :       Optional[int]             =	[torch.tensor(UpperCAmelCase__							)]

   _a  :       Union[str, Any]             =	[[1764, 2646]]

   _a  :       List[str]             =	[[683, 1024]]
   _a  :       Optional[int]             =	processor.post_process_masks(
       UpperCAmelCase__							,						UpperCAmelCase__							,						UpperCAmelCase__							,						return_tensors="""tf"""							)
   _a  :       List[str]             =	processor.post_process_masks(
       UpperCAmelCase__							,						UpperCAmelCase__							,						UpperCAmelCase__							,						return_tensors="""pt"""							)

   self.assertTrue(np.all(tf_masks[0].numpy() == pt_masks[0].numpy()							)							)







 @is_pt_tf_cross_test
 def _lowercase	(     self		:						str							)					->       Optional[Any]:
   _a  :       List[Any]             =	self.get_image_processor()

   _a  :       Any             =	SamProcessor(image_processor=UpperCAmelCase__							)

   _a  :       Dict             =	self.prepare_image_inputs()

   _a  :       List[str]             =	image_processor(UpperCAmelCase__							,						return_tensors="""pt"""							)["""pixel_values"""].numpy()
   _a  :       str             =	processor(images=UpperCAmelCase__							,						return_tensors="""pt"""							)["""pixel_values"""].numpy()

   _a  :       Optional[Any]             =	image_processor(UpperCAmelCase__							,						return_tensors="""tf"""							)["""pixel_values"""].numpy()
   _a  :       Optional[int]             =	processor(images=UpperCAmelCase__							,						return_tensors="""tf"""							)["""pixel_values"""].numpy()

   self.assertTrue(np.allclose(UpperCAmelCase__							,						UpperCAmelCase__							)							)
   self.assertTrue(np.allclose(UpperCAmelCase__							,						UpperCAmelCase__							)							)
   self.assertTrue(np.allclose(UpperCAmelCase__							,						UpperCAmelCase__							)							)

