Skip to content

Latest commit

 

History

History
 
 

text_detection

README.md

Text Detection in TensorFlow*

This repository contains inference and training code for PixelLink-like model networks. Models code is designed to enable export to a frozen graph and inference on CPU via OpenVINO™.

NOTE: Refer to the original implementation for details.

Trained models

Setup

Prerequisites

  • Ubuntu* 16.04
  • Python* 3.5.2
  • TensorFlow* 2.0
  • CUDA* 10.0

Installation

  1. Create virtual environment:

    virtualenv venv -p python3 --prompt="(td)"

  2. Activate virtual environment and set up OpenVINO™ variables:

    . venv/bin/activate

  3. Install the module:

    pip3 install -e .

Dataset

Sources

A toy dataset located in ./data. You can use it to do all steps including the following:

  • data preparation
  • model training
  • model evaluation

NOTE: This dataset is considerably small. It is highly recommended to use at least several thousand images datasets like ICDAR-* (ICDAR 2013, ICDAR 2015, ICDAR 2017, ...), COCO-TEXT, MSRA-TD500, and others.

Conversion

Use the annotation.py module to convert the datasets listed above to the internal format and create TFRecord that is suitable for training.

To convert a toy dataset located in the ./data folder,run the command:

python tools/prepare_annotation.py \
  --type toy \
  --images data \
  --out_annotation data/annotation.json

To create TFRecordDataset, run the command:

python tools/create_dataset.py \
  --input_datasets data/annotation.json \
  --output data/annotation.tfrecord

Training

To run training, run the following:

python tools/train.py \
  --learning_rate 0.001 \
  --train_dir model \
  --train_dataset data/annotation.tfrecord \
  --epochs_per_evaluation 100 \
  --test_dataset data/annotation.tfrecord \
  --model_type mobilenet_v2_ext \
  --config configs/config.yaml

Parameter description

  • train_dir - training directory where all snapshots and logs are stored
  • learning_rate - estimation of how fast the model weighs are updated. Too high value might result in divergence, while too low value might be a reason of slow convergence.
  • train_dataset - path to the TFRecord dataset that can be created using steps listed in the dataset section of this document. It is used during training.
  • batch_size - batch size.
  • test_dataset - path to the TFRecord dataset that can be created using steps listed in the dataset section of this document. It is used during validation to compute F1-score, precision, and recall.
  • epochs_per_evaluation - the value showing how often the model is saved/evaluated

Optional:

  • weights - weights of a pretrained model. Can increase convergence speed and result in a better model.
python tools/train.py \
  --learning_rate 0.0001 \
  --train_dir model \
  --train_dataset data/annotation.tfrecord \
  --epochs_per_evaluation 100 \
  --test_dataset data/annotation.tfrecord \
  --model_type mobilenet_v2_ext \
  --config configs/config.yaml \
  --weights init_weights/model_mobilenet_v2_ext/weights/model-523.save_weights

Evaluation

python tools/test.py \
  --config model/configuration.yaml \
  --dataset data/annotation.tfrecord \
  --weights model/weights/model-500.save_weights

Export models to OpenVINO™ (IR)

  1. Freeze your model:

    NOTE: Use the configuration file that appears in train_dir during training.

    python tools/export.py \
  --resolution 1280 768 \
  --config model/configuration.yaml \
  --weights model/weights/model-500.save_weights

    The command prints information about the frozen model and getting IR:

    Operations number: 51.075934092 GFlops

Output tensor names for using in InferenceEngine:
    model/link_logits_/add
    model/segm_logits/add
Run model_optimizer to get IR: mo.py --input_model model/weights/export/frozen_graph.pb --reverse_input_channels

  2. Run the Model Optimizer.

    NOTE You need to install TF1.13 to use the Model Optimizer.

    1. Create and activate new virtual environment:
    virtualenv venv_mo -p python3 --prompt="(td-mo)"
. venv_mo/bin/activate
    1. Install modules and activate environment for OpenVINO™ :
    pip3 install -r requirements-mo.txt
source /opt/intel/openvino/bin/setupvars.sh
    1. Run the Model Optimizer tool to export frozen graph to IR:
    mo.py --model_name text_detection \
  --input_model model/weights/export/frozen_graph.pb \
  --reverse_input_channels \
  --data_type FP32 \
  --input_shape="[1,768,1280,3]" \
  --output_dir IR

Demo in OpenVINO™

See https://github.com/opencv/open_model_zoo/tree/master/demos/text_detection_demo.