





import argparse
from collections import OrderedDict
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision.transforms import functional as F

from transformers import DetrImageProcessor, TableTransformerConfig, TableTransformerForObjectDetection
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase			       =      logging.get_logger(__name__)

# here we list all keys to be renamed (original name on the left, our name on the right)
_lowerCamelCase			       =      []
for i in range(6):
       # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
       rename_keys.append(
           (F'''transformer.encoder.layers.{i}.self_attn.out_proj.weight''', F'''encoder.layers.{i}.self_attn.out_proj.weight''')
       )
       rename_keys.append(
           (F'''transformer.encoder.layers.{i}.self_attn.out_proj.bias''', F'''encoder.layers.{i}.self_attn.out_proj.bias''')
       )
       rename_keys.append((F'''transformer.encoder.layers.{i}.linear1.weight''', F'''encoder.layers.{i}.fc1.weight'''))
       rename_keys.append((F'''transformer.encoder.layers.{i}.linear1.bias''', F'''encoder.layers.{i}.fc1.bias'''))
       rename_keys.append((F'''transformer.encoder.layers.{i}.linear2.weight''', F'''encoder.layers.{i}.fc2.weight'''))
       rename_keys.append((F'''transformer.encoder.layers.{i}.linear2.bias''', F'''encoder.layers.{i}.fc2.bias'''))
       rename_keys.append(
           (F'''transformer.encoder.layers.{i}.norm1.weight''', F'''encoder.layers.{i}.self_attn_layer_norm.weight''')
       )
       rename_keys.append((F'''transformer.encoder.layers.{i}.norm1.bias''', F'''encoder.layers.{i}.self_attn_layer_norm.bias'''))
       rename_keys.append((F'''transformer.encoder.layers.{i}.norm2.weight''', F'''encoder.layers.{i}.final_layer_norm.weight'''))
       rename_keys.append((F'''transformer.encoder.layers.{i}.norm2.bias''', F'''encoder.layers.{i}.final_layer_norm.bias'''))
       # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
       rename_keys.append(
           (F'''transformer.decoder.layers.{i}.self_attn.out_proj.weight''', F'''decoder.layers.{i}.self_attn.out_proj.weight''')
       )
       rename_keys.append(
           (F'''transformer.decoder.layers.{i}.self_attn.out_proj.bias''', F'''decoder.layers.{i}.self_attn.out_proj.bias''')
       )
       rename_keys.append(
           (
               F'''transformer.decoder.layers.{i}.multihead_attn.out_proj.weight''',
               F'''decoder.layers.{i}.encoder_attn.out_proj.weight''',
           )
       )
       rename_keys.append(
           (
               F'''transformer.decoder.layers.{i}.multihead_attn.out_proj.bias''',
               F'''decoder.layers.{i}.encoder_attn.out_proj.bias''',
           )
       )
       rename_keys.append((F'''transformer.decoder.layers.{i}.linear1.weight''', F'''decoder.layers.{i}.fc1.weight'''))
       rename_keys.append((F'''transformer.decoder.layers.{i}.linear1.bias''', F'''decoder.layers.{i}.fc1.bias'''))
       rename_keys.append((F'''transformer.decoder.layers.{i}.linear2.weight''', F'''decoder.layers.{i}.fc2.weight'''))
       rename_keys.append((F'''transformer.decoder.layers.{i}.linear2.bias''', F'''decoder.layers.{i}.fc2.bias'''))
       rename_keys.append(
           (F'''transformer.decoder.layers.{i}.norm1.weight''', F'''decoder.layers.{i}.self_attn_layer_norm.weight''')
       )
       rename_keys.append((F'''transformer.decoder.layers.{i}.norm1.bias''', F'''decoder.layers.{i}.self_attn_layer_norm.bias'''))
       rename_keys.append(
           (F'''transformer.decoder.layers.{i}.norm2.weight''', F'''decoder.layers.{i}.encoder_attn_layer_norm.weight''')
       )
       rename_keys.append(
           (F'''transformer.decoder.layers.{i}.norm2.bias''', F'''decoder.layers.{i}.encoder_attn_layer_norm.bias''')
       )
       rename_keys.append((F'''transformer.decoder.layers.{i}.norm3.weight''', F'''decoder.layers.{i}.final_layer_norm.weight'''))
       rename_keys.append((F'''transformer.decoder.layers.{i}.norm3.bias''', F'''decoder.layers.{i}.final_layer_norm.bias'''))

# convolutional projection + query embeddings + layernorm of encoder + layernorm of decoder + class and bounding box heads
rename_keys.extend(
    [
        ('input_proj.weight', 'input_projection.weight'),
        ('input_proj.bias', 'input_projection.bias'),
        ('query_embed.weight', 'query_position_embeddings.weight'),
        ('transformer.encoder.norm.weight', 'encoder.layernorm.weight'),
        ('transformer.encoder.norm.bias', 'encoder.layernorm.bias'),
        ('transformer.decoder.norm.weight', 'decoder.layernorm.weight'),
        ('transformer.decoder.norm.bias', 'decoder.layernorm.bias'),
        ('class_embed.weight', 'class_labels_classifier.weight'),
        ('class_embed.bias', 'class_labels_classifier.bias'),
        ('bbox_embed.layers.0.weight', 'bbox_predictor.layers.0.weight'),
        ('bbox_embed.layers.0.bias', 'bbox_predictor.layers.0.bias'),
        ('bbox_embed.layers.1.weight', 'bbox_predictor.layers.1.weight'),
        ('bbox_embed.layers.1.bias', 'bbox_predictor.layers.1.bias'),
        ('bbox_embed.layers.2.weight', 'bbox_predictor.layers.2.weight'),
        ('bbox_embed.layers.2.bias', 'bbox_predictor.layers.2.bias'),
    ]
)




def    __UpperCAmelCase(      lowercase_ ,    lowercase_ ,    lowercase_      ):
       _lowerCamelCase							:     str	=		state_dict.pop(lowercase_      )
       _lowerCamelCase							:     Any	=		val




def    __UpperCAmelCase(      lowercase_      ):
       _lowerCamelCase							:     int	=		OrderedDict()
       for key, value in state_dict.items():
              if "backbone.0.body" in key:
                     _lowerCamelCase							:     Any	=		key.replace('''backbone.0.body''' ,    '''backbone.conv_encoder.model'''      )
                     _lowerCamelCase							:     Union[str, Any]	=		value
              else:
                     _lowerCamelCase							:     Optional[int]	=		value

       return new_state_dict




def    __UpperCAmelCase(      lowercase_      ):
       _lowerCamelCase							:     Any	=		''''''

       # first: transformer encoder
       for i in range(6      ):
              # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
              _lowerCamelCase							:     Union[str, Any]	=		state_dict.pop(F"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"""      )
              _lowerCamelCase							:     List[str]	=		state_dict.pop(F"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias"""      )
              # next, add query, keys and values (in that order) to the state dict
              _lowerCamelCase							:     Union[str, Any]	=		in_proj_weight[:2_56, :]
              _lowerCamelCase							:     str	=		in_proj_bias[:2_56]
              _lowerCamelCase							:     Any	=		in_proj_weight[2_56:5_12, :]
              _lowerCamelCase							:     Optional[Any]	=		in_proj_bias[2_56:5_12]
              _lowerCamelCase							:     Dict	=		in_proj_weight[-2_56:, :]
              _lowerCamelCase							:     str	=		in_proj_bias[-2_56:]
       # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
       for i in range(6      ):
              # read in weights + bias of input projection layer of self-attention
              _lowerCamelCase							:     Tuple	=		state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight"""      )
              _lowerCamelCase							:     Tuple	=		state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias"""      )
              # next, add query, keys and values (in that order) to the state dict
              _lowerCamelCase							:     Optional[int]	=		in_proj_weight[:2_56, :]
              _lowerCamelCase							:     Optional[Any]	=		in_proj_bias[:2_56]
              _lowerCamelCase							:     List[str]	=		in_proj_weight[2_56:5_12, :]
              _lowerCamelCase							:     List[str]	=		in_proj_bias[2_56:5_12]
              _lowerCamelCase							:     int	=		in_proj_weight[-2_56:, :]
              _lowerCamelCase							:     Any	=		in_proj_bias[-2_56:]
              # read in weights + bias of input projection layer of cross-attention
              _lowerCamelCase							:     Any	=		state_dict.pop(
                  F"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"""      )
              _lowerCamelCase							:     int	=		state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias"""      )
              # next, add query, keys and values (in that order) of cross-attention to the state dict
              _lowerCamelCase							:     List[Any]	=		in_proj_weight_cross_attn[:2_56, :]
              _lowerCamelCase							:     Any	=		in_proj_bias_cross_attn[:2_56]
              _lowerCamelCase							:     Union[str, Any]	=		in_proj_weight_cross_attn[2_56:5_12, :]
              _lowerCamelCase							:     Tuple	=		in_proj_bias_cross_attn[2_56:5_12]
              _lowerCamelCase							:     Any	=		in_proj_weight_cross_attn[-2_56:, :]
              _lowerCamelCase							:     Optional[int]	=		in_proj_bias_cross_attn[-2_56:]




def    __UpperCAmelCase(      lowercase_ ,    lowercase_      ):
       _lowerCamelCase,  _lowerCamelCase							:     Any	=		image.size
       _lowerCamelCase							:     Dict	=		max(lowercase_ ,    lowercase_      )
       _lowerCamelCase							:     Dict	=		8_00 if '''detection''' in checkpoint_url else 10_00
       _lowerCamelCase							:     str	=		target_max_size / current_max_size
       _lowerCamelCase							:     Any	=		image.resize((int(round(scale * width      )      ), int(round(scale * height      )      ))      )

       return resized_image




def    __UpperCAmelCase(      lowercase_      ):
       _lowerCamelCase							:     Any	=		F.to_tensor(lowercase_      )
       _lowerCamelCase							:     Dict	=		F.normalize(lowercase_ ,    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]      )
       return image




@torch.no_grad()
def    __UpperCAmelCase(      lowercase_ ,    lowercase_ ,    lowercase_      ):
       logger.info('''Converting model...'''      )

       # load original state dict
       _lowerCamelCase							:     Dict	=		torch.hub.load_state_dict_from_url(lowercase_ ,    map_location='''cpu'''      )
       # rename keys
       for src, dest in rename_keys:
              rename_key(lowercase_ ,    lowercase_ ,    lowercase_      )
       _lowerCamelCase							:     Optional[int]	=		rename_backbone_keys(lowercase_      )
       # query, key and value matrices need special treatment
       read_in_q_k_v(lowercase_      )
       # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
       _lowerCamelCase							:     List[Any]	=		'''model.'''
       for key in state_dict.copy().keys():
              if not key.startswith('''class_labels_classifier'''      ) and not key.startswith('''bbox_predictor'''      ):
                     _lowerCamelCase							:     Tuple	=		state_dict.pop(lowercase_      )
                     _lowerCamelCase							:     Dict	=		val
    # create HuggingFace model and load state dict
       _lowerCamelCase							:     Optional[Any]	=		TableTransformerConfig(
           backbone='''resnet18''' ,    mask_loss_coefficient=1 ,    dice_loss_coefficient=1 ,    ce_loss_coefficient=1 ,    bbox_loss_coefficient=5 ,    giou_loss_coefficient=2 ,    eos_coefficient=0.4 ,    class_cost=1 ,    bbox_cost=5 ,    giou_cost=2 ,    )

       if "detection" in checkpoint_url:
              _lowerCamelCase							:     Any	=		15
              _lowerCamelCase							:     Optional[int]	=		2
              _lowerCamelCase							:     Any	=		{0: '''table''', 1: '''table rotated'''}
              _lowerCamelCase							:     List[str]	=		idalabel
              _lowerCamelCase							:     List[Any]	=		{v: k for k, v in idalabel.items()}
       else:
              _lowerCamelCase							:     Union[str, Any]	=		1_25
              _lowerCamelCase							:     Any	=		6
              _lowerCamelCase							:     List[str]	=		{
                  0: '''table''',
                  1: '''table column''',
                  2: '''table row''',
                  3: '''table column header''',
                  4: '''table projected row header''',
                  5: '''table spanning cell''',
              }
              _lowerCamelCase							:     Optional[int]	=		idalabel
              _lowerCamelCase							:     Tuple	=		{v: k for k, v in idalabel.items()}

       _lowerCamelCase							:     Union[str, Any]	=		DetrImageProcessor(
           format='''coco_detection''' ,    max_size=8_00 if '''detection''' in checkpoint_url else 10_00      )
       _lowerCamelCase							:     Optional[Any]	=		TableTransformerForObjectDetection(lowercase_      )
       model.load_state_dict(lowercase_      )
       model.eval()

       # verify our conversion
       _lowerCamelCase							:     Union[str, Any]	=		'''example_pdf.png''' if '''detection''' in checkpoint_url else '''example_table.png'''
       _lowerCamelCase							:     Optional[Any]	=		hf_hub_download(repo_id='''nielsr/example-pdf''' ,    repo_type='''dataset''' ,    filename=lowercase_      )
       _lowerCamelCase							:     Union[str, Any]	=		Image.open(lowercase_      ).convert('''RGB'''      )
       _lowerCamelCase							:     Any	=		normalize(resize(lowercase_ ,    lowercase_      )      ).unsqueeze(0      )

       _lowerCamelCase							:     Tuple	=		model(lowercase_      )

       if "detection" in checkpoint_url:
              _lowerCamelCase							:     Tuple	=		(1, 15, 3)
              _lowerCamelCase							:     List[Any]	=		torch.tensor(
                  [[-6.7_8_9_7, -1_6.9_9_8_5, 6.7_9_3_7], [-8.0_1_8_6, -2_2.2_1_9_2, 6.9_6_7_7], [-7.3_1_1_7, -2_1.0_7_0_8, 7.4_0_5_5]]      )
              _lowerCamelCase							:     int	=		torch.tensor([[0.4_8_6_7, 0.1_7_6_7, 0.6_7_3_2], [0.6_7_1_8, 0.4_4_7_9, 0.3_8_3_0], [0.4_7_1_6, 0.1_7_6_0, 0.6_3_6_4]]      )

       else:
              _lowerCamelCase							:     List[str]	=		(1, 1_25, 7)
              _lowerCamelCase							:     Optional[Any]	=		torch.tensor(
                  [[-1_8.1_4_3_0, -8.3_2_1_4, 4.8_2_7_4], [-1_8.4_6_8_5, -7.1_3_6_1, -4.2_6_6_7], [-2_6.3_6_9_3, -9.3_4_2_9, -4.9_9_6_2]]      )
              _lowerCamelCase							:     Tuple	=		torch.tensor([[0.4_9_8_3, 0.5_5_9_5, 0.9_4_4_0], [0.4_9_1_6, 0.6_3_1_5, 0.5_9_5_4], [0.6_1_0_8, 0.8_6_3_7, 0.1_1_3_5]]      )

       assert outputs.logits.shape == expected_shape
       assert torch.allclose(outputs.logits[0, :3, :3] ,    lowercase_ ,    atol=1e-4      )
       assert torch.allclose(outputs.pred_boxes[0, :3, :3] ,    lowercase_ ,    atol=1e-4      )
       print('''Looks ok!'''      )

       if pytorch_dump_folder_path is not None:
              # Save model and image processor
              logger.info(F"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""      )
              Path(lowercase_      ).mkdir(exist_ok=lowercase_      )
              model.save_pretrained(lowercase_      )
              image_processor.save_pretrained(lowercase_      )

       if push_to_hub:
              # Push model to HF hub
              logger.info('''Pushing model to the hub...'''      )
              _lowerCamelCase							:     List[Any]	=		(
                  '''microsoft/table-transformer-detection'''
                  if '''detection''' in checkpoint_url
                  else '''microsoft/table-transformer-structure-recognition'''
              )
              model.push_to_hub(lowercase_      )
              image_processor.push_to_hub(lowercase_      )


if __name__ == "__main__":
       _lowerCamelCase			       =      argparse.ArgumentParser()

       parser.add_argument(
           '--checkpoint_url',
           default='https://pubtables1m.blob.core.windows.net/model/pubtables1m_detection_detr_r18.pth',
           type=str,
           choices=[
               'https://pubtables1m.blob.core.windows.net/model/pubtables1m_detection_detr_r18.pth',
               'https://pubtables1m.blob.core.windows.net/model/pubtables1m_structure_detr_r18.pth',
           ],
           help='URL of the Table Transformer checkpoint you\'d like to convert.',
       )
       parser.add_argument(
           '--pytorch_dump_folder_path', default=None, type=str, help='Path to the folder to output PyTorch model.'
       )
       parser.add_argument(
           '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
       )
       _lowerCamelCase			       =      parser.parse_args()
       convert_table_transformer_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)








import pytest

import datasets


# Import fixture modules as plugins
_lowerCamelCase			       =      ['tests.fixtures.files', 'tests.fixtures.hub', 'tests.fixtures.fsspec']




def    __UpperCAmelCase(      lowercase_ ,    lowercase_      ):
       # Mark tests as "unit" by default if not marked as "integration" (or already marked as "unit")
       for item in items:
              if any(marker in item.keywords for marker in ['''integration''', '''unit''']      ):
                     continue
              item.add_marker(pytest.mark.unit      )




def    __UpperCAmelCase(      lowercase_      ):
       config.addinivalue_line('''markers''' ,    '''torchaudio_latest: mark test to run with torchaudio>=0.12'''      )




@pytest.fixture(autouse=lowercase_      )
def    __UpperCAmelCase(      lowercase_ ,    lowercase_      ):
       # test_hf_cache_home = tmp_path_factory.mktemp("cache")  # TODO: why a cache dir per test function does not work?
       _lowerCamelCase							:     Optional[Any]	=		tmp_path_factory.getbasetemp() / '''cache'''
       _lowerCamelCase							:     Optional[Any]	=		test_hf_cache_home / '''datasets'''
       _lowerCamelCase							:     Union[str, Any]	=		test_hf_cache_home / '''metrics'''
       _lowerCamelCase							:     Dict	=		test_hf_cache_home / '''modules'''
       monkeypatch.setattr('''datasets.config.HF_DATASETS_CACHE''' ,    str(lowercase_      )      )
       monkeypatch.setattr('''datasets.config.HF_METRICS_CACHE''' ,    str(lowercase_      )      )
       monkeypatch.setattr('''datasets.config.HF_MODULES_CACHE''' ,    str(lowercase_      )      )
       _lowerCamelCase							:     str	=		test_hf_datasets_cache / '''downloads'''
       monkeypatch.setattr('''datasets.config.DOWNLOADED_DATASETS_PATH''' ,    str(lowercase_      )      )
       _lowerCamelCase							:     Optional[int]	=		test_hf_datasets_cache / '''downloads''' / '''extracted'''
       monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_PATH''' ,    str(lowercase_      )      )




@pytest.fixture(autouse=lowercase_ ,    scope='''session'''      )
def    __UpperCAmelCase(      ):
       datasets.disable_progress_bar()




@pytest.fixture(autouse=lowercase_      )
def    __UpperCAmelCase(      lowercase_      ):
       # don't take tests into account when counting downloads
       monkeypatch.setattr('''datasets.config.HF_UPDATE_DOWNLOAD_COUNTS''' ,    lowercase_      )




@pytest.fixture
def    __UpperCAmelCase(      lowercase_      ):
       # Required to suppress RemovedIn20Warning when feature(s) are not compatible with SQLAlchemy 2.0
       # To be removed once SQLAlchemy 2.0 supported
       monkeypatch.setattr('''sqlalchemy.util.deprecations.SILENCE_UBER_WARNING''' ,    lowercase_      )

