Awesome

Leveraging Text Localization for Scene Text Removal via Text-aware Masked Image Modeling

This is a pytorch implementation for paper TMIM

Installation

1.Requirements

Python==3.8.12
Pytorch==1.11.0
CUDA==11.3

conda create -n tmim python==3.8.12
conda activate tmim
pip install --upgrade pip
pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html 
pip install -r requirements.txt

2.Datasets

Create a "data" folder. Download text removal dataset (SCUT-Enstext) and text detection datasets(TextOCR, Total-Text, ICDAR2015, COCO-Text, MLT19, ArT, lsvt(fullly annotated), ReCTS).
Create the coco-style annotations for text detection datasets with the code in utils/prepare_dataset/ (or download them from here(data.zip).

The structure of the data folder is shown below.

data
├── text_det
│   ├── art
│   │   ├── train_images
│   │   └── annotation.json
│   ├── cocotext
│   │   ├── train2014
│   │   └── cocotext.v2.json
│   ├── ic15
│   │   ├── train_images
│   │   └── annotation.json 
│   ├── lsvt
│   │   ├── train_images
│   │   └── annotation.json 
│   ├── mlt19
│   │   ├── train_images
│   │   └── annotation.json 
│   ├── rects
│   │   ├── img
│   │   └── annotation.json 
│   ├── textocr
│   │   ├── train_images
│   │   ├── TextOCR_0.1_train.json 
│   │   └── TextOCR_0.1_val.json 
│   └── totaltext
│       ├── train_images
│       └── annotation.json
└── text_rmv
    └── SCUT-EnsText
        ├── train
        │   ├── all_images
        │   ├── all_labels
        │   └── mask
        └── test
            ├── all_images
            ├── all_labels
            └── mask

Models

Model	Method	PSNR	MSSIM	MSE	AGE	Download
Uformer-B	Pretrained	36.66	97.66	0.0637	1.70	uformer_b_tmim.pth
Uformer-B	Fintuned	37.42	97.70	0.0459	1.52	uformer_b_tmim_str.pth
PERT	Pretrained	34.51	96.63	0.1231	2.11	pert_tmim.pth
PERT	Fintuned	35.66	97.18	0.0729	1.76	pert_tmim_str.pth
EraseNet	Pretrained	34.25	97.03	0.1141	2.23	erasenet_tmim.pth
EraseNet	Fintuned	35.47	97.30	0.0765	1.95	erasenet_tmim_str.pth

Inference

Download the pretrained models and run the following command for inference.

python -m torch.distributed.launch --master_port 29501 --nproc_per_node=1 demo.py --cfg configs/uformer_b_str.py --resume path/to/uformer_b_tmim_str.pth --test-dir path/to/image/folder --visualize-dir path/to/result/folder

Training and Testing

Set the "snapshot_dir"(The location for saving the checkpoints) and "dataroot"(The location of the datasets) in configs/*.py
Erasenet and Pert require 4 1080ti GPUs. Uformer requires 8 1080ti GPUs

1.Pretraining

Run the following command to pretrain the model on text detection datasets.

python -m torch.distributed.launch --master_port 29501 --nproc_per_node=8 train.py --cfg configs/uformer_b_tmim.py --ckpt-name uformer_b_tmim --save-log

Run the following command to test the performance of the pretrained model.

python test.py --cfg configs/uformer_b_tmim.py --ckpt-name uformer_b_tmim/latest.pth --save-log --visualize

2.Finetuning

Run the following command to finetune the model on text removal datasets.

python -m torch.distributed.launch --master_port 29501} --nproc_per_node=8 train.py --cfg configs/uformer_b_str.py --ckpt-name uformer_b_tmim_str --save-log --resume 'ckpt/uformer_b_tmim/latest.pth'

Run the following command to test the performance of the finetuned model.

python test.py --cfg configs/uformer_b_str.py --ckpt-name uformer_b_tmim_str/latest.pth --save-log --visualize