





"""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 YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
__A											=						logging.get_logger(__name__)
def 					__A  (_SCREAMING_SNAKE_CASE     )   ->YolosConfig:





  """simple docstring"""


  lowerCAmelCase__     :List[str]											=				YolosConfig()

  # size of the architecture
  if "yolos_ti" in yolos_name:
    lowerCAmelCase__     :Tuple											=				192
    lowerCAmelCase__     :List[str]											=				768
    lowerCAmelCase__     :Optional[int]											=				12
    lowerCAmelCase__     :int											=				3
    lowerCAmelCase__     :List[str]											=				[800, 1333]
    lowerCAmelCase__     :Optional[Any]											=				False
  elif yolos_name == "yolos_s_dWr":
    lowerCAmelCase__     :List[Any]											=				330
    lowerCAmelCase__     :str											=				14
    lowerCAmelCase__     :str											=				6
    lowerCAmelCase__     :Dict											=				1320
  elif "yolos_s" in yolos_name:
    lowerCAmelCase__     :int											=				384
    lowerCAmelCase__     :int											=				1536
    lowerCAmelCase__     :int											=				12
    lowerCAmelCase__     :List[str]											=				6
  elif "yolos_b" in yolos_name:
    lowerCAmelCase__     :Tuple											=				[800, 1344]

  lowerCAmelCase__     :List[str]											=				91
  lowerCAmelCase__     :Dict											=				'huggingface/label-files'
  lowerCAmelCase__     :Union[str, Any]											=				'coco-detection-id2label.json'
  lowerCAmelCase__     :List[str]											=				json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  repo_type='dataset'     )	,  'r'     )     )
  lowerCAmelCase__     :Optional[Any]											=				{int(_SCREAMING_SNAKE_CASE     ): v for k, v in idalabel.items()}
  lowerCAmelCase__     :int											=				idalabel
  lowerCAmelCase__     :List[Any]											=				{v: k for k, v in idalabel.items()}

  return config
def 					__A  (_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE = False     )   ->List[str]:





  """simple docstring"""


  for i in range(config.num_hidden_layers     ):
    # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
    lowerCAmelCase__     :Optional[Any]											=				state_dict.pop(F"blocks.{i}.attn.qkv.weight"     )
    lowerCAmelCase__     :Union[str, Any]											=				state_dict.pop(F"blocks.{i}.attn.qkv.bias"     )
    # next, add query, keys and values (in that order) to the state dict
    lowerCAmelCase__     :int											=				in_proj_weight[: config.hidden_size, :]
    lowerCAmelCase__     :int											=				in_proj_bias[: config.hidden_size]
    lowerCAmelCase__     :Any											=				in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    lowerCAmelCase__     :Optional[Any]											=				in_proj_bias[
        config.hidden_size : config.hidden_size * 2
    ]
    lowerCAmelCase__     :Dict											=				in_proj_weight[-config.hidden_size :, :]
    lowerCAmelCase__     :List[str]											=				in_proj_bias[-config.hidden_size :]
def 					__A  (_SCREAMING_SNAKE_CASE     )   ->str:





  """simple docstring"""


  if "backbone" in name:
    lowerCAmelCase__     :str											=				name.replace('backbone'	,  'vit'     )
  if "cls_token" in name:
    lowerCAmelCase__     :str											=				name.replace('cls_token'	,  'embeddings.cls_token'     )
  if "det_token" in name:
    lowerCAmelCase__     :Tuple											=				name.replace('det_token'	,  'embeddings.detection_tokens'     )
  if "mid_pos_embed" in name:
    lowerCAmelCase__     :Dict											=				name.replace('mid_pos_embed'	,  'encoder.mid_position_embeddings'     )
  if "pos_embed" in name:
    lowerCAmelCase__     :List[str]											=				name.replace('pos_embed'	,  'embeddings.position_embeddings'     )
  if "patch_embed.proj" in name:
    lowerCAmelCase__     :Optional[Any]											=				name.replace('patch_embed.proj'	,  'embeddings.patch_embeddings.projection'     )
  if "blocks" in name:
    lowerCAmelCase__     :List[Any]											=				name.replace('blocks'	,  'encoder.layer'     )
  if "attn.proj" in name:
    lowerCAmelCase__     :str											=				name.replace('attn.proj'	,  'attention.output.dense'     )
  if "attn" in name:
    lowerCAmelCase__     :List[Any]											=				name.replace('attn'	,  'attention.self'     )
  if "norm1" in name:
    lowerCAmelCase__     :List[Any]											=				name.replace('norm1'	,  'layernorm_before'     )
  if "norm2" in name:
    lowerCAmelCase__     :Union[str, Any]											=				name.replace('norm2'	,  'layernorm_after'     )
  if "mlp.fc1" in name:
    lowerCAmelCase__     :str											=				name.replace('mlp.fc1'	,  'intermediate.dense'     )
  if "mlp.fc2" in name:
    lowerCAmelCase__     :int											=				name.replace('mlp.fc2'	,  'output.dense'     )
  if "class_embed" in name:
    lowerCAmelCase__     :Tuple											=				name.replace('class_embed'	,  'class_labels_classifier'     )
  if "bbox_embed" in name:
    lowerCAmelCase__     :Tuple											=				name.replace('bbox_embed'	,  'bbox_predictor'     )
  if "vit.norm" in name:
    lowerCAmelCase__     :Tuple											=				name.replace('vit.norm'	,  'vit.layernorm'     )

  return name
def 					__A  (_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE     )   ->dict:





  """simple docstring"""


  for key in orig_state_dict.copy().keys():
    lowerCAmelCase__     :Any											=				orig_state_dict.pop(_SCREAMING_SNAKE_CASE     )

    if "qkv" in key:
      lowerCAmelCase__     :str											=				key.split('.'     )
      lowerCAmelCase__     :Any											=				int(key_split[2]     )
      lowerCAmelCase__     :Union[str, Any]											=				model.vit.encoder.layer[layer_num].attention.attention.all_head_size
      if "weight" in key:
        lowerCAmelCase__     :Dict											=				val[:dim, :]
        lowerCAmelCase__     :Optional[Any]											=				val[
            dim : dim * 2, :
        ]
        lowerCAmelCase__     :Union[str, Any]											=				val[-dim:, :]
      else:
        lowerCAmelCase__     :Dict											=				val[:dim]
        lowerCAmelCase__     :List[Any]											=				val[dim : dim * 2]
        lowerCAmelCase__     :Any											=				val[-dim:]
    else:
      lowerCAmelCase__     :int											=				val

  return orig_state_dict
def 					__A  ()   ->torch.Tensor:





  """simple docstring"""


  lowerCAmelCase__     :Any											=				'http://images.cocodataset.org/val2017/000000039769.jpg'
  lowerCAmelCase__     :Union[str, Any]											=				Image.open(requests.get(_SCREAMING_SNAKE_CASE	,  stream=_SCREAMING_SNAKE_CASE     ).raw     )
  return im
@torch.no_grad()
def 					__A  (_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE = False     )   ->Tuple:





  """simple docstring"""


  lowerCAmelCase__     :str											=				get_yolos_config(_SCREAMING_SNAKE_CASE     )

  # load original state_dict
  lowerCAmelCase__     :Optional[Any]											=				torch.load(_SCREAMING_SNAKE_CASE	,  map_location='cpu'     )['model']

  # load 🤗 model
  lowerCAmelCase__     :Optional[Any]											=				YolosForObjectDetection(_SCREAMING_SNAKE_CASE     )
  model.eval()
  lowerCAmelCase__     :str											=				convert_state_dict(_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE     )
  model.load_state_dict(_SCREAMING_SNAKE_CASE     )

  # Check outputs on an image, prepared by YolosImageProcessor
  lowerCAmelCase__     :Dict											=				800 if yolos_name != 'yolos_ti' else 512
  lowerCAmelCase__     :Dict											=				YolosImageProcessor(format='coco_detection'	,  size=_SCREAMING_SNAKE_CASE     )
  lowerCAmelCase__     :Optional[int]											=				image_processor(images=prepare_img()	,  return_tensors='pt'     )
  lowerCAmelCase__     :List[Any]											=				model(**_SCREAMING_SNAKE_CASE     )
  lowerCAmelCase__				,							lowerCAmelCase__     :List[str]											=				outputs.logits, outputs.pred_boxes

  lowerCAmelCase__				,							lowerCAmelCase__     :int											=				None, None
  if yolos_name == "yolos_ti":
    lowerCAmelCase__     :List[str]											=				torch.tensor(
        [[-3_9.5_0_2_2, -1_1.9_8_2_0, -1_7.6_8_8_8], [-2_9.9_5_7_4, -9.9_7_6_9, -1_7.7_6_9_1], [-4_2.3_2_8_1, -2_0.7_2_0_0, -3_0.6_2_9_4]]     )
    lowerCAmelCase__     :Union[str, Any]											=				torch.tensor(
        [[0.4_0_2_1, 0.0_8_3_6, 0.7_9_7_9], [0.0_1_8_4, 0.2_6_0_9, 0.0_3_6_4], [0.1_7_8_1, 0.2_0_0_4, 0.2_0_9_5]]     )
  elif yolos_name == "yolos_s_200_pre":
    lowerCAmelCase__     :int											=				torch.tensor(
        [[-2_4.0_2_4_8, -1_0.3_0_2_4, -1_4.8_2_9_0], [-4_2.0_3_9_2, -1_6.8_2_0_0, -2_7.4_3_3_4], [-2_7.2_7_4_3, -1_1.8_1_5_4, -1_8.7_1_4_8]]     )
    lowerCAmelCase__     :Tuple											=				torch.tensor(
        [[0.2_5_5_9, 0.5_4_5_5, 0.4_7_0_6], [0.2_9_8_9, 0.7_2_7_9, 0.1_8_7_5], [0.7_7_3_2, 0.4_0_1_7, 0.4_4_6_2]]     )
  elif yolos_name == "yolos_s_300_pre":
    lowerCAmelCase__     :List[str]											=				torch.tensor(
        [[-3_6.2_2_2_0, -1_4.4_3_8_5, -2_3.5_4_5_7], [-3_5.6_9_7_0, -1_4.7_5_8_3, -2_1.3_9_3_5], [-3_1.5_9_3_9, -1_3.6_0_4_2, -1_6.8_0_4_9]]     )
    lowerCAmelCase__     :Dict											=				torch.tensor(
        [[0.7_6_1_4, 0.2_3_1_6, 0.4_7_2_8], [0.7_1_6_8, 0.4_4_9_5, 0.3_8_5_5], [0.4_9_9_6, 0.1_4_6_6, 0.9_9_9_6]]     )
  elif yolos_name == "yolos_s_dWr":
    lowerCAmelCase__     :Dict											=				torch.tensor(
        [[-4_2.8_6_6_8, -2_4.1_0_4_9, -4_1.1_6_9_0], [-3_4.7_4_5_6, -1_4.1_2_7_4, -2_4.9_1_9_4], [-3_3.7_8_9_8, -1_2.1_9_4_6, -2_5.6_4_9_5]]     )
    lowerCAmelCase__     :Optional[int]											=				torch.tensor(
        [[0.5_5_8_7, 0.2_7_7_3, 0.0_6_0_5], [0.5_0_0_4, 0.3_0_1_4, 0.9_9_9_4], [0.4_9_9_9, 0.1_5_4_8, 0.9_9_9_4]]     )
  elif yolos_name == "yolos_base":
    lowerCAmelCase__     :str											=				torch.tensor(
        [[-4_0.6_0_6_4, -2_4.3_0_8_4, -3_2.6_4_4_7], [-5_5.1_9_9_0, -3_0.7_7_1_9, -3_5.5_8_7_7], [-5_1.4_3_1_1, -3_3.3_5_0_7, -3_5.6_4_6_2]]     )
    lowerCAmelCase__     :Any											=				torch.tensor(
        [[0.5_5_5_5, 0.2_7_9_4, 0.0_6_5_5], [0.9_0_4_9, 0.2_6_6_4, 0.1_8_9_4], [0.9_1_8_3, 0.1_9_8_4, 0.1_6_3_5]]     )
  else:
    raise ValueError(F"Unknown yolos_name: {yolos_name}"     )

  assert torch.allclose(logits[0, :3, :3]	,  _SCREAMING_SNAKE_CASE	,  atol=1e-4     )
  assert torch.allclose(pred_boxes[0, :3, :3]	,  _SCREAMING_SNAKE_CASE	,  atol=1e-4     )

  Path(_SCREAMING_SNAKE_CASE     ).mkdir(exist_ok=_SCREAMING_SNAKE_CASE     )
  print(F"Saving model {yolos_name} to {pytorch_dump_folder_path}"     )
  model.save_pretrained(_SCREAMING_SNAKE_CASE     )
  print(F"Saving image processor to {pytorch_dump_folder_path}"     )
  image_processor.save_pretrained(_SCREAMING_SNAKE_CASE     )

  if push_to_hub:
    lowerCAmelCase__     :Union[str, Any]											=				{
        'yolos_ti': 'yolos-tiny',
        'yolos_s_200_pre': 'yolos-small',
        'yolos_s_300_pre': 'yolos-small-300',
        'yolos_s_dWr': 'yolos-small-dwr',
        'yolos_base': 'yolos-base',
    }

    print('Pushing to the hub...'     )
    lowerCAmelCase__     :Union[str, Any]											=				model_mapping[yolos_name]
    image_processor.push_to_hub(_SCREAMING_SNAKE_CASE	,  organization='hustvl'     )
    model.push_to_hub(_SCREAMING_SNAKE_CASE	,  organization='hustvl'     )


if __name__ == "__main__":
  __A											=						argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--yolos_name""",
      default="""yolos_s_200_pre""",
      type=str,
      help=(
          """Name of the YOLOS model you'd like to convert. Should be one of 'yolos_ti', 'yolos_s_200_pre',"""
          """ 'yolos_s_300_pre', 'yolos_s_dWr', 'yolos_base'."""
      ),
  )
  parser.add_argument(
      """--checkpoint_path""", default=None, type=str, help="""Path to the original state dict (.pth file)."""
  )
  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 or not to push the converted model to the 🤗 hub."""
  )

  __A											=						parser.parse_args()
  convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)








import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


lowercase    				=   logging.get_logger(__name__)




class 					UpperCamelCase_ (     snake_case_							):



	'''simple docstring'''




	lowerCAmelCase						=					['''input_features''', '''is_longer''']




	def __init__(			self			,     a=64			,     a=4_80_00			,     a=4_80			,     a=10			,     a=10_24			,     a=0.0			,     a=False			,     a = 0			,     a = 1_40_00			,     a = None			,     a = "fusion"			,     a = "repeatpad"			,     **a			,     )			->       List[str]:
								super().__init__(
								    feature_size=a			,     sampling_rate=a			,     padding_value=a			,     return_attention_mask=a			,     **a			,     )
								snake_case_           =				top_db
								snake_case_           =				truncation
								snake_case_           =				padding
								snake_case_           =				fft_window_size
								snake_case_           =				(fft_window_size >> 1) + 1
								snake_case_           =				hop_length
								snake_case_           =				max_length_s
								snake_case_           =				max_length_s * sampling_rate
								snake_case_           =				sampling_rate
								snake_case_           =				frequency_min
								snake_case_           =				frequency_max
								snake_case_           =				mel_filter_bank(
								    num_frequency_bins=self.nb_frequency_bins			,     num_mel_filters=a			,     min_frequency=a			,     max_frequency=a			,     sampling_rate=a			,     norm=a			,     mel_scale='htk'			,     )
								snake_case_           =				mel_filter_bank(
								    num_frequency_bins=self.nb_frequency_bins			,     num_mel_filters=a			,     min_frequency=a			,     max_frequency=a			,     sampling_rate=a			,     norm='slaney'			,     mel_scale='slaney'			,     )




	def 				_UpperCamelCase     (			self   )			->       Dict[str, Any]:
								snake_case_           =				copy.deepcopy(self.__dict__   )
								snake_case_           =				self.__class__.__name__
								if "mel_filters" in output:
															del output["mel_filters"]
								if "mel_filters_slaney" in output:
															del output["mel_filters_slaney"]
								return output




	def 				_UpperCamelCase     (			self			,     a			,     a = None   )			->       np.ndarray:
								snake_case_           =				spectrogram(
								    a			,     window_function(self.fft_window_size			,     'hann'   )			,     frame_length=self.fft_window_size			,     hop_length=self.hop_length			,     power=2.0			,     mel_filters=a			,     log_mel='dB'			,     )
								return log_mel_spectrogram.T




	def 				_UpperCamelCase     (			self			,     a			,     a			,     a   )			->       Tuple:
								snake_case_           =				np.array_split(list(range(0			,     total_frames - chunk_frames + 1   )   )			,     3   )
								if len(ranges[1]   ) == 0:
															# if the audio is too short, we just use the first chunk
															snake_case_           =				[0]
								if len(ranges[2]   ) == 0:
															# if the audio is too short, we just use the first chunk
															snake_case_           =				[0]
								# randomly choose index for each part
								snake_case_           =				np.random.choice(ranges[0]   )
								snake_case_           =				np.random.choice(ranges[1]   )
								snake_case_           =				np.random.choice(ranges[2]   )

								snake_case_           =				mel[idx_front : idx_front + chunk_frames, :]
								snake_case_           =				mel[idx_middle : idx_middle + chunk_frames, :]
								snake_case_           =				mel[idx_back : idx_back + chunk_frames, :]

								snake_case_           =				torch.tensor(mel[None, None, :]   )
								snake_case_           =				torch.nn.functional.interpolate(
								    a			,     size=[chunk_frames, 64]			,     mode='bilinear'			,     align_corners=a   )
								snake_case_           =				mel_shrink[0][0].numpy()
								snake_case_           =				np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back]			,     axis=0   )
								return mel_fusion




	def 				_UpperCamelCase     (			self			,     a			,     a			,     a			,     a   )			->       np.array:
								if waveform.shape[0] > max_length:
															if truncation == "rand_trunc":
																						snake_case_           =				True
																						# random crop to max_length (for compatibility) -> this should be handled by self.pad
																						snake_case_           =				len(a   ) - max_length
																						snake_case_           =				np.random.randint(0			,     overflow + 1   )
																						snake_case_           =				waveform[idx : idx + max_length]
																						snake_case_           =				self._np_extract_fbank_features(a			,     self.mel_filters_slaney   )[None, :]
															elif truncation == "fusion":
																						snake_case_           =				self._np_extract_fbank_features(a			,     self.mel_filters   )
																						snake_case_           =				max_length // self.hop_length + 1  # the +1 related to how the spectrogram is computed
																						snake_case_           =				mel.shape[0]
																						if chunk_frames == total_frames:
																													# there is a corner case where the audio length is larger than max_length but smaller than max_length+hop_length.
																													# In this case, we just use the whole audio.
																													snake_case_           =				np.stack([mel, mel, mel, mel]			,     axis=0   )
																													snake_case_           =				False
																						else:
																													snake_case_           =				self._random_mel_fusion(a			,     a			,     a   )
																													snake_case_           =				True
															else:
																						raise NotImplementedError(F'''data_truncating {truncation} not implemented'''   )

								else:
															snake_case_           =				False
															# only use repeat as a new possible value for padding. you repeat the audio before applying the usual max_length padding
															if waveform.shape[0] < max_length:
																						if padding == "repeat":
																													snake_case_           =				int(max_length / len(a   )   )
																													snake_case_           =				np.stack(np.tile(a			,     n_repeat + 1   )   )[:max_length]
																						if padding == "repeatpad":
																													snake_case_           =				int(max_length / len(a   )   )
																													snake_case_           =				np.stack(np.tile(a			,     a   )   )
																						snake_case_           =				np.pad(a			,     (0, max_length - waveform.shape[0])			,     mode='constant'			,     constant_values=0   )

															if truncation == "fusion":
																						snake_case_           =				self._np_extract_fbank_features(a			,     self.mel_filters   )
																						snake_case_           =				np.stack([input_mel, input_mel, input_mel, input_mel]			,     axis=0   )
															else:
																						snake_case_           =				self._np_extract_fbank_features(a			,     self.mel_filters_slaney   )[None, :]

								return input_mel, longer






	def __call__(			self			,     a			,     a = None			,     a = None			,     a = None			,     a = None			,     a = None			,     **a			,     )			->       BatchFeature:
								snake_case_           =				truncation if truncation is not None else self.truncation
								snake_case_           =				padding if padding else self.padding

								if sampling_rate is not None:
															if sampling_rate != self.sampling_rate:
																						raise ValueError(
																						    F'''The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a'''
																						    F''' sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input'''
																						    F''' was sampled with {self.sampling_rate} and not {sampling_rate}.'''   )
								else:
															logger.warning(
															    'It is strongly recommended to pass the `sampling_rate` argument to this function. '
															    'Failing to do so can result in silent errors that might be hard to debug.'   )

								snake_case_           =				isinstance(a			,     np.ndarray   ) and len(raw_speech.shape   ) > 1
								if is_batched_numpy and len(raw_speech.shape   ) > 2:
															raise ValueError(F'''Only mono-channel audio is supported for input to {self}'''   )
								snake_case_           =				is_batched_numpy or (
								    isinstance(a			,     (list, tuple)   ) and (isinstance(raw_speech[0]			,     (np.ndarray, tuple, list)   ))
								)

								if is_batched:
															snake_case_           =				[np.asarray(a			,     dtype=np.floataa   ) for speech in raw_speech]
								elif not is_batched and not isinstance(a			,     np.ndarray   ):
															snake_case_           =				np.asarray(a			,     dtype=np.floataa   )
								elif isinstance(a			,     np.ndarray   ) and raw_speech.dtype is np.dtype(np.floataa   ):
															snake_case_           =				raw_speech.astype(np.floataa   )

								# always return batch
								if not is_batched:
															snake_case_           =				[np.asarray(a   )]

								# convert to mel spectrogram, truncate and pad if needed.
								snake_case_           =				[
								    self._get_input_mel(a			,     max_length if max_length else self.nb_max_samples			,     a			,     a   )
								    for waveform in raw_speech
								]

								snake_case_           =				[]
								snake_case_           =				[]
								for mel, longer in padded_inputs:
															input_mel.append(a   )
															is_longer.append(a   )

								if truncation == "fusion" and sum(a   ) == 0:
															# if no audio is longer than 10s, then randomly select one audio to be longer
															snake_case_           =				np.random.randint(0			,     len(a   )   )
															snake_case_           =				True

								if isinstance(input_mel[0]			,     a   ):
															snake_case_           =				[np.asarray(a			,     dtype=np.floataa   ) for feature in input_mel]

								# is_longer is a list of bool
								snake_case_           =				[[longer] for longer in is_longer]

								snake_case_           =				{'input_features': input_mel, 'is_longer': is_longer}
								snake_case_           =				BatchFeature(a   )

								if return_tensors is not None:
															snake_case_           =				input_features.convert_to_tensors(a   )

								return input_features





