codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert dpt checkpoints from the original repository url https github comislorgdpt import argparse import json from pathlib import path import requests import torch from huggingfacehub import cacheddownload hfhuburl from pil import image from transformers import dptconfig dptfordepthestimation dptforsemanticsegmentation dptimageprocessor from transformers utils import logging logging setverbosityinfo logger logging getloggername def getdptconfigcheckpointurl config dptconfig if large in checkpointurl config hiddensize 1024 config intermediatesize 4096 config numhiddenlayers 24 config numattentionheads 16 config backboneoutindices 5 11 17 23 config neckhiddensizes 256 512 1024 1024 expectedshape 1 384 384 if ade in checkpointurl config usebatchnorminfusionresidual true config numlabels 150 repoid huggingfacelabelfiles filename ade20kid2label json id2label json loadopencacheddownloadhfhuburlrepoid filename repotypedataset r id2label intk v for k v in id2label items config id2label id2label config label2id v k for k v in id2label items expectedshape 1 150 480 480 return config expectedshape def removeignorekeysstatedict ignorekeys pretrained model head weight pretrained model head bias for k in ignorekeys statedict popk none def renamekeyname if pretrained model in name and clstoken not in name and posembed not in name and patchembed not in name name name replacepretrained model dpt encoder if pretrained model in name name name replacepretrained model dpt embeddings if patchembed in name name name replacepatchembed patchembeddings if posembed in name name name replaceposembed positionembeddings if attn proj in name name name replaceattn proj attention output dense if proj in name and project not in name name name replaceproj projection if blocks in name name name replaceblocks layer if mlp fc1 in name name name replacemlp fc1 intermediate dense if mlp fc2 in name name name replacemlp fc2 output dense if norm1 in name name name replacenorm1 layernormbefore if norm2 in name name name replacenorm2 layernormafter if scratch outputconv in name name name replacescratch outputconv head if scratch in name name name replacescratch neck if layer1rn in name name name replacelayer1rn convs 0 if layer2rn in name name name replacelayer2rn convs 1 if layer3rn in name name name replacelayer3rn convs 2 if layer4rn in name name name replacelayer4rn convs 3 if refinenet in name layeridx intnamelenneck refinenet lenneck refinenet 1 tricky here we need to map 4 to 0 3 to 1 2 to 2 and 1 to 3 name name replacefrefinenetlayeridx ffusionstage layers abslayeridx4 if outconv in name name name replaceoutconv projection if resconfunit1 in name name name replaceresconfunit1 residuallayer1 if resconfunit2 in name name name replaceresconfunit2 residuallayer2 if conv1 in name name name replaceconv1 convolution1 if conv2 in name name name replaceconv2 convolution2 readout blocks if pretrained actpostprocess1 0 project 0 in name name name replacepretrained actpostprocess1 0 project 0 neck reassemblestage readoutprojects 0 0 if pretrained actpostprocess2 0 project 0 in name name name replacepretrained actpostprocess2 0 project 0 neck reassemblestage readoutprojects 1 0 if pretrained actpostprocess3 0 project 0 in name name name replacepretrained actpostprocess3 0 project 0 neck reassemblestage readoutprojects 2 0 if pretrained actpostprocess4 0 project 0 in name name name replacepretrained actpostprocess4 0 project 0 neck reassemblestage readoutprojects 3 0 resize blocks if pretrained actpostprocess1 3 in name name name replacepretrained actpostprocess1 3 neck reassemblestage layers 0 projection if pretrained actpostprocess1 4 in name name name replacepretrained actpostprocess1 4 neck reassemblestage layers 0 resize if pretrained actpostprocess2 3 in name name name replacepretrained actpostprocess2 3 neck reassemblestage layers 1 projection if pretrained actpostprocess2 4 in name name name replacepretrained actpostprocess2 4 neck reassemblestage layers 1 resize if pretrained actpostprocess3 3 in name name name replacepretrained actpostprocess3 3 neck reassemblestage layers 2 projection if pretrained actpostprocess4 3 in name name name replacepretrained actpostprocess4 3 neck reassemblestage layers 3 projection if pretrained actpostprocess4 4 in name name name replacepretrained actpostprocess4 4 neck reassemblestage layers 3 resize if pretrained in name name name replacepretrained dpt if bn in name name name replacebn batchnorm if head in name name name replacehead head head if encoder norm in name name name replaceencoder norm layernorm if auxlayer in name name name replaceauxlayer auxiliaryhead head return name we split up the matrix of each encoder layer into queries keys and values def readinqkvstatedict config for i in rangeconfig numhiddenlayers read in weights bias of input projection layer in timm this is a single matrix bias inprojweight statedict popfdpt encoder layer i attn qkv weight inprojbias statedict popfdpt encoder layer i attn qkv bias next add query keys and values in that order to the state dict statedictfdpt encoder layer i attention attention query weight inprojweight config hiddensize statedictfdpt encoder layer i attention attention query bias inprojbias config hiddensize statedictfdpt encoder layer i attention attention key weight inprojweight config hiddensize config hiddensize 2 statedictfdpt encoder layer i attention attention key bias inprojbias config hiddensize config hiddensize 2 statedictfdpt encoder layer i attention attention value weight inprojweight config hiddensize statedictfdpt encoder layer i attention attention value bias inprojbiasconfig hiddensize we will verify our results on an image of cute cats def prepareimg url http images cocodataset orgval2017000000039769 jpg im image openrequests geturl streamtrue raw return im torch nograd def convertdptcheckpointcheckpointurl pytorchdumpfolderpath pushtohub modelname define dpt configuration based on url config expectedshape getdptconfigcheckpointurl load original statedict from url statedict torch hub loadstatedictfromurlcheckpointurl maplocationcpu remove certain keys removeignorekeysstatedict rename keys for key in statedict copy keys val statedict popkey statedictrenamekeykey val read in qkv matrices readinqkvstatedict config load huggingface model model dptforsemanticsegmentationconfig if ade in checkpointurl else dptfordepthestimationconfig model loadstatedictstatedict model eval check outputs on an image size 480 if ade in checkpointurl else 384 imageprocessor dptimageprocessorsizesize image prepareimg encoding imageprocessorimage returntensorspt forward pass outputs modelencoding logits if ade in checkpointurl else modelencoding predicteddepth assert logits expectedslice torch tensor6 3199 6 3629 6 4148 6 3850 6 3615 6 4166 6 3519 6 3176 6 3575 if ade in checkpointurl expectedslice torch tensor4 0480 4 2420 4 4360 4 3124 4 5693 4 8261 4 5768 4 8965 5 2163 assert outputs shape torch sizeexpectedshape assert torch allcloseoutputs0 0 3 3 expectedslice atol1e4 if ade in checkpointurl else torch allcloseoutputs0 3 3 expectedslice printlooks ok if pytorchdumpfolderpath is not none pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if pushtohub printpushing model to hub model pushtohub repopathornamepathpytorchdumpfolderpath modelname organizationnielsr commitmessageadd model usetempdirtrue imageprocessor pushtohub repopathornamepathpytorchdumpfolderpath modelname organizationnielsr commitmessageadd image processor usetempdirtrue if name main parser argparse argumentparser required parameters parser addargument checkpointurl defaulthttps github comintelisldptreleasesdownload10dptlargemidas2f21e586 pt typestr helpurl of the original dpt checkpoint you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr requiredfalse helppath to the output pytorch model directory parser addargument pushtohub actionstoretrue parser addargument modelname defaultdptlarge typestr requiredfalse helpname of the model in case you re pushing to the hub args parser parseargs convertdptcheckpointargs checkpointurl args pytorchdumpfolderpath args pushtohub args modelname coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert dpt checkpoints from the original repository url https github com isl org dpt tricky here we need to map 4 to 0 3 to 1 2 to 2 and 1 to 3 readout blocks resize blocks we split up the matrix of each encoder layer into queries keys and values read in weights bias of input projection layer in timm this is a single matrix bias next add query keys and values in that order to the state dict we will verify our results on an image of cute cats copy paste tweak model s weights to our dpt structure define dpt configuration based on url load original state_dict from url remove certain keys rename keys read in qkv matrices load huggingface model check outputs on an image forward pass assert logits required parameters
import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_url
from PIL import Image

from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor
from transformers.utils import logging


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


def get_dpt_config(checkpoint_url):
    config = DPTConfig()

    if "large" in checkpoint_url:
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
        config.backbone_out_indices = [5, 11, 17, 23]
        config.neck_hidden_sizes = [256, 512, 1024, 1024]
        expected_shape = (1, 384, 384)

    if "ade" in checkpoint_url:
        config.use_batch_norm_in_fusion_residual = True

        config.num_labels = 150
        repo_id = "huggingface/label-files"
        filename = "ade20k-id2label.json"
        id2label = json.load(open(cached_download(hf_hub_url(repo_id, filename, repo_type="dataset")), "r"))
        id2label = {int(k): v for k, v in id2label.items()}
        config.id2label = id2label
        config.label2id = {v: k for k, v in id2label.items()}
        expected_shape = [1, 150, 480, 480]

    return config, expected_shape


def remove_ignore_keys_(state_dict):
    ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"]
    for k in ignore_keys:
        state_dict.pop(k, None)


def rename_key(name):
    if (
        "pretrained.model" in name
        and "cls_token" not in name
        and "pos_embed" not in name
        and "patch_embed" not in name
    ):
        name = name.replace("pretrained.model", "dpt.encoder")
    if "pretrained.model" in name:
        name = name.replace("pretrained.model", "dpt.embeddings")
    if "patch_embed" in name:
        name = name.replace("patch_embed", "patch_embeddings")
    if "pos_embed" in name:
        name = name.replace("pos_embed", "position_embeddings")
    if "attn.proj" in name:
        name = name.replace("attn.proj", "attention.output.dense")
    if "proj" in name and "project" not in name:
        name = name.replace("proj", "projection")
    if "blocks" in name:
        name = name.replace("blocks", "layer")
    if "mlp.fc1" in name:
        name = name.replace("mlp.fc1", "intermediate.dense")
    if "mlp.fc2" in name:
        name = name.replace("mlp.fc2", "output.dense")
    if "norm1" in name:
        name = name.replace("norm1", "layernorm_before")
    if "norm2" in name:
        name = name.replace("norm2", "layernorm_after")
    if "scratch.output_conv" in name:
        name = name.replace("scratch.output_conv", "head")
    if "scratch" in name:
        name = name.replace("scratch", "neck")
    if "layer1_rn" in name:
        name = name.replace("layer1_rn", "convs.0")
    if "layer2_rn" in name:
        name = name.replace("layer2_rn", "convs.1")
    if "layer3_rn" in name:
        name = name.replace("layer3_rn", "convs.2")
    if "layer4_rn" in name:
        name = name.replace("layer4_rn", "convs.3")
    if "refinenet" in name:
        layer_idx = int(name[len("neck.refinenet") : len("neck.refinenet") + 1])
        
        name = name.replace(f"refinenet{layer_idx}", f"fusion_stage.layers.{abs(layer_idx-4)}")
    if "out_conv" in name:
        name = name.replace("out_conv", "projection")
    if "resConfUnit1" in name:
        name = name.replace("resConfUnit1", "residual_layer1")
    if "resConfUnit2" in name:
        name = name.replace("resConfUnit2", "residual_layer2")
    if "conv1" in name:
        name = name.replace("conv1", "convolution1")
    if "conv2" in name:
        name = name.replace("conv2", "convolution2")
    
    if "pretrained.act_postprocess1.0.project.0" in name:
        name = name.replace("pretrained.act_postprocess1.0.project.0", "neck.reassemble_stage.readout_projects.0.0")
    if "pretrained.act_postprocess2.0.project.0" in name:
        name = name.replace("pretrained.act_postprocess2.0.project.0", "neck.reassemble_stage.readout_projects.1.0")
    if "pretrained.act_postprocess3.0.project.0" in name:
        name = name.replace("pretrained.act_postprocess3.0.project.0", "neck.reassemble_stage.readout_projects.2.0")
    if "pretrained.act_postprocess4.0.project.0" in name:
        name = name.replace("pretrained.act_postprocess4.0.project.0", "neck.reassemble_stage.readout_projects.3.0")
    
    if "pretrained.act_postprocess1.3" in name:
        name = name.replace("pretrained.act_postprocess1.3", "neck.reassemble_stage.layers.0.projection")
    if "pretrained.act_postprocess1.4" in name:
        name = name.replace("pretrained.act_postprocess1.4", "neck.reassemble_stage.layers.0.resize")
    if "pretrained.act_postprocess2.3" in name:
        name = name.replace("pretrained.act_postprocess2.3", "neck.reassemble_stage.layers.1.projection")
    if "pretrained.act_postprocess2.4" in name:
        name = name.replace("pretrained.act_postprocess2.4", "neck.reassemble_stage.layers.1.resize")
    if "pretrained.act_postprocess3.3" in name:
        name = name.replace("pretrained.act_postprocess3.3", "neck.reassemble_stage.layers.2.projection")
    if "pretrained.act_postprocess4.3" in name:
        name = name.replace("pretrained.act_postprocess4.3", "neck.reassemble_stage.layers.3.projection")
    if "pretrained.act_postprocess4.4" in name:
        name = name.replace("pretrained.act_postprocess4.4", "neck.reassemble_stage.layers.3.resize")
    if "pretrained" in name:
        name = name.replace("pretrained", "dpt")
    if "bn" in name:
        name = name.replace("bn", "batch_norm")
    if "head" in name:
        name = name.replace("head", "head.head")
    if "encoder.norm" in name:
        name = name.replace("encoder.norm", "layernorm")
    if "auxlayer" in name:
        name = name.replace("auxlayer", "auxiliary_head.head")

    return name



def read_in_q_k_v(state_dict, config):
    for i in range(config.num_hidden_layers):
        
        in_proj_weight = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.weight")
        in_proj_bias = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.bias")
        
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[: config.hidden_size, :]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
            -config.hidden_size :, :
        ]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :]



def prepare_img():
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    im = Image.open(requests.get(url, stream=True).raw)
    return im


@torch.no_grad()
def convert_dpt_checkpoint(checkpoint_url, pytorch_dump_folder_path, push_to_hub, model_name):
    

    
    config, expected_shape = get_dpt_config(checkpoint_url)
    
    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
    
    remove_ignore_keys_(state_dict)
    
    for key in state_dict.copy().keys():
        val = state_dict.pop(key)
        state_dict[rename_key(key)] = val
    
    read_in_q_k_v(state_dict, config)

    
    model = DPTForSemanticSegmentation(config) if "ade" in checkpoint_url else DPTForDepthEstimation(config)
    model.load_state_dict(state_dict)
    model.eval()

    
    size = 480 if "ade" in checkpoint_url else 384
    image_processor = DPTImageProcessor(size=size)

    image = prepare_img()
    encoding = image_processor(image, return_tensors="pt")

    
    outputs = model(**encoding).logits if "ade" in checkpoint_url else model(**encoding).predicted_depth

    
    expected_slice = torch.tensor([[6.3199, 6.3629, 6.4148], [6.3850, 6.3615, 6.4166], [6.3519, 6.3176, 6.3575]])
    if "ade" in checkpoint_url:
        expected_slice = torch.tensor([[4.0480, 4.2420, 4.4360], [4.3124, 4.5693, 4.8261], [4.5768, 4.8965, 5.2163]])
    assert outputs.shape == torch.Size(expected_shape)
    assert (
        torch.allclose(outputs[0, 0, :3, :3], expected_slice, atol=1e-4)
        if "ade" in checkpoint_url
        else torch.allclose(outputs[0, :3, :3], expected_slice)
    )
    print("Looks ok!")

    if pytorch_dump_folder_path is not None:
        Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
        print(f"Saving model to {pytorch_dump_folder_path}")
        model.save_pretrained(pytorch_dump_folder_path)
        print(f"Saving image processor to {pytorch_dump_folder_path}")
        image_processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing model to hub...")
        model.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add model",
            use_temp_dir=True,
        )
        image_processor.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add image processor",
            use_temp_dir=True,
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--checkpoint_url",
        default="https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt",
        type=str,
        help="URL of the original DPT checkpoint you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=str,
        required=False,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument(
        "--push_to_hub",
        action="store_true",
    )
    parser.add_argument(
        "--model_name",
        default="dpt-large",
        type=str,
        required=False,
        help="Name of the model, in case you're pushing to the hub.",
    )

    args = parser.parse_args()
    convert_dpt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name)