




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     _SCREAMING_SNAKE_CASE    (							lowercase			: Union[str, Any]  ):


      '''simple docstring'''



      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_	  =					1_28
      elif "large" in model_name:
            lowerCamelCase_	  =					1_92
      elif "xlarge" in model_name:
            lowerCamelCase_	  =					2_56
      elif "huge" in model_name:
            lowerCamelCase_	  =					3_52

      # 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(lowercase			,							lowercase			,							repo_type='dataset'  )			,							'r'  )  )
      lowerCamelCase_	  =					{int(lowercase  ): v for k, v in idalabel.items()}
      lowerCamelCase_	  =					{v: k for k, v in idalabel.items()}

      lowerCamelCase_	  =					FocalNetConfig(
          embed_dim=lowercase			,							depths=lowercase			,							focal_levels=lowercase			,							focal_windows=lowercase			,							use_conv_embed=lowercase			,							idalabel=lowercase			,							labelaid=lowercase			,							use_post_layernorm=lowercase			,							use_layerscale=lowercase			,							)

      return config

def     _SCREAMING_SNAKE_CASE    (							lowercase			: str  ):


      '''simple docstring'''



      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     _SCREAMING_SNAKE_CASE    (							lowercase			: Optional[int]			,							lowercase			: Tuple			,							lowercase			: List[str]=False  ):


      '''simple docstring'''



      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: '			,							lowercase  )
      lowerCamelCase_	  =					torch.hub.load_state_dict_from_url(lowercase			,							map_location='cpu'  )['model']

      # rename keys
      for key in state_dict.copy().keys():
            lowerCamelCase_	  =					state_dict.pop(lowercase  )
            lowerCamelCase_	  =					val

      lowerCamelCase_	  =					get_focalnet_config(lowercase  )
      lowerCamelCase_	  =					FocalNetForImageClassification(lowercase  )
      model.eval()

      # load state dict
      model.load_state_dict(lowercase  )

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

      lowerCamelCase_	  =					BitImageProcessor(
          do_resize=lowercase			,							size={'shortest_edge': 2_56}			,							resample=PILImageResampling.BILINEAR			,							do_center_crop=lowercase			,							crop_size=2_24			,							do_normalize=lowercase			,							image_mean=lowercase			,							image_std=lowercase			,							)
      lowerCamelCase_	  =					Image.open(requests.get(lowercase			,							stream=lowercase  ).raw  )
      lowerCamelCase_	  =					processor(images=lowercase			,							return_tensors='pt'  )

      lowerCamelCase_	  =					transforms.Compose(
          [
              transforms.Resize(2_56  ),
              transforms.CenterCrop(2_24  ),
              transforms.ToTensor(),
              transforms.Normalize(mean=[0.485, 0.456, 0.406]			,							std=[0.229, 0.224, 0.225]  ),
          ]  )

      lowerCamelCase_	  =					image_transforms(lowercase  ).unsqueeze(0  )

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

      lowerCamelCase_	  =					model(**lowercase  )

      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.2166, -0.4368, 0.2191]  )
      elif model_name == "focalnet-tiny-lrf":
            lowerCamelCase_	  =					torch.tensor([1.1669, 0.0125, -0.1695]  )
      elif model_name == "focalnet-small":
            lowerCamelCase_	  =					torch.tensor([0.4917, -0.0430, 0.1341]  )
      elif model_name == "focalnet-small-lrf":
            lowerCamelCase_	  =					torch.tensor([-0.2588, -0.5342, -0.2331]  )
      elif model_name == "focalnet-base":
            lowerCamelCase_	  =					torch.tensor([-0.1655, -0.4090, -0.1730]  )
      elif model_name == "focalnet-base-lrf":
            lowerCamelCase_	  =					torch.tensor([0.5306, -0.0483, -0.3928]  )
      assert torch.allclose(outputs.logits[0, :3]			,							lowercase			,							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(lowercase  )
            processor.save_pretrained(lowercase  )

      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__":
 lowerCamelCase       :	int									=     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.",
 )

 lowerCamelCase       :	Optional[Any]									=     parser.parse_args()
 convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









from typing import Dict, Iterable, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


lowerCamelCase       :	List[Any]									=     logging.get_logger(__name__)







class 						A(		UpperCamelCase     ):



 '''simple docstring'''


 UpperCamelCase		     =		['''pixel_values''']

 def __init__(       self      :						List[Any]		,							A_      :						bool = True		,							A_      :						Dict[str, int] = None		,							A_      :						PILImageResampling = PILImageResampling.BICUBIC		,							A_      :						bool = True		,							A_      :						Dict[str, int] = None		,							A_      :						bool = True		,							A_      :						Union[int, float] = 1 / 255		,							A_      :						bool = True		,							A_      :						Optional[Union[float, Iterable[float]]] = IMAGENET_DEFAULT_MEAN		,							A_      :						Optional[Union[float, Iterable[float]]] = IMAGENET_DEFAULT_STD		,							**A_      :						Tuple		,							)     ->			None:

       """simple docstring"""






       super().__init__(**A_						)
       lowerCamelCase_	  =					size if size is not None else {'shortest_edge': 224}
       lowerCamelCase_	  =					get_size_dict(A_		,							default_to_square=A_						)
       lowerCamelCase_	  =					crop_size if crop_size is not None else {'height': 224, 'width': 224}
       lowerCamelCase_	  =					get_size_dict(A_		,							param_name='crop_size'						)

       lowerCamelCase_	  =					do_resize
       lowerCamelCase_	  =					size
       lowerCamelCase_	  =					resample
       lowerCamelCase_	  =					do_center_crop
       lowerCamelCase_	  =					crop_size
       lowerCamelCase_	  =					do_rescale
       lowerCamelCase_	  =					rescale_factor
       lowerCamelCase_	  =					do_normalize
       lowerCamelCase_	  =					image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
       lowerCamelCase_	  =					image_std if image_std is not None else IMAGENET_DEFAULT_STD

 def        a__				(       self      :						Optional[Any]		,							A_      :						np.ndarray		,							A_      :						Dict[str, int]		,							A_      :						PILImageResampling = PILImageResampling.BICUBIC		,							A_      :						Optional[Union[str, ChannelDimension]] = None		,							**A_      :						Tuple		,							)     ->			np.ndarray:

       """simple docstring"""






       lowerCamelCase_	  =					get_size_dict(A_		,							default_to_square=A_						)
       # size_dict is a dict with either keys "height" and "width" or "shortest_edge"
       if "shortest_edge" in size:
             lowerCamelCase_	  =					int((256 / 224) * size['shortest_edge']						)
             lowerCamelCase_	  =					get_resize_output_image_size(A_		,							size=A_		,							default_to_square=A_						)
             lowerCamelCase_	  =					{'height': output_size[0], 'width': output_size[1]}
       if "height" not in size_dict or "width" not in size_dict:
             raise ValueError(
                 f"""Size dict must have keys 'height' and 'width' or 'shortest_edge'. Got {size_dict.keys()}"""						)
       return resize(
           A_		,							size=(size_dict['height'], size_dict['width'])		,							resample=A_		,							data_format=A_		,							**A_						)

 def        a__				(       self      :						Any		,							A_      :						np.ndarray		,							A_      :						Dict[str, int]		,							A_      :						Optional[Union[str, ChannelDimension]] = None		,							**A_      :						Any		,							)     ->			np.ndarray:

       """simple docstring"""






       lowerCamelCase_	  =					get_size_dict(A_						)
       if "height" not in size or "width" not in size:
             raise ValueError(f"""Size dict must have keys 'height' and 'width'. Got {size.keys()}"""						)
       return center_crop(A_		,							size=(size['height'], size['width'])		,							data_format=A_		,							**A_						)

 def        a__				(       self      :						Optional[Any]		,							A_      :						np.ndarray		,							A_      :						Union[int, float]		,							A_      :						Optional[Union[str, ChannelDimension]] = None		,							**A_      :						Optional[int]		,							)     ->			np.ndarray:

       """simple docstring"""






       return rescale(A_		,							scale=A_		,							data_format=A_		,							**A_						)

 def        a__				(       self      :						List[str]		,							A_      :						np.ndarray		,							A_      :						Union[float, List[float]]		,							A_      :						Union[float, List[float]]		,							A_      :						Optional[Union[str, ChannelDimension]] = None		,							**A_      :						str		,							)     ->			np.ndarray:

       """simple docstring"""






       return normalize(A_		,							mean=A_		,							std=A_		,							data_format=A_		,							**A_						)





 def        a__				(       self      :						Optional[int]		,							A_      :						ImageInput		,							A_      :						Optional[bool] = None		,							A_      :						Optional[Dict[str, int]] = None		,							A_      :						PILImageResampling = None		,							A_      :						Optional[bool] = None		,							A_      :						Optional[Dict[str, int]] = None		,							A_      :						Optional[bool] = None		,							A_      :						Optional[float] = None		,							A_      :						Optional[bool] = None		,							A_      :						Optional[Union[float, Iterable[float]]] = None		,							A_      :						Optional[Union[float, Iterable[float]]] = None		,							A_      :						Optional[TensorType] = None		,							A_      :						ChannelDimension = ChannelDimension.FIRST		,							**A_      :						List[Any]		,							)     ->			BatchFeature:

       """simple docstring"""






       lowerCamelCase_	  =					do_resize if do_resize is not None else self.do_resize
       lowerCamelCase_	  =					resample if resample is not None else self.resample
       lowerCamelCase_	  =					do_center_crop if do_center_crop is not None else self.do_center_crop
       lowerCamelCase_	  =					do_rescale if do_rescale is not None else self.do_rescale
       lowerCamelCase_	  =					rescale_factor if rescale_factor is not None else self.rescale_factor
       lowerCamelCase_	  =					do_normalize if do_normalize is not None else self.do_normalize
       lowerCamelCase_	  =					image_mean if image_mean is not None else self.image_mean
       lowerCamelCase_	  =					image_std if image_std is not None else self.image_std

       lowerCamelCase_	  =					size if size is not None else self.size
       lowerCamelCase_	  =					get_size_dict(A_		,							default_to_square=A_						)
       lowerCamelCase_	  =					crop_size if crop_size is not None else self.crop_size
       lowerCamelCase_	  =					get_size_dict(A_		,							param_name='crop_size'						)

       lowerCamelCase_	  =					make_list_of_images(A_						)

       if not valid_images(A_						):
             raise ValueError(
                 'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '
                 'torch.Tensor, tf.Tensor or jax.ndarray.'						)

       if do_resize and size is None:
             raise ValueError('Size must be specified if do_resize is True.'						)

       if do_center_crop and crop_size is None:
             raise ValueError('Crop size must be specified if do_center_crop is True.'						)

       if do_rescale and rescale_factor is None:
             raise ValueError('Rescale factor must be specified if do_rescale is True.'						)

       if do_normalize and (image_mean is None or image_std is None):
             raise ValueError('Image mean and std must be specified if do_normalize is True.'						)

       # All transformations expect numpy arrays.
       lowerCamelCase_	  =					[to_numpy_array(A_						) for image in images]

       if do_resize:
             lowerCamelCase_	  =					[self.resize(A_		,							A_		,							A_						) for image in images]

       if do_center_crop:
             lowerCamelCase_	  =					[self.center_crop(A_		,							A_						) for image in images]

       if do_rescale:
             lowerCamelCase_	  =					[self.rescale(A_		,							A_						) for image in images]

       if do_normalize:
             lowerCamelCase_	  =					[self.normalize(A_		,							A_		,							A_						) for image in images]

       lowerCamelCase_	  =					[to_channel_dimension_format(A_		,							A_						) for image in images]

       lowerCamelCase_	  =					{'pixel_values': images}
       return BatchFeature(data=A_		,							tensor_type=A_						)



