/EAST

PyTorch Re-Implementation of EAST: An Efficient and Accurate Scene Text Detector

Primary LanguagePythonMIT LicenseMIT

Description

This is a PyTorch Re-Implementation of EAST: An Efficient and Accurate Scene Text Detector.

  • Only RBOX part is implemented.
  • Using dice loss instead of class-balanced cross-entropy loss. Some codes refer to argman/EAST and songdejia/EAST
  • The pre-trained model provided achieves 82.79 F-score on ICDAR 2015 Challenge 4 using only the 1000 images. see here for the detailed results.
Model Loss Recall Precision F-score
Original CE 72.75 80.46 76.41
Re-Implement Dice 81.27 84.36 82.79

Prerequisites

Only tested on

  • Anaconda3
  • Python 3.7.1
  • PyTorch 1.0.1
  • Shapely 1.6.4
  • opencv-python 4.0.0.21
  • lanms 1.0.2

When running the script, if some module is not installed you will see a notification and installation instructions. if you failed to install lanms, please update gcc and binutils. The update under conda environment is:

conda install -c omgarcia gcc-6
conda install -c conda-forge binutils

The original lanms code has a bug in normalize_poly that the ref vertices are not fixed when looping the p's ordering to calculate the minimum distance. We fixed this bug in LANMS so that anyone could compile the correct lanms. However, this repo still uses the original lanms.

Installation

1. Clone the repo

git clone https://github.com/SakuraRiven/EAST.git
cd EAST

2. Data & Pre-Trained Model

  • Download Train and Test Data: ICDAR 2015 Challenge 4. Cut the data into four parts: train_img, train_gt, test_img, test_gt.

  • Download pre-trained VGG16 from PyTorch: VGG16 and our trained EAST model: EAST. Make a new folder pths and put the download pths into pths

mkdir pths
mv east_vgg16.pth vgg16_bn-6c64b313.pth pths/

Here is an example:

.
├── EAST
│   ├── evaluate
│   └── pths
└── ICDAR_2015
    ├── test_gt
    ├── test_img
    ├── train_gt
    └── train_img

Train

Modify the parameters in train.py and run:

CUDA_VISIBLE_DEVICES=0,1 python train.py

Detect

Modify the parameters in detect.py and run:

CUDA_VISIBLE_DEVICES=0 python detect.py

Evaluate

  • The evaluation scripts are from ICDAR Offline evaluation and have been modified to run successfully with Python 3.7.1.
  • Change the evaluate/gt.zip if you test on other datasets.
  • Modify the parameters in eval.py and run:
CUDA_VISIBLE_DEVICES=0 python eval.py