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
![]()