[论文解读] Bayesian Image Reconstruction using Deep Generative Models

Razvan Marinescu, Daniel Moyer+1 more|arXiv (Cornell University)|Dec 8, 2020

Advanced Image Processing Techniques参考文献 58被引用 11

一句话总结

本文提出贝叶斯生成模型重构（BRGM），一种利用预训练的StyleGAN2生成器作为先验，通过潜在向量上的贝叶斯最大后验（MAP）估计实现图像重构的方法。在不微调的情况下，BRGM在超分辨率和图像修复任务中均表现出色，涵盖医学影像等多样化数据集，通过变分推理在推理时优化潜在码，以采样多种合理的重构结果。

ABSTRACT

Machine learning models are commonly trained end-to-end and in a supervised setting, using paired (input, output) data. Examples include recent super-resolution methods that train on pairs of (low-resolution, high-resolution) images. However, these end-to-end approaches require re-training every time there is a distribution shift in the inputs (e.g., night images vs daylight) or relevant latent variables (e.g., camera blur or hand motion). In this work, we leverage state-of-the-art (SOTA) generative models (here StyleGAN2) for building powerful image priors, which enable application of Bayes' theorem for many downstream reconstruction tasks. Our method, Bayesian Reconstruction through Generative Models (BRGM), uses a single pre-trained generator model to solve different image restoration tasks, i.e., super-resolution and in-painting, by combining it with different forward corruption models. We keep the weights of the generator model fixed, and reconstruct the image by estimating the Bayesian maximum a-posteriori (MAP) estimate over the input latent vector that generated the reconstructed image. We further use variational inference to approximate the posterior distribution over the latent vectors, from which we sample multiple solutions. We demonstrate BRGM on three large and diverse datasets: (i) 60,000 images from the Flick Faces High Quality dataset (ii) 240,000 chest X-rays from MIMIC III and (iii) a combined collection of 5 brain MRI datasets with 7,329 scans. Across all three datasets and without any dataset-specific hyperparameter tuning, our simple approach yields performance competitive with current task-specific state-of-the-art methods on super-resolution and in-painting, while being more generalisable and without requiring any training. Our source code and pre-trained models are available online: https://razvanmarinescu.github.io/brgm/.

研究动机与目标

• 为解决端到端监督学习在图像重构中的局限性，特别是面对分布偏移或新退化模型时需要重新训练的问题。
• 通过单一预训练的生成模型实现对多种图像恢复任务的泛化，避免针对特定任务的微调。
• 提供一种贝叶斯框架，不仅估计单一重构结果，还通过后验采样估计一组合理的解的分布。
• 在多样化的真实世界数据集（包括医学影像）上展示该方法的有效性，且无需针对数据集调整超参数。
• 通过利用预训练的生成模型作为强图像先验，减少对大规模成对训练数据和计算资源的依赖。

提出的方法

• 该方法使用预训练的StyleGAN2生成器作为先验，以建模自然图像的分布，且在推理过程中保持生成器权重固定。
• 对于给定的退化输入图像 I，该方法通过贝叶斯MAP估计，寻找使后验 p(w|I) 最大的潜在向量 w*：w* = argmax_w p(w)p(I|f∘G(w))。
• 前向退化模型 f 作用于生成图像 G(w)，以模拟观测到的输入 I，且 f 可以是任意已知的退化过程（如下采样、掩码、模糊等）。
• 使用变分推理近似真实后验 p(w|I)，从而实现从潜在空间中采样多种合理的重构结果。
• 该方法在推理时应用，无需任何模型微调，因此可灵活适应新的退化模型或输入分布。
• 该方法在三个数据集上进行了评估：FlickFaces、MIMIC-III胸部X光片和脑部MRI扫描，各项任务中表现一致。

实验结果

研究问题

• RQ1 是否可以使用单一预训练的生成模型在不重新训练的情况下解决多种图像恢复任务？
• RQ2 与端到端监督方法相比，使用深度生成先验的贝叶斯MAP估计在重构质量与多样性方面表现如何？
• RQ3 在预训练生成器的潜在空间上进行变分推理，能否提供有意义的、合理的重构解的分布？
• RQ4 该方法在多样化数据分布（包括存在领域偏移的医学影像）上的泛化能力如何？
• RQ5 该方法是否在超分辨率和图像修复任务中优于最先进模型，同时避免了像素级损失方法常见的平均效应？

主要发现

• 在FFHQ超分辨率基准测试中，16²分辨率下，BRGM的LPIPS为0.24，RMSE为25.66，优于PULSE（0.29/27.14）和ESRGAN（0.35/29.32）。
• 在64²分辨率下，BRGM在FFHQ上的LPIPS为0.36，RMSE为16.07，RMSE优于SRFBN（0.23/9.40），但LPIPS略逊于SRFBN。
• 在图像修复任务中，BRGM在FFHQ上的LPIPS为0.19，RMSE为24.28，优于SN-PatchGAN（0.24/30.75）的两项指标。
• 在人工评估中，BRGM在16²分辨率下获得42%的投票认为其为最佳重构结果，显著优于PULSE（32%）和ESRGAN（11%）。
• 在MIMIC-III X光片数据集上，BRGM在16²分辨率下的LPIPS为0.18，RMSE为11.61，优于ESRGAN（0.32/14.67）和SRFBN（0.37/12.28）。
• 在脑部MRI扫描数据上，BRGM在16²分辨率下的LPIPS为0.12，RMSE为12.42，显著优于ESRGAN（0.34/22.81）和SRFBN（0.33/12.57）。