第 12 章 · 真实图片扩散去噪代码实验¶

# 真实图片去噪：把一张花朵照片切成图块，作为训练样本。
photo_clean = np.asarray(Image.fromarray(load_sample_image("flower.jpg")).resize((96, 96))).astype("float32") / 255.0
patch_size = 12
stride = 6

clean_patches = []
patch_positions = []
for row in range(0, photo_clean.shape[0] - patch_size + 1, stride):
    for col in range(0, photo_clean.shape[1] - patch_size + 1, stride):
        clean_patches.append(photo_clean[row:row + patch_size, col:col + patch_size].reshape(-1))
        patch_positions.append((row, col))
clean_patches = np.array(clean_patches)

display(pd.DataFrame({
    "图片": ["flower.jpg"],
    "图片尺寸": [f"{photo_clean.shape[0]}x{photo_clean.shape[1]}"],
    "图块尺寸": [f"{patch_size}x{patch_size}"],
    "训练图块数": [len(clean_patches)],
}))

# 前向加噪：同一张真实图片在不同噪声强度下逐渐丢失结构。
rng = np.random.default_rng(11)
noise = rng.normal(size=photo_clean.shape)
noise_levels = [0.00, 0.12, 0.28, 0.50]
photo_forward = []
for level in noise_levels:
    noisy = np.sqrt(1 - level) * photo_clean + np.sqrt(level) * noise
    photo_forward.append(np.clip(noisy, 0, 1))

photo_forward_df = pd.DataFrame({
    "噪声强度": noise_levels,
    "相对原图 MSE": [mean_squared_error(photo_clean.reshape(-1), img.reshape(-1)) for img in photo_forward],
    "像素标准差": [float(img.std()) for img in photo_forward],
})
display(photo_forward_df.round(4))

# 用局部去噪器看修复目标，再用预训练 DDPM 看真实反向采样轨迹。
train_rng = np.random.default_rng(12)
condition_levels = [0.08, 0.18, 0.32, 0.50]
train_inputs = []
train_targets = []
for level in condition_levels:
    noisy = np.clip(np.sqrt(1 - level) * clean_patches + np.sqrt(level) * train_rng.normal(size=clean_patches.shape), 0, 1)
    level_column = np.full((len(noisy), 1), level)
    train_inputs.append(np.hstack([noisy, level_column]))
    train_targets.append(clean_patches)
train_inputs = np.vstack(train_inputs)
train_targets = np.vstack(train_targets)

photo_denoiser = MLPRegressor(hidden_layer_sizes=(192,), max_iter=220, random_state=12)
photo_denoiser.fit(train_inputs, train_targets)

test_level = noise_levels[-1]
noisy_photo = photo_forward[-1]
noisy_patches = []
for row, col in patch_positions:
    noisy_patches.append(noisy_photo[row:row + patch_size, col:col + patch_size].reshape(-1))
noisy_patches = np.array(noisy_patches)
test_level_column = np.full((len(noisy_patches), 1), test_level)
predicted_patches = np.clip(photo_denoiser.predict(np.hstack([noisy_patches, test_level_column])), 0, 1)

denoised_photo = np.zeros_like(photo_clean)
weight = np.zeros(photo_clean.shape[:2] + (1,), dtype=float)
for (row, col), patch in zip(patch_positions, predicted_patches):
    denoised_photo[row:row + patch_size, col:col + patch_size] += patch.reshape(patch_size, patch_size, 3)
    weight[row:row + patch_size, col:col + patch_size] += 1
denoised_photo = denoised_photo / np.maximum(weight, 1)

photo_denoise_summary = pd.DataFrame(
    [
        {"图像": "高噪声输入", "噪声强度": test_level, "相对原图均方误差": mean_squared_error(photo_clean.reshape(-1), noisy_photo.reshape(-1))},
        {"图像": "条件去噪输出", "噪声强度": test_level, "相对原图均方误差": mean_squared_error(photo_clean.reshape(-1), denoised_photo.reshape(-1))},
    ]
)
condition_summary = pd.DataFrame({
    "训练噪声强度": condition_levels,
    "每档图块数": [len(clean_patches)] * len(condition_levels),
})
display(condition_summary)
display(photo_denoise_summary.round(4))

ddpm_model_name = "google/ddpm-cifar10-32"
ddpm_reverse_images = []
ddpm_reverse_titles = []
ddpm_error = ""
try:
    diffusion_packages = {"torch": "torch>=2.2", "diffusers": "diffusers>=0.30", "accelerate": "accelerate>=0.30"}
    missing = [package for module, package in diffusion_packages.items() if importlib.util.find_spec(module) is None]
    install_packages(missing)

    import torch
    from diffusers import DDPMPipeline
    from diffusers.utils import logging as diffusers_logging

    diffusers_logging.set_verbosity_error()
    diffusers_logging.disable_progress_bar()

    with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
        ddpm_pipe = DDPMPipeline.from_pretrained(ddpm_model_name, use_safetensors=False)
    ddpm_pipe.set_progress_bar_config(disable=True)
    ddpm_pipe.to("cpu")
    ddpm_pipe.unet.eval()
    ddpm_pipe.scheduler.set_timesteps(60)
    generator = torch.Generator(device="cpu").manual_seed(24)
    sample_size = ddpm_pipe.unet.config.sample_size
    if isinstance(sample_size, int):
        sample_size = (sample_size, sample_size)
    sample = torch.randn(
        (1, ddpm_pipe.unet.config.in_channels, sample_size[0], sample_size[1]),
        generator=generator,
    )
    snapshot_indices = {0, 14, 34, 59}
    with torch.no_grad():
        for step_index, timestep in enumerate(ddpm_pipe.scheduler.timesteps):
            model_output = ddpm_pipe.unet(sample, timestep).sample
            sample = ddpm_pipe.scheduler.step(model_output, timestep, sample, generator=generator).prev_sample
            if step_index in snapshot_indices:
                image = (sample[0].permute(1, 2, 0).cpu().numpy() / 2 + 0.5).clip(0, 1)
                ddpm_reverse_images.append(image)
                ddpm_reverse_titles.append(f"反向步 {step_index}")
except Exception as exc:
    ddpm_error = str(exc)[:180]

display(pd.DataFrame([{
    "预训练反向模型": ddpm_model_name,
    "轨迹帧数": len(ddpm_reverse_images),
    "回退信息": ddpm_error,
}]))

# 绘制真实图片的前向加噪、局部去噪和预训练 DDPM 反向轨迹。
has_ddpm = len(ddpm_reverse_images) > 0
row_count = 3 if has_ddpm else 2
fig = plt.figure(figsize=(11.2, 8.8 if has_ddpm else 6.2))
gs = fig.add_gridspec(row_count, 4, height_ratios=[1.0, 1.05, 0.9][:row_count], hspace=0.30, wspace=0.08)

for idx, (img, level) in enumerate(zip(photo_forward, noise_levels)):
    ax = fig.add_subplot(gs[0, idx])
    ax.imshow(img)
    ax.set_title(f"噪声强度 {level:.2f}", fontweight="bold")
    ax.set_xticks([])
    ax.set_yticks([])

for idx, (img, title) in enumerate([
    (photo_clean, "原图"),
    (noisy_photo, "高噪声输入"),
    (denoised_photo, "图块去噪输出"),
    (np.abs(photo_clean - denoised_photo) * 3, "误差 x3"),
]):
    ax = fig.add_subplot(gs[1, idx])
    ax.imshow(np.clip(img, 0, 1))
    ax.set_title(title, fontweight="bold")
    ax.set_xticks([])
    ax.set_yticks([])

if has_ddpm:
    for idx, (img, title) in enumerate(zip(ddpm_reverse_images, ddpm_reverse_titles)):
        ax = fig.add_subplot(gs[2, idx])
        ax.imshow(img)
        ax.set_title(title, fontweight="bold")
        ax.set_xticks([])
        ax.set_yticks([])
    for idx in range(len(ddpm_reverse_images), 4):
        fig.add_subplot(gs[2, idx]).axis("off")

fig.suptitle("扩散去噪：真实图片前向加噪，预训练 DDPM 反向生成", x=0.08, ha="left", fontsize=14, fontweight="bold", color="#0f172a")
plt.tight_layout()
plt.show()