AOT-GAN用于高分辨率图像修复

Arxiv论文 |

AOT-GAN：用于高分辨率图像修复的聚合上下文转换<br> 曾艳红、付剑龙、巢宏杨和郭百宁。<br>

引用

如果我们的论文和代码对您的工作有帮助， 请慷慨地引用并为我们点赞 :kissing_heart: :kissing_heart: :kissing_heart: ！

@inproceedings{yan2021agg,
  author = {Zeng, Yanhong and Fu, Jianlong and Chao, Hongyang and Guo, Baining},
  title = {Aggregated Contextual Transformations for High-Resolution Image Inpainting},
  booktitle = {Arxiv},
  pages={-},
  year = {2020}
}

简介

尽管取得了一些有希望的结果，但现有的图像修复方法在填充高分辨率图像（如512x512）中的大面积缺失区域时仍然面临挑战。我们分析认为，这些困难主要来自于同时推断缺失内容和为极大的缺失区域合成精细纹理。 我们提出了一个基于GAN的模型，通过以下方式提高性能：

1. 通过生成器中的AOT块增强上下文推理。 AOT块聚合了不同感受野的上下文转换，允许捕捉信息丰富的远距离上下文和丰富的感兴趣模式进行上下文推理。
2. 通过判别器中的SoftGAN增强纹理合成。 我们通过定制的掩码预测任务改进了判别器的训练。经过增强的判别器被优化以区分真实和合成补丁的详细外观，从而反过来促使生成器合成更逼真的纹理。

结果

先决条件

• python 3.8.8
• pytorch (在Release 1.8.1上测试)

安装

克隆此仓库。

git clone git@github.com:researchmm/AOT-GAN-for-Inpainting.git
cd AOT-GAN-for-Inpainting/

对于完整的所需Python包集，我们建议从提供的YAML创建一个Conda环境，例如：

conda env create -f environment.yml
conda activate inpainting

数据集

1. 下载图像和掩码
2. 通过--dir_image和--dir_mask指定训练数据的路径。

入门指南

1. 训练：
   • 我们的代码基于Pytorch的分布式训练构建。
   • 运行

   cd src
   python train.py
   

2. 恢复训练：

   cd src
   python train.py --resume
   

3. 测试：

   cd src
   python test.py --pre_train [预训练模型路径]
   

4. 评估：

   cd src
   python eval.py --real_dir [真实图像] --fake_dir [修复结果] --metric mae psnr ssim fid
   

预训练模型

CELEBA-HQ | Places2

下载模型目录并将其放在experiments/下

演示

1. 下载预训练模型参数并将其放在experiments/下
2. 运行

cd src
python demo.py --dir_image [图像文件夹] --pre_train [预训练模型路径] --painter [bbox|freeform]

3. 按"+"或"-"控制画笔粗细。
4. 按"r"重置掩码；"k"保留现有修改；"s"保存结果。
5. 按空格键执行修复；"n"移至下一张图像；"Esc"退出演示。

TensorBoard

支持在TensorBoard上进行训练可视化。

运行tensorboard --logdir [日志文件夹] --bind_all并打开浏览器查看训练进度。

许可证

本项目基于Apache 2.0许可发布。请查看LICENSE文件了解更多信息。

致谢

我们要感谢edge-connect和EDSR_PyTorch。