MAE-FAR

Codes of Learning Prior Feature and Attention Enhanced Image Inpainting (ECCV2022)

Updates

Codes about MAE pre-training/inference
Codes about ACR
Pre-trained MAE weights
Release weights trained on Places2

Preparation

You can download irregular/coco masks from here. Of course, you can use your own masks with a txt index as link.

Then download models for perceptual loss from LaMa:

mkdir -p ade20k/ade20k-resnet50dilated-ppm_deepsup/
wget -P ade20k/ade20k-resnet50dilated-ppm_deepsup/ http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth

MAE for Inpainting

Pre-trained MAE for Inpainting

FFHQ: link

Places2: link

Pre-training MAE

python -m torch.distributed.launch --nproc_per_node=2 --use_env mae_pretrain.py \
    --data_path ${IMAGE_FILES_TXT} \
    --mask_path ${IRR_MASK_TXT} ${COCO_MASK_TXT} \
    --batch_size 256 \
    --model mae_vit_base_patch16 \
    --mask_ratio 0.75 \
    --epochs 200 \
    --warmup_epochs 10 \
    --blr 1.5e-4 --weight_decay 0.05 \
    --num_workers 16 \
    --output_dir ./ckpts/mae_wo_cls_wo_pixnorm \
    --log_dir ./ckpts/mae_wo_cls_wo_pixnorm

mask_path can also be set as one file with --mask_path ${YOUR_MASK_TXT}.

You can also set --finetune and --random_mask for different MAE pre-training settings (not recommended in inpainting). Details are discussed in the paper.

Simple Inference

See simple_test.ipynb.

ACR

TIPS: Now we recommend to use features from layer6 of the MAE instead of layer8 to enjoy superior performance.

Ensure you have downloaded pre-trained resnet50dilated from LaMa.

Training

If multiple gpus (>1) are used, codes will work in DDP automatically .

python train.py --config configs/config_FAR_places2.yml \
                --exp_name ${EXP_NAME} \
                --resume_mae ${MAE_PATH}

Finetuning for 512x512~256x256

python finetune.py --config configs/config_FAR_places2_finetune_512.yml \
                   --exp_name ${EXP_NAME} \
                   --pl_resume ${PL_MODEL_PATH} \
                   --dynamic_size # if you need dynamic size training from 256 to 512

Testing

Download weights from link

This model is re-trained by the new code.

CUDA_VISIBLE_DEVICES=0 python test.py \
  --resume ${PL_MODEL_PATH} \
  --config ${CONFIG_PATH} \
  --output_path ${OUTPUT_PATH} \
  --image_size ${TEST_IMG_SCALE} \
  --load_pl

Inference for specific image/mask paths. MASK_PATH and IMAGE_PATH should contain equivalent images with the same names.

CUDA_VISIBLE_DEVICES=0 python test_custom.py \
  --resume ${PL_MODEL_PATH} \
  --config ${CONFIG_PATH} \
  --image_path ${IMAGE_PATH} \
  --mask_path ${MASK_PATH} \
  --output_path ${OUTPUT_PATH} \
  --image_size ${TEST_IMG_SCALE} \
  --load_pl

Acknowledgments

Our codes are based on LaMa and MAE.

Cite

If you found our program helpful, please consider citing:

@inproceedings{cao2022learning,
      title={Learning Prior Feature and Attention Enhanced Image Inpainting}, 
      author={Cao, Chenjie and Dong, Qiaole and Fu, Yanwei},
      journal={{ECCV}},
      year={2022}
}

ewrfcas / MAE-FAR