QDM: Quadtree-Based Region-Adaptive Sparse Diffusion Models for Efficient Image

15 阅读5分钟

QDM模型(github.com/linYDTHU/QD…

Quadtree-Based Region-Adaptive Sparse Diffusion Models (基于四叉树的区域自适应稀疏扩散模型)

论文核心观点是因为扩散模型计算太慢,特别是在不需要进行复杂计算的区域比如说背景,蓝天白云等,而在具体需要细节的地方比如物体核心区域。

QDM: 基于四叉树区域自适应的稀疏扩散模型用于高效图像超分

摘要: 基于深度学习的图像超分方法通常需要像素级的计算整张图片,甚至在一些同质化区域(就是噪声,纹理等相同的区域,比如纯黑色的背景)。QDM(四叉树扩散模型)区域自适应的(就是通过简单的计算方法比如方差判断出图片某个patch是否属于复杂区域)扩散模型使用一个四叉树结构区别计算同质区域和发杂区域。通过四叉树对模糊图的提取来引导扩散模型进行计算,四叉树生成的叶子节点就是需要扩散模型计算的地方。

环境准备: 环境安装,readme有 模型需要,第一阶段需要的VAE自编码器,灰度图使用下载Medical SR Task,彩色图下载Real-world SR Task。在项目根目录新建文件夹weights,把模型放在这里面。 预训练模型根据模型大小不同下载不同版本,GitHub上给的是L版本。

自制数据集准备:

(1)灰度图的准备,需要把图片变为正方形 import cv2 import os import numpy as np import shutil from pathlib import Path

def make_padded_dataset(input_folder, output_root, target_size=128, scale=4): """ input_folder: 原始图片文件夹 output_root: 输出路径 target_size: 目标正方形大小 (建议 128 或 256) scale: 缩放倍数 (4) """

# 路径准备
train_hr = os.path.join(output_root, 'train/HR')
train_lr = os.path.join(output_root, 'train/LR')
val_hr = os.path.join(output_root, 'val/HR')
val_lr = os.path.join(output_root, 'val/LR')

# 暴力清空旧数据 (防止冲突)
if os.path.exists(output_root):
    shutil.rmtree(output_root)

for p in [train_hr, train_lr, val_hr, val_lr]:
    os.makedirs(p, exist_ok=True)
    
valid_exts = {'.jpg', '.jpeg', '.png', '.bmp'}
files = [f for f in os.listdir(input_folder) if Path(f).suffix.lower() in valid_exts]

print(f"正在处理 {len(files)} 张图片...")
print(f"🎯 策略: 保持原图比例,周围补黑边,统一成 {target_size}x{target_size}")

# 前8张做验证集
val_count = 8

for i, filename in enumerate(files):
    img_path = os.path.join(input_folder, filename)
    
    try:
        # 读取原图
        img = cv2.imread(img_path)
        if img is None: continue
        
        h, w = img.shape[:2]
        
        # --- 核心逻辑: 补黑边 (Padding) ---
        
        # 1. 如果图片比目标还大,先等比例缩小 (防止塞不进去)
        if h > target_size or w > target_size:
            scale_down = min(target_size / h, target_size / w)
            new_h, new_w = int(h * scale_down), int(w * scale_down)
            img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
            h, w = new_h, new_w
        
        # 2. 创建黑色画布
        # 如果是彩色图用 (target_size, target_size, 3)
        # 如果是灰度图用 (target_size, target_size)
        if len(img.shape) == 3:
            canvas = np.zeros((target_size, target_size, 3), dtype=np.uint8)
        else:
            canvas = np.zeros((target_size, target_size), dtype=np.uint8)
        
        # 3. 计算居中坐标
        y_offset = (target_size - h) // 2
        x_offset = (target_size - w) // 2
        
        # 4. 把图贴上去
        canvas[y_offset:y_offset+h, x_offset:x_offset+w] = img
        img_hr_final = canvas
        
        # --- 生成 LR ---
        lr_size = target_size // scale
        img_lr_final = cv2.resize(img_hr_final, (lr_size, lr_size), interpolation=cv2.INTER_CUBIC)
        
        # 保存
        save_hr_dir = val_hr if i < val_count else train_hr
        save_lr_dir = val_lr if i < val_count else train_lr
        
        cv2.imwrite(os.path.join(save_hr_dir, filename), img_hr_final)
        cv2.imwrite(os.path.join(save_lr_dir, filename), img_lr_final)
            
    except Exception as e:
        print(f"Error processing {filename}: {e}")

print(f"\n✅ 处理完成!数据已生成在: {output_root}")
print("👉 现在的图片都是标准的正方形了,可以直接去训练!")

(2)彩色图准备

import cv2 import os import numpy as np import shutil from pathlib import Path

def make_color_padded_dataset(input_folder, output_root, target_size=128, scale=4): """ Color 版数据处理: 1. 保持 RGB 彩色,不转灰度。 2. 依然补黑边变成正方形 (适配 DiT)。 """

# 路径准备
train_hr = os.path.join(output_root, 'train/HR')
train_lr = os.path.join(output_root, 'train/LR')
val_hr = os.path.join(output_root, 'val/HR')
val_lr = os.path.join(output_root, 'val/LR')

if os.path.exists(output_root):
    shutil.rmtree(output_root)

for p in [train_hr, train_lr, val_hr, val_lr]:
    os.makedirs(p, exist_ok=True)
    
valid_exts = {'.jpg', '.jpeg', '.png', '.bmp'}
files = [f for f in os.listdir(input_folder) if Path(f).suffix.lower() in valid_exts]

print(f"🌈 正在处理 {len(files)} 张图片 (彩色模式)...")

val_count = 8 # 验证集数量

for i, filename in enumerate(files):
    img_path = os.path.join(input_folder, filename)
    
    try:
        # 读取 (OpenCV 默认就是 BGR 彩色)
        img = cv2.imread(img_path)
        if img is None: continue
        
        # 【重要】如果原图真的是单通道灰度,强制转成 3 通道,防止报错
        if len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        
        h, w = img.shape[:2]
        
        # --- 补黑边逻辑 (和之前一样,但画布是彩色的) ---
        if h > target_size or w > target_size:
            scale_down = min(target_size / h, target_size / w)
            new_h, new_w = int(h * scale_down), int(w * scale_down)
            img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
            h, w = new_h, new_w
        
        # 创建 3 通道黑色画布 (target_size, target_size, 3)
        canvas = np.zeros((target_size, target_size, 3), dtype=np.uint8)
        
        y_offset = (target_size - h) // 2
        x_offset = (target_size - w) // 2
        
        canvas[y_offset:y_offset+h, x_offset:x_offset+w] = img
        img_hr_final = canvas
        
        # --- 生成 LR ---
        lr_size = target_size // scale
        img_lr_final = cv2.resize(img_hr_final, (lr_size, lr_size), interpolation=cv2.INTER_CUBIC)
        
        # 保存
        save_hr_dir = val_hr if i < val_count else train_hr
        save_lr_dir = val_lr if i < val_count else train_lr
        
        cv2.imwrite(os.path.join(save_hr_dir, filename), img_hr_final)
        cv2.imwrite(os.path.join(save_lr_dir, filename), img_lr_final)
            
    except Exception as e:
        print(f"Error: {e}")

print(f"\n✅ 彩色数据处理完成!位置: {output_root}")

启动训练命令,找到~/.../Super-Resolution/QDM-main/configs/train/medx4_qdm_l.yaml 文件修改参数,数据路径,训练参数等 trainer: target: trainer.TrainerQDM

启动训练 CUDA_VISIBLE_DEVICES=0 torchrun --standalone --nproc_per_node=1 main.py --cfg_path configs/train/medx4_qdm_l.yaml --save_dir experiments/finetune

  


![]()