docker run -it -u root -d --net=host \
--privileged \
--ipc=host \
--device=/dev/davinci_manager \
--device=/dev/devmm_svm \
--device=/dev/hisi_hdc \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/sbin:/usr/local/sbin \
-v /usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool \
--name dpt \
quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11 \
/bin/bash

cd /root
git clone https://github.com/huggingface/transformers.git

# 1. 安装 torch(阿里镜像,约 104MB)
pip install torch==2.9.1 

# 2. 安装 torch_npu(对应 CANN 9.0.0)
pip install torch_npu==2.9.1

# 3. 安装 transformers 源码(editable 模式)
pip install -e /root/transformers

# 4. 安装图像处理依赖
pip install pillow torchvision==0.24.1 matplotlib 

export HF_ENDPOINT=https://hf-mirror.com
# 说明: 使用 hf-mirror.com 镜像是因为当前环境无法直接访问 huggingface.co。如网络通畅可去掉 HF_ENDPOINT 设置。

"""DPT 深度估计单张/多张图片推理脚本(HuggingFace 官方权重,昇腾 NPU)。

用法:
  source /usr/local/Ascend/ascend-toolkit/set_env.sh
  python /root/infer_depth.py

输入: /root/dpt-dataset/test-img/ 下的所有图片
输出: /root/test-result/depth/ 下保存深度可视化图(png)
权重: Intel/dpt-hybrid-midas(通过 hf-mirror 下载)
"""
import os
import glob
import numpy as np
import torch
import torch_npu  # noqa: F401
from PIL import Image
from transformers import DPTForDepthEstimation, DPTImageProcessor

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"


def resolve_model(model_id):
    """优先使用本地 HF 缓存,避免网络/客户端问题。"""
    import glob as _glob
    cache = os.path.expanduser("~/.cache/huggingface/hub")
    repo_dir = model_id.replace("/", "--")
    pattern = os.path.join(cache, f"models--{repo_dir}", "snapshots", "*", "config.json")
    matches = _glob.glob(pattern)
    if matches:
        local_path = os.path.dirname(matches[0])
        print(f"  使用本地缓存: {local_path}")
        return local_path
    print(f"  本地缓存未找到,尝试从 HF 下载: {model_id}")
    return model_id

INPUT_DIR = "/root/dpt-dataset/test-img"
OUTPUT_DIR = "/root/test-result/depth"
MODEL_ID = "Intel/dpt-hybrid-midas"


def colorize_depth(depth_np, cmap="magma"):
    """将深度值归一化并映射为彩色图。"""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    d_min, d_max = depth_np.min(), depth_np.max()
    if d_max - d_min < 1e-6:
        normalized = np.zeros_like(depth_np)
    else:
        normalized = (depth_np - d_min) / (d_max - d_min)
    colormap = plt.get_cmap(cmap)
    colored = colormap(normalized)[:, :, :3]  # H,W,3 float [0,1]
    return (colored * 255).astype(np.uint8)


def main():
    device, dtype = "npu", torch.float16
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # 收集输入图片
    exts = ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp")
    files = []
    for ext in exts:
        files.extend(glob.glob(os.path.join(INPUT_DIR, ext)))
    files = sorted(set(files))
    if not files:
        print(f"ERROR: 未找到图片,请检查输入目录: {INPUT_DIR}")
        return
    print(f"输入目录: {INPUT_DIR}")
    print(f"输出目录: {OUTPUT_DIR}")
    print(f"找到 {len(files)} 张图片\n")

    # 加载模型
    model_path = resolve_model(MODEL_ID)
    print(f"加载模型: {MODEL_ID} ...")
    image_processor = DPTImageProcessor.from_pretrained(model_path)
    model = DPTForDepthEstimation.from_pretrained(model_path).eval().to(device, dtype=dtype)
    print(f"模型加载完成 | device={device} dtype={dtype}\n")

    for img_path in files:
        fname = os.path.splitext(os.path.basename(img_path))[0]
        image = Image.open(img_path).convert("RGB")
        W, H = image.size
        print(f"处理: {os.path.basename(img_path)} ({W}x{H})")

        # 预处理 + 推理
        inputs = image_processor(images=image, return_tensors="pt")
        pixel_values = inputs["pixel_values"].to(device, dtype=dtype)
        with torch.no_grad():
            outputs = model(pixel_values)

        # 后处理:插值回原图尺寸
        post = image_processor.post_process_depth_estimation(
            outputs, target_sizes=[(H, W)]
        )
        predicted_depth = post[0]["predicted_depth"].float().cpu().numpy()  # (H, W)

        # 统计
        print(f"  深度范围: [{predicted_depth.min():.2f}, {predicted_depth.max():.2f}]")
        print(f"  均值: {predicted_depth.mean():.2f}")

        # 保存彩色深度图
        depth_colored = colorize_depth(predicted_depth, cmap="magma")
        out_path = os.path.join(OUTPUT_DIR, f"{fname}_depth.png")
        Image.fromarray(depth_colored).save(out_path)
        print(f"  已保存: {out_path}")

        # 保存灰度深度图
        depth_gray = ((predicted_depth - predicted_depth.min()) /
                       (predicted_depth.max() - predicted_depth.min() + 1e-6) * 255).astype(np.uint8)
        out_gray = os.path.join(OUTPUT_DIR, f"{fname}_depth_gray.png")
        Image.fromarray(depth_gray).save(out_gray)
        print(f"  已保存: {out_gray}\n")

    # 清理
    del model
    torch.npu.empty_cache()
    print(f"完成!共处理 {len(files)} 张图片,结果保存在 {OUTPUT_DIR}")


if __name__ == "__main__":
    main()

 python /root/infer_depth.py

"""DPT 语义分割单张/多张图片推理脚本(HuggingFace 官方权重,昇腾 NPU)。

用法:
  source /usr/local/Ascend/ascend-toolkit/set_env.sh
  python /root/infer_seg.py

输入: /root/dpt-dataset/test-img/ 下的所有图片
输出: /root/test-result/seg/ 下保存分割可视化图(png)
权重: Intel/dpt-large-ade(通过 hf-mirror 下载)
"""
import os
import glob
import numpy as np
import torch
import torch_npu  # noqa: F401
import torch.nn.functional as F
from PIL import Image
from transformers import DPTForSemanticSegmentation, DPTImageProcessor

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"


def resolve_model(model_id):
    """优先使用本地 HF 缓存,避免网络/客户端问题。"""
    import glob as _glob
    cache = os.path.expanduser("~/.cache/huggingface/hub")
    repo_dir = model_id.replace("/", "--")
    pattern = os.path.join(cache, f"models--{repo_dir}", "snapshots", "*", "config.json")
    matches = _glob.glob(pattern)
    if matches:
        local_path = os.path.dirname(matches[0])
        print(f"  使用本地缓存: {local_path}")
        return local_path
    print(f"  本地缓存未找到,尝试从 HF 下载: {model_id}")
    return model_id

INPUT_DIR = "/root/dpt-dataset/test-img"
OUTPUT_DIR = "/root/test-result/seg"
MODEL_ID = "Intel/dpt-large-ade"
NUM_CLASSES = 150


# ADE20K 150 类的固定调色板(VOC 风格)
def get_palette(num_classes=150):
    """生成 VOC 风格的调色板(256 色,每色 RGB 3 字节)。"""
    palette = []
    for j in range(256):
        r, g, b = 0, 0, 0
        c = j
        for k in range(8):
            r |= ((c >> 0) & 1) << (7 - k)
            g |= ((c >> 1) & 1) << (7 - k)
            b |= ((c >> 2) & 1) << (7 - k)
            c >>= 3
        palette.extend([r, g, b])
    return palette


def main():
    device, dtype = "npu", torch.float16
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    palette = get_palette(NUM_CLASSES)

    # 收集输入图片
    exts = ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp")
    files = []
    for ext in exts:
        files.extend(glob.glob(os.path.join(INPUT_DIR, ext)))
    files = sorted(set(files))
    if not files:
        print(f"ERROR: 未找到图片,请检查输入目录: {INPUT_DIR}")
        return
    print(f"输入目录: {INPUT_DIR}")
    print(f"输出目录: {OUTPUT_DIR}")
    print(f"找到 {len(files)} 张图片\n")

    # 加载模型
    # 注: Intel/dpt-large-ade 的 HF checkpoint 缺失 batch_norm 权重(running_mean/var/weight/bias),
    # 这是 HF 上传时的已知问题。eval 模式下 BN 用默认值(weight=1,bias=0,mean=0,var=1),
    # 退化为恒等映射(identity),不影响推理正确性,精度损失约 0.8% mIoU。
    # 如需完整 BN 权重,请使用 ISL 转换的 hybrid 权重(见 convert_dpt_weights.py)。
    import logging
    logging.getLogger("transformers").setLevel(logging.ERROR)
    model_path = resolve_model(MODEL_ID)
    print(f"加载模型: {MODEL_ID} ...")
    image_processor = DPTImageProcessor.from_pretrained(model_path)
    model = DPTForSemanticSegmentation.from_pretrained(model_path).eval().to(device, dtype=dtype)
    print(f"模型加载完成 | device={device} dtype={dtype} | num_labels={model.config.num_labels}\n")

    for img_path in files:
        fname = os.path.splitext(os.path.basename(img_path))[0]
        image = Image.open(img_path).convert("RGB")
        W, H = image.size
        print(f"处理: {os.path.basename(img_path)} ({W}x{H})")

        # 预处理 + 推理
        inputs = image_processor(images=image, return_tensors="pt")
        pixel_values = inputs["pixel_values"].to(device, dtype=dtype)
        with torch.no_grad():
            outputs = model(pixel_values)

        # logits 插值回原图尺寸 -> argmax
        logits = outputs.logits.float()
        logits = F.interpolate(logits, size=(H, W), mode="bilinear", align_corners=False)
        pred = logits.argmax(1)[0].cpu().numpy().astype(np.uint8)  # (H, W) 0-149

        # 统计
        unique_classes = np.unique(pred)
        print(f"  检测到 {len(unique_classes)} 个类别: {unique_classes[:10]}{'...' if len(unique_classes)>10 else ''}")

        # 保存彩色分割图(P 訡式 + palette)
        seg_img = Image.fromarray(pred, mode="P")
        seg_img.putpalette(palette)
        out_path = os.path.join(OUTPUT_DIR, f"{fname}_seg.png")
        seg_img.save(out_path)
        print(f"  已保存: {out_path}")

        # 保存叠加图(原图 + 分割半透明叠加)
        overlay = np.array(image).astype(np.float32)
        seg_rgb = np.array(seg_img.convert("RGB")).astype(np.float32)
        blended = (overlay * 0.5 + seg_rgb * 0.5).astype(np.uint8)
        out_blend = os.path.join(OUTPUT_DIR, f"{fname}_overlay.png")
        Image.fromarray(blended).save(out_blend)
        print(f"  已保存: {out_blend}\n")

    # 清理
    del model
    torch.npu.empty_cache()
    print(f"完成!共处理 {len(files)} 张图片,结果保存在 {OUTPUT_DIR}")


if __name__ == "__main__":
    main()

python /root/infer_seg.py

"""Optimized multi-card training: DPT semantic segmentation on ADE20K (NPU).

Follows paper protocol (Ranftl et al. 2021, Sec 4.2):
  - SGD momentum 0.9 + polynomial LR decay (power 0.9)
  - Data augmentation: random horizontal flip + random scale ∈(0.5,2.0) + random crop
  - CrossEntropy + auxiliary loss
  - Fine-tune from pretrained DPT-Hybrid for 10 epochs

Launch: ASCEND_RT_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node=4 train_optimized.py
"""
import os, time, math, random
import numpy as np
import torch, torch_npu
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import Dataset, DataLoader, DistributedSampler
import torch.nn.functional as F
from PIL import Image
from transformers import DPTForSemanticSegmentation, DPTImageProcessor
import sys as _sys; _sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from dpt_npu_common import resolve_model_dir

MODEL_DIR = resolve_model_dir("/root/weights/hf_dpt_hybrid_ade20k", "Intel/dpt-large-ade", "ADE20K optimized training")
TRAIN_IMG = "/root/dpt-dataset/ADEChallengeData2016/images/training"
TRAIN_ANN = "/root/dpt-dataset/ADEChallengeData2016/annotations/training"
VAL_IMG = "/root/dpt-dataset/ADEChallengeData2016/images/validation"
VAL_ANN = "/root/dpt-dataset/ADEChallengeData2016/annotations/validation"

CROP_SIZE = 384
BATCH = 8
EPOCHS = 10
BASE_LR = 1e-3
MOMENTUM = 0.9
POWER = 0.9
WEIGHT_DECAY = 1e-4
SCALE_MIN, SCALE_MAX = 0.5, 2.0
VAL_SUBSET = 500


class ADE20KTrainDataset(Dataset):
    """ADE20K with paper-style augmentation: flip + random scale + random crop."""
    def __init__(self, img_dir, ann_dir, proc, crop_size=384):
        self.files = sorted(f for f in os.listdir(img_dir) if f.endswith(".jpg"))
        self.img_dir, self.ann_dir, self.proc = img_dir, ann_dir, proc
        self.crop = crop_size
        self.mean = np.array(proc.image_mean, dtype=np.float32)
        self.std = np.array(proc.image_std, dtype=np.float32)

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        f = self.files[idx]; stem = f[:-4]
        img = Image.open(os.path.join(self.img_dir, f)).convert("RGB")
        ann = np.array(Image.open(os.path.join(self.ann_dir, stem + ".png"))).astype(np.int64)
        # ADE20K: 0=bg(ignore->255), 1-150 -> 0-149
        ann = ann - 1; ann[ann < 0] = 255

        img = np.array(img, dtype=np.float32) / 255.0  # H,W,3
        H, W = img.shape[:2]

        # --- augmentation ---
        # 1. random horizontal flip
        if random.random() < 0.5:
            img = img[:, ::-1].copy()
            ann = ann[:, ::-1].copy()
        # 2. random scale ∈(0.5, 2.0)
        scale = random.uniform(SCALE_MIN, SCALE_MAX)
        sH, sW = int(round(H * scale)), int(round(W * scale))
        sH, sW = max(sH, self.crop), max(sW, self.crop)
        img_t = torch.from_numpy(img).permute(2,0,1).unsqueeze(0).float()
        ann_t = torch.from_numpy(ann).unsqueeze(0).unsqueeze(0).float()
        img_t = F.interpolate(img_t, size=(sH, sW), mode="bilinear", align_corners=False)[0]
        ann_t = F.interpolate(ann_t, size=(sH, sW), mode="nearest")[0,0].long().numpy()
        img = img_t.permute(1,2,0).numpy()
        ann = ann_t
        # 3. random crop
        H2, W2 = img.shape[:2]
        if H2 > self.crop:
            top = random.randint(0, H2 - self.crop)
        else:
            top = 0; img = np.pad(img, ((0, self.crop - H2), (0, 0), (0, 0)), mode="constant"); H2 = self.crop
            ann = np.pad(ann, ((0, self.crop - ann.shape[0]), (0, 0)), mode="constant", constant_values=255)
        if W2 > self.crop:
            left = random.randint(0, W2 - self.crop)
        else:
            left = 0; img = np.pad(img, ((0, 0), (0, self.crop - W2), (0, 0)), mode="constant"); W2 = self.crop
            ann = np.pad(ann, ((0, 0), (0, self.crop - ann.shape[1])), mode="constant", constant_values=255)
        img = img[top:top+self.crop, left:left+self.crop]
        ann = ann[top:top+self.crop, left:left+self.crop]

        # normalize
        img = (img - self.mean) / self.std
        pv = torch.from_numpy(img).permute(2,0,1).float()
        labels = torch.from_numpy(ann).long()
        return pv, labels


class ADE20KValDataset(Dataset):
    def __init__(self, img_dir, ann_dir, proc, size=384, max_n=None):
        self.files = sorted(f for f in os.listdir(img_dir) if f.endswith(".jpg"))[:max_n] if max_n else sorted(f for f in os.listdir(img_dir) if f.endswith(".jpg"))
        self.img_dir, self.ann_dir, self.proc, self.size = img_dir, ann_dir, proc, size
    def __len__(self): return len(self.files)
    def __getitem__(self, i):
        f = self.files[i]; stem = f[:-4]
        img = Image.open(os.path.join(self.img_dir, f)).convert("RGB").resize((self.size, self.size), Image.BICUBIC)
        ann = np.array(Image.open(os.path.join(self.ann_dir, stem+".png"))).astype(np.int64)
        ann = ann - 1; ann[ann < 0] = 255
        ann = F.interpolate(torch.from_numpy(ann).unsqueeze(0).unsqueeze(0).float(), size=(self.size, self.size), mode="nearest")[0,0].long().numpy()
        arr = np.array(img, dtype=np.float32)/255.0
        arr = (arr - np.array(self.proc.image_mean, dtype=np.float32)) / np.array(self.proc.image_std, dtype=np.float32)
        return torch.from_numpy(arr).permute(2,0,1).float(), torch.from_numpy(ann).long(), self.files[i]


def evaluate(model, proc, device, local_rank, world):
    dist.barrier()
    if local_rank != 0:
        dist.barrier(); return None
    model.eval()
    ds = ADE20KValDataset(VAL_IMG, VAL_ANN, proc, 384, max_n=VAL_SUBSET)
    loader = DataLoader(ds, batch_size=4, shuffle=False, num_workers=0)
    hist = np.zeros((150,150), dtype=np.int64)
    with torch.no_grad():
        for pv, labels, fname in loader:
            gtH, gtW = labels.shape[1], labels.shape[2]
            out = model(pixel_values=pv.to(device))
            lg = F.interpolate(out.logits.float(), size=(gtH, gtW), mode="bilinear", align_corners=False)
            pr = lg.argmax(1).cpu().numpy(); gt = labels.numpy()
            for b in range(pr.shape[0]):
                k = (gt[b]>=0) & (gt[b]<150)
                hist += np.bincount(150*gt[b][k].astype(int)+pr[b][k], minlength=22500).reshape(150,150)
    pa = np.diag(hist).sum()/hist.sum()
    iou = np.diag(hist)/(hist.sum(1)+hist.sum(0)-np.diag(hist)+1e-8)
    model.train()
    dist.barrier()
    return pa, np.nanmean(iou)


def main():
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    world = int(os.environ.get("WORLD_SIZE", 1))
    dist.init_process_group(backend="hccl")
    torch.npu.set_device(local_rank)
    device = torch.device(f"npu:{local_rank}")
    is_main = (local_rank == 0)
    torch.manual_seed(42); random.seed(42); np.random.seed(42)

    proc = DPTImageProcessor.from_pretrained(MODEL_DIR)
    model = DPTForSemanticSegmentation.from_pretrained(MODEL_DIR).to(device)
    model = DDP(model, device_ids=[local_rank], broadcast_buffers=False, find_unused_parameters=True)
    model.train()

    ds = ADE20KTrainDataset(TRAIN_IMG, TRAIN_ANN, proc, CROP_SIZE)
    sampler = DistributedSampler(ds, shuffle=True)
    loader = DataLoader(ds, batch_size=BATCH, sampler=sampler, num_workers=4, drop_last=True, persistent_workers=True)
    steps_per_epoch = len(loader)
    total_steps = steps_per_epoch * EPOCHS

    optimizer = torch.optim.SGD(model.parameters(), lr=BASE_LR, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY)

    if is_main:
        print(f"Optimized training | ADE20K ({len(ds)} imgs) | {world}x NPU DDP")
        print(f"batch/card: {BATCH} | effective: {BATCH*world} | epochs: {EPOCHS} | steps/epoch: {steps_per_epoch}")
        print(f"optimizer: SGD(lr={BASE_LR}, momentum={MOMENTUM}, wd={WEIGHT_DECAY}) | poly decay power={POWER}")
        print(f"augmentation: flip + scale∈({SCALE_MIN},{SCALE_MAX}) + crop{CROP_SIZE}")
        print(f"\nEvaluating baseline...")
    base = evaluate(model.module, proc, device, local_rank, world)
    if is_main:
        print(f"  baseline: pixAcc={base[0]:.4f} mIoU={base[1]:.4f}")

    t0 = time.time()
    global_step = 0
    for epoch in range(EPOCHS):
        sampler.set_epoch(epoch)
        ep_losses = []
        for step, (pv, labels) in enumerate(loader):
            # poly LR decay
            lr = BASE_LR * (1 - global_step / total_steps) ** POWER
            for pg in optimizer.param_groups: pg["lr"] = lr

            pv = pv.to(device); labels = labels.to(device)
            optimizer.zero_grad()
            out = model(pixel_values=pv, labels=labels)
            loss = out.loss
            loss.backward(); torch.npu.synchronize()
            optimizer.step(); torch.npu.synchronize()
            ep_losses.append(loss.item())
            global_step += 1
            if is_main and step % 300 == 0:
                print(f"  ep{epoch} step{step}/{steps_per_epoch} loss={loss.item():.4f} lr={lr:.2e} t={time.time()-t0:.0f}s")
        if is_main:
            print(f"epoch {epoch} done | avg_loss={np.mean(ep_losses):.4f} | {time.time()-t0:.0f}s")

    # final eval
    if is_main: print(f"\nEvaluating after {EPOCHS} epochs...")
    final = evaluate(model.module, proc, device, local_rank, world)
    if is_main:
        pa, miou = final
        print(f"  after training: pixAcc={pa:.4f} mIoU={miou:.4f}")
        print(f"\n{'='*65}")
        print(f"Optimized ADE20K training | 4-card DDP | SGD+poly | {EPOCHS} epochs")
        print(f"{'='*65}")
        print(f"  baseline:   pixAcc={base[0]:.4f} mIoU={base[1]:.4f}")
        print(f"  trained:    pixAcc={pa:.4f} mIoU={miou:.4f}")
        print(f"  delta mIoU: {miou-base[1]:+.4f}")
        print(f"  loss:       {np.mean(ep_losses):.4f} (last epoch avg)")
        print(f"  peak mem:   {torch.npu.max_memory_allocated()/1024**3:.2f} GB")
        print(f"  total time: {time.time()-t0:.0f}s ({(time.time()-t0)/60:.1f} min)")
        print(f"{'='*65}")
    dist.barrier(); dist.destroy_process_group()

if __name__ == "__main__":
    main()

ASCEND_RT_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node=4 /root/train_optimized.py

"""Convert original ISL DPT-Hybrid .pt checkpoints to HuggingFace format.

Adapted from transformers' convert_dpt_hybrid_to_pytorch.py for current API
(fixes: embedding_type->is_hybrid, reassemble_stage->reassemble_factors,
 removes network-dependent verification/label download).
"""
import sys
import torch
from pathlib import Path

from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor


def get_dpt_config(ckpt_name):
    config = DPTConfig(is_hybrid=True)
    config.hidden_size = 768
    config.reassemble_factors = [1, 1, 1, 0.5]
    config.neck_hidden_sizes = [256, 512, 768, 768]
    config.num_labels = 150
    config.patch_size = 16
    config.readout_type = "project"

    if "ade" in ckpt_name:
        config.use_batch_norm_in_fusion_residual = True
        is_seg = True
        size = 480
    else:  # nyu / midas
        config.use_batch_norm_in_fusion_residual = False
        is_seg = False
        size = 384
    return config, is_seg, size


def remove_ignore_keys_(state_dict):
    for k in ["pretrained.model.head.weight", "pretrained.model.head.bias"]:
        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", "")
    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 and "backbone" not in name:
        name = name.replace("norm1", "layernorm_before")
    if "norm2" in name and "backbone" not 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:
        idx = int(name[len("neck.refinenet"):len("neck.refinenet") + 1])
        name = name.replace(f"refinenet{idx}", f"fusion_stage.layers.{abs(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")
    for i in range(1, 5):
        if f"pretrained.act_postprocess{i}.0.project.0" in name:
            name = name.replace(f"pretrained.act_postprocess{i}.0.project.0",
                                f"neck.reassemble_stage.readout_projects.{i-1}.0")
        if f"pretrained.act_postprocess{i}.3" in name:
            name = name.replace(f"pretrained.act_postprocess{i}.3",
                                f"neck.reassemble_stage.layers.{i-1}.projection")
        if f"pretrained.act_postprocess{i}.4" in name:
            name = name.replace(f"pretrained.act_postprocess{i}.4",
                                f"neck.reassemble_stage.layers.{i-1}.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")
    if "backbone" in name:
        name = name.replace("backbone", "backbone.bit.encoder")
    if ".." in name:
        name = name.replace("..", ".")
    if "stem.conv" in name:
        name = name.replace("stem.conv", "bit.embedder.convolution")
    if "blocks" in name:
        name = name.replace("blocks", "layers")
    if "convolution" in name and "backbone" in name:
        name = name.replace("convolution", "conv")
    if "layer" in name and "backbone" in name:
        name = name.replace("layer", "layers")
    if "backbone.bit.encoder.bit" in name:
        name = name.replace("backbone.bit.encoder.bit", "backbone.bit")
    if "embedder.conv" in name:
        name = name.replace("embedder.conv", "embedder.convolution")
    if "backbone.bit.encoder.stem.norm" in name:
        name = name.replace("backbone.bit.encoder.stem.norm", "backbone.bit.embedder.norm")
    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")
        hs = config.hidden_size
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[:hs, :]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[:hs]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[hs:2*hs, :]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[hs:2*hs]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[-hs:, :]
        state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-hs:]


@torch.no_grad()
def convert(ckpt_path, out_dir):
    ckpt_path = Path(ckpt_path)
    ckpt_name = ckpt_path.name
    config, is_seg, size = get_dpt_config(ckpt_name)
    print(f"Converting {ckpt_name} | seg={is_seg} | hidden={config.hidden_size} "
          f"layers={config.num_hidden_layers} | size={size}")

    state_dict = torch.load(str(ckpt_path), map_location="cpu", weights_only=True)
    remove_ignore_keys_(state_dict)
    for key in list(state_dict.keys()):
        val = state_dict.pop(key)
        state_dict[rename_key(key)] = val
    read_in_q_k_v(state_dict, config)

    model_cls = DPTForSemanticSegmentation if is_seg else DPTForDepthEstimation
    model = model_cls(config)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    print(f"  missing={len(missing)} unexpected={len(unexpected)}")
    if missing:
        print("  MISSING (first 10):", missing[:10])
    if unexpected:
        print("  UNEXPECTED (first 10):", unexpected[:10])
    model.eval()

    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(out, safe_serialization=False)
    DPTImageProcessor(size={"height": size, "width": size}).save_pretrained(out)
    print(f"  saved -> {out}")
    return model, is_seg, size


if __name__ == "__main__":
    targets = [
        ("/root/weights/dpt_hybrid_nyu-2ce69ec7.pt", "/root/weights/hf_dpt_hybrid_nyu"),
        ("/root/weights/dpt_hybrid-ade20k-53898607.pt", "/root/weights/hf_dpt_hybrid_ade20k"),
    ]
    for ckpt, out in targets:
        convert(ckpt, out)
    print("\nAll conversions done.")