🦙 LaMa: 基于傅里叶卷积的分辨率鲁棒大型掩码修复
作者:Roman Suvorov, Elizaveta Logacheva, Anton Mashikhin, Anastasia Remizova, Arsenii Ashukha, Aleksei Silvestrov, Naejin Kong, Harshith Goka, Kiwoong Park, Victor Lempitsky。
🔥🔥🔥
LaMa 在远超训练时分辨率(~2k❗️)的情况下表现出令人惊讶的泛化能力(训练时为256x256),即使在具有挑战性的场景中(如周期性结构的补全),也能实现卓越的性能。
[项目页面] [arXiv] [补充材料] [BibTeX] [Casual GAN Papers Summary]
LaMa 开发
(欢迎通过创建 Issue 分享您的论文)
- https://github.com/geekyutao/Inpaint-Anything --- Inpaint Anything: Segment Anything 与图像修复的结合
- 提高高分辨率图像修复的特征精炼 / 视频 / 代码 https://github.com/advimman/lama/pull/112 / 由 Geomagical Labs (geomagical.com)
非官方第三方应用:
(欢迎通过创建 Issue 分享您的应用/实现/演示)
- https://github.com/enesmsahin/simple-lama-inpainting - 一个用于 LaMa 修复的简单 pip 包。
- https://github.com/mallman/CoreMLaMa - 苹果 Core ML 模型格式。
- https://cleanup.pictures - 由 @cyrildiagne 创建的简单交互式对象移除工具。
- lama-cleaner 由 @Sanster 创建,是 https://cleanup.pictures 的自托管版本。
- 集成到 Huggingface Spaces 和 Gradio 中。参见演示: 由 @AK391 创建。
- Telegram 机器人 @MagicEraserBot 由 @Moldoteck 创建,代码。
- Auto-LaMa = DE:TR 对象检测 + LaMa 修复,由 @andy971022 创建。
- LAMA-Magic-Eraser-Local = 一个基于 PyQt5 构建的独立修复应用程序,由 @zhaoyun0071 创建。
- Hama - 使用智能刷子简化掩码绘制的对象移除工具。
- ModelScope = 由 @chenbinghui1 创建的中文最大模型社区。
- LaMa with MaskDINO = MaskDINO 对象检测 + LaMa 修复与优化,由 @qwopqwop200 创建。
- CoreMLaMa - 一个将 LaMa Cleaner 的 LaMa 端口转换为苹果 Core ML 模型格式的脚本。
环境配置
克隆仓库:
git clone https://github.com/advimman/lama.git
有三种环境选项:
-
Python virtualenv:
virtualenv inpenv --python=/usr/bin/python3 source inpenv/bin/activate pip install torch==1.8.0 torchvision==0.9.0 cd lama pip install -r requirements.txt
-
Conda
% 安装 Linux 的 conda,其他操作系统请从 https://docs.conda.io/en/latest/miniconda.html 下载 miniconda wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh bash Miniconda3-latest-Linux-x86_64.sh -b -p $HOME/miniconda $HOME/miniconda/bin/conda init bash cd lama conda env create -f conda_env.yml conda activate lama conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch -y pip install pytorch-lightning==1.2.9
-
Docker: 不需要任何操作 🎉。
推理
运行
cd lama
export TORCH_HOME=$(pwd) && export PYTHONPATH=$(pwd)
1. 下载预训练模型
最佳模型(Places2,Places Challenge):
curl -LJO https://huggingface.co/smartywu/big-lama/resolve/main/big-lama.zip
unzip big-lama.zip
所有模型(Places & CelebA-HQ):
下载 [https://drive.google.com/drive/folders/1B2x7eQDgecTL0oh3LSIBDGj0fTxs6Ips?usp=drive_link]
unzip lama-models.zip
2. 准备图像和掩码
下载测试图像:
unzip LaMa_test_images.zip
或者准备您的数据:
1) 创建名为 `[images_name]_maskXXX[image_suffix]` 的掩码,将图像和掩码放在同一文件夹中。- 您可以使用 脚本 生成随机掩码。
- 检查文件格式:
image1_mask001.png image1.png image2_mask001.png image2.png
- 在
configs/prediction/default.yaml
中指定image_suffix
,例如.png
或.jpg
或_input.jpg
。
3. 预测
假设我们想在 512x512 验证数据集上评估并选择最佳模型,使用粗/细/中等掩码,并且您的图像具有 .jpg 扩展名:
python3 bin/gen_mask_dataset.py \
$(pwd)/configs/data_gen/random_<size>_512.yaml \ # 粗、细、中等
my_dataset/val_source/ \
my_dataset/val/random_<size>_512.yaml \# 粗、细、中等
--ext jpg
因此,掩码生成器将会:
- 调整验证图像大小并裁剪,并将它们保存为 .png
- 生成掩码
ls my_dataset/val/random_medium_512/
image1_crop000_mask000.png
image1_crop000.png
image2_crop000_mask000.png
image2_crop000.png
...
为 visual_test 文件夹生成粗、细、中等掩码:
python3 bin/gen_mask_dataset.py \
$(pwd)/configs/data_gen/random_<size>_512.yaml \ #粗、细、中等
my_dataset/visual_test_source/ \
my_dataset/visual_test/random_<size>_512/ \ #粗、细、中等
--ext jpg
ls my_dataset/visual_test/random_thick_512/
image1_crop000_mask000.png
image1_crop000.png
image2_crop000_mask000.png
image2_crop000.png
...
对 eval_source 图像文件夹执行相同的操作:
python3 bin/gen_mask_dataset.py \
$(pwd)/configs/data_gen/random_<size>_512.yaml \ #粗、细、中等
my_dataset/eval_source/ \
my_dataset/eval/random_<size>_512/ \ #粗、细、中等
--ext jpg
生成定位配置文件以定位这些文件夹:
touch my_dataset.yaml
echo "data_root_dir: $(pwd)/my_dataset/" >> my_dataset.yaml
echo "out_root_dir: $(pwd)/experiments/" >> my_dataset.yaml
echo "tb_dir: $(pwd)/tb_logs/" >> my_dataset.yaml
mv my_dataset.yaml ${PWD}/configs/training/location/
检查数据配置是否与 my_dataset 文件夹结构一致:
$ cat ${PWD}/configs/training/data/abl-04-256-mh-dist
...
train:
indir: ${location.data_root_dir}/train
...
val:
indir: ${location.data_root_dir}/val
img_suffix: .png
visual_test:
indir: ${location.data_root_dir}/visual_test
img_suffix: .png
运行训练:
python3 bin/train.py -cn lama-fourier location=my_dataset data.batch_size=10
评估:LaMa 训练过程根据 my_dataset/val/ 上的评分选择最好的几个模型。
要在之前未见过的 my_dataset/eval 数据集上评估其中一个最佳模型(例如 epoch=32),请对细、粗和中等掩码执行以下操作:
推理:
python3 bin/predict.py \
model.path=$(pwd)/experiments/<user>_<date:time>_lama-fourier_/ \
indir=$(pwd)/my_dataset/eval/random_<size>_512/ \
outdir=$(pwd)/inference/my_dataset/random_<size>_512 \
model.checkpoint=epoch32.ckpt
指标计算:
python3 bin/evaluate_predicts.py \
$(pwd)/configs/eval2_gpu.yaml \
$(pwd)/my_dataset/eval/random_<size>_512/ \
$(pwd)/inference/my_dataset/random_<size>_512 \
$(pwd)/inference/my_dataset/random_<size>_512_metrics.csv
或者在 Docker 中:
TODO: train
TODO: eval
提示
生成不同种类的掩码: 以下命令将执行一个生成随机掩码的脚本。
bash docker/1_generate_masks_from_raw_images.sh \
configs/data_gen/random_medium_512.yaml \
/directory_with_input_images \
/directory_where_to_store_images_and_masks \
--ext png
测试数据生成命令将以适合预测的格式存储图像。
下表描述了我们用于生成论文中不同测试集的配置。
请注意,我们没有固定随机种子,因此每次的结果会略有不同。
Places 512x512 | CelebA 256x256 | |
---|---|---|
窄 | random_thin_512.yaml | random_thin_256.yaml |
中等 | random_medium_512.yaml | random_medium_256.yaml |
宽 | random_thick_512.yaml | random_thick_256.yaml |
可以随意将配置路径(参数#1)更改为 configs/data_gen
中的任何其他配置,或调整配置文件本身。
覆盖配置中的参数: 您还可以像这样覆盖配置中的参数:
python3 bin/train.py -cn <config> data.batch_size=10 run_title=my-title
其中 .yaml 文件扩展名被省略。
模型选项: 论文中的模型配置名称(替换到训练命令中):
- big-lama
- big-lama-regular
- lama-fourier
- lama-regular
- lama_small_train_masks
它们位于 configs/training/folder 中。
链接:
- 所有数据(模型、测试图像等)https://disk.yandex.ru/d/AmdeG-bIjmvSug
- 论文中的测试图像 https://disk.yandex.ru/d/xKQJZeVRk5vLlQ
- 预训练模型 https://disk.yandex.ru/d/EgqaSnLohjuzAg
- 感知损失模型 https://disk.yandex.ru/d/ncVmQlmT_kTemQ
- 我们的训练日志可在 https://disk.yandex.ru/d/9Bt1wNSDS4jDkQ 获取。
训练时间与资源:
TODO
致谢
- 分割代码和模型来自 CSAILVision。
- LPIPS 指标来自 richzhang。
- SSIM 来自 Po-Hsun-Su。
- FID 来自 mseitzer。
引用
如果您觉得此代码有帮助,请考虑引用:
@article{suvorov2021resolution,
title={Resolution-robust Large Mask Inpainting with Fourier Convolutions},
author={Suvorov, Roman and Logacheva, Elizaveta and Mashikhin, Anton and Remizova, Anastasia and Ashukha, Arsenii and Silvestrov, Aleksei and Kong, Naejin and Goka, Harshith and Park, Kiwoong and Lempitsky, Victor},
journal={arXiv preprint arXiv:2109.07161},
year={2021}
}