update docs and fix conflic

ppstructure/table/README.md

@@ -41,7 +41,7 @@ wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_tab
 wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_structure_infer.tar && tar xf en_ppocr_mobile_v2.0_table_structure_infer.tar
 cd ..
 # run
-python3 table/predict_table.py --det_model_dir=inference/en_ppocr_mobile_v2.0_table_det_infer --rec_model_dir=inference/en_ppocr_mobile_v2.0_table_rec_infer --table_model_dir=inference/en_ppocr_mobile_v2.0_table_structure_infer --image_dir=../doc/table/table.jpg --rec_char_dict_path=../ppocr/utils/ppocr_keys_v1.txt --table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt --rec_char_type=ch --det_limit_side_len=736 --det_limit_type=min --output ../output/table
+python3 table/predict_table.py --det_model_dir=inference/en_ppocr_mobile_v2.0_table_det_infer --rec_model_dir=inference/en_ppocr_mobile_v2.0_table_rec_infer --table_model_dir=inference/en_ppocr_mobile_v2.0_table_structure_infer --image_dir=../doc/table/table.jpg --rec_char_dict_path=../ppocr/utils/dict/table_dict.txt --table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt --rec_char_type=EN --det_limit_side_len=736 --det_limit_type=min --output ../output/table
 ```
 Note: The above model is trained on the PubLayNet dataset and only supports English scanning scenarios. If you need to identify other scenarios, you need to train the model yourself and replace the three fields `det_model_dir`, `rec_model_dir`, `table_model_dir`.
 

ppstructure/table/README_ch.md

@@ -43,7 +43,7 @@ wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_tab
 wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_structure_infer.tar && tar xf en_ppocr_mobile_v2.0_table_structure_infer.tar
 cd ..
 # 执行预测
-python3 table/predict_table.py --det_model_dir=inference/en_ppocr_mobile_v2.0_table_det_infer --rec_model_dir=inference/en_ppocr_mobile_v2.0_table_rec_infer --table_model_dir=inference/en_ppocr_mobile_v2.0_table_structure_infer --image_dir=../doc/table/table.jpg --rec_char_dict_path=../ppocr/utils/ppocr_keys_v1.txt --table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt --rec_char_type=ch --det_limit_side_len=736 --det_limit_type=min --output ../output/table
+python3 table/predict_table.py --det_model_dir=inference/en_ppocr_mobile_v2.0_table_det_infer --rec_model_dir=inference/en_ppocr_mobile_v2.0_table_rec_infer --table_model_dir=inference/en_ppocr_mobile_v2.0_table_structure_infer --image_dir=../doc/table/table.jpg --rec_char_dict_path=../ppocr/utils/dict/table_dict.txt --table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt --rec_char_type=EN --det_limit_side_len=736 --det_limit_type=min --output ../output/table
 ```
 运行完成后，每张图片的excel表格会保存到output字段指定的目录下
 

requirements.txt

@@ -8,3 +8,6 @@ numpy
 visualdl
 python-Levenshtein
 opencv-contrib-python==4.4.0.46
+lxml
+premailer
+openpyxl
\ No newline at end of file

tests/ocr_det_params.txt

@@ -4,7 +4,7 @@ python:python3.7
 gpu_list:0|0,1
 Global.use_gpu:True|True
 Global.auto_cast:null
-Global.epoch_num:lite_train_infer=2|whole_train_infer=300
+Global.epoch_num:lite_train_infer=1|whole_train_infer=300
 Global.save_model_dir:./output/
 Train.loader.batch_size_per_card:lite_train_infer=2|whole_train_infer=4
 Global.pretrained_model:null
@@ -15,7 +15,7 @@ null:null
 trainer:norm_train|pact_train
 norm_train:tools/train.py -c configs/det/det_mv3_db.yml -o Global.pretrained_model=./pretrain_models/MobileNetV3_large_x0_5_pretrained
 pact_train:deploy/slim/quantization/quant.py -c configs/det/det_mv3_db.yml -o
-fpgm_train:null
+fpgm_train:deploy/slim/prune/sensitivity_anal.py -c configs/det/det_mv3_db.yml -o Global.pretrained_model=./pretrain_models/det_mv3_db_v2.0_train/best_accuracy
 distill_train:null
 null:null
 null:null
@@ -29,7 +29,7 @@ Global.save_inference_dir:./output/
 Global.pretrained_model:
 norm_export:tools/export_model.py -c configs/det/det_mv3_db.yml -o 
 quant_export:deploy/slim/quantization/export_model.py -c configs/det/det_mv3_db.yml -o 
-fpgm_export:deploy/slim/prune/export_prune_model.py
+fpgm_export:deploy/slim/prune/export_prune_model.py -c configs/det/det_mv3_db.yml -o 
 distill_export:null
 export1:null
 export2:null
@@ -49,4 +49,19 @@ inference:tools/infer/predict_det.py
 --save_log_path:null
 --benchmark:True
 null:null
+===========================cpp_infer_params===========================
+use_opencv:True
+infer_model:./inference/ch_ppocr_mobile_v2.0_det_infer/
+infer_quant:False
+inference:./deploy/cpp_infer/build/ppocr det
+--use_gpu:True|False
+--enable_mkldnn:True|False
+--cpu_threads:1|6
+--rec_batch_num:1
+--use_tensorrt:False|True
+--precision:fp32|fp16
+--det_model_dir:
+--image_dir:./inference/ch_det_data_50/all-sum-510/
+--save_log_path:null
+--benchmark:True
 

tests/ocr_det_server_params.txt

+===========================train_params===========================
+model_name:ocr_server_det
+python:python3.7
+gpu_list:0|0,1
+Global.use_gpu:True|True
+Global.auto_cast:null
+Global.epoch_num:lite_train_infer=2|whole_train_infer=300
+Global.save_model_dir:./output/
+Train.loader.batch_size_per_card:lite_train_infer=2|whole_train_infer=4
+Global.pretrained_model:null
+train_model_name:latest
+train_infer_img_dir:./train_data/icdar2015/text_localization/ch4_test_images/
+null:null
+##
+trainer:norm_train|pact_train
+norm_train:tools/train.py -c configs/det/det_r50_vd_db.yml -o Global.pretrained_model=""
+pact_train:null
+fpgm_train:null
+distill_train:null
+null:null
+null:null
+##
+===========================eval_params=========================== 
+eval:tools/eval.py -c configs/det/det_mv3_db.yml -o
+null:null
+##
+===========================infer_params===========================
+Global.save_inference_dir:./output/
+Global.pretrained_model:
+norm_export:tools/export_model.py -c configs/det/det_r50_vd_db.yml -o 
+quant_export:null 
+fpgm_export:null
+distill_export:null
+export1:null
+export2:null
+##
+infer_model:./inference/ch_ppocr_server_v2.0_det_infer/
+infer_export:null
+infer_quant:False
+inference:tools/infer/predict_det.py
+--use_gpu:True|False
+--enable_mkldnn:True|False
+--cpu_threads:1|6
+--rec_batch_num:1
+--use_tensorrt:False|True
+--precision:fp32|fp16|int8
+--det_model_dir:
+--image_dir:./inference/ch_det_data_50/all-sum-510/
+--save_log_path:null
+--benchmark:True
+null:null
+

tests/prepare.sh

 #!/bin/bash
 FILENAME=$1
-# MODE be one of ['lite_train_infer' 'whole_infer' 'whole_train_infer', 'infer']
+# MODE be one of ['lite_train_infer' 'whole_infer' 'whole_train_infer', 'infer', 'cpp_infer']
 MODE=$2
 
 dataline=$(cat ${FILENAME})
@@ -34,10 +34,13 @@ MODE=$2
 if [ ${MODE} = "lite_train_infer" ];then
     # pretrain lite train data
     wget -nc -P  ./pretrain_models/ https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV3_large_x0_5_pretrained.pdparams
+    wget -nc -P ./pretrain_models/  https://paddleocr.bj.bcebos.com/dygraph_v2.0/en/det_mv3_db_v2.0_train.tar
+    cd ./pretrain_models/ && tar xf det_mv3_db_v2.0_train.tar && cd ../
     rm -rf ./train_data/icdar2015
     rm -rf ./train_data/ic15_data
     wget -nc -P ./train_data/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/icdar2015_lite.tar
     wget -nc -P ./train_data/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/ic15_data.tar # todo change to bcebos
+    wget -nc -P ./deploy/slim/prune https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/sen.pickle
     
     cd ./train_data/ && tar xf icdar2015_lite.tar && tar xf ic15_data.tar
     ln -s ./icdar2015_lite ./icdar2015
@@ -58,13 +61,17 @@ elif [ ${MODE} = "whole_infer" ];then
     cd ./train_data/ && tar xf icdar2015_infer.tar && tar xf ic15_data.tar
     ln -s ./icdar2015_infer ./icdar2015
     cd ../
-else
+elif [ ${MODE} = "infer" ] || [ ${MODE} = "cpp_infer" ];then
     if [ ${model_name} = "ocr_det" ]; then
         eval_model_name="ch_ppocr_mobile_v2.0_det_infer"
         rm -rf ./train_data/icdar2015
-        wget -nc -P ./train_data https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/ch_det_data_50.tar
+        wget -nc -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/ch_det_data_50.tar
         wget -nc  -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar
         cd ./inference && tar xf ${eval_model_name}.tar && tar xf ch_det_data_50.tar && cd ../
+    elif [ ${model_name} = "ocr_server_det" ]; then
+        wget -nc -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_server_v2.0_det_infer.tar
+        wget -nc -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/ch_det_data_50.tar
+        cd ./inference && tar xf ch_ppocr_server_v2.0_det_infer.tar && tar xf ch_det_data_50.tar && cd ../
     else 
         rm -rf ./train_data/ic15_data
         eval_model_name="ch_ppocr_mobile_v2.0_rec_infer"
@@ -74,3 +81,72 @@ else
     fi 
 fi
 
+if [ ${MODE} = "cpp_infer" ];then
+    cd deploy/cpp_infer
+    use_opencv=$(func_parser_value "${lines[52]}")
+    if [ ${use_opencv} = "True" ]; then
+        echo "################### build opencv ###################"
+        rm -rf 3.4.7.tar.gz opencv-3.4.7/
+        wget https://github.com/opencv/opencv/archive/3.4.7.tar.gz
+        tar -xf 3.4.7.tar.gz
+
+        cd opencv-3.4.7/
+        install_path=$(pwd)/opencv-3.4.7/opencv3
+
+        rm -rf build
+        mkdir build
+        cd build
+
+        cmake .. \
+            -DCMAKE_INSTALL_PREFIX=${install_path} \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DBUILD_SHARED_LIBS=OFF \
+            -DWITH_IPP=OFF \
+            -DBUILD_IPP_IW=OFF \
+            -DWITH_LAPACK=OFF \
+            -DWITH_EIGEN=OFF \
+            -DCMAKE_INSTALL_LIBDIR=lib64 \
+            -DWITH_ZLIB=ON \
+            -DBUILD_ZLIB=ON \
+            -DWITH_JPEG=ON \
+            -DBUILD_JPEG=ON \
+            -DWITH_PNG=ON \
+            -DBUILD_PNG=ON \
+            -DWITH_TIFF=ON \
+            -DBUILD_TIFF=ON
+
+        make -j
+        make install
+        cd ../
+        echo "################### build opencv finished ###################"
+    fi
+
+
+    echo "################### build PaddleOCR demo ####################"
+    if [ ${use_opencv} = "True" ]; then
+        OPENCV_DIR=$(pwd)/opencv-3.4.7/opencv3/
+    else
+        OPENCV_DIR=''
+    fi
+    LIB_DIR=$(pwd)/Paddle/build/paddle_inference_install_dir/
+    CUDA_LIB_DIR=$(dirname `find /usr -name libcudart.so`)
+    CUDNN_LIB_DIR=$(dirname `find /usr -name libcudnn.so`)
+    
+    BUILD_DIR=build
+    rm -rf ${BUILD_DIR}
+    mkdir ${BUILD_DIR}
+    cd ${BUILD_DIR}
+    cmake .. \
+        -DPADDLE_LIB=${LIB_DIR} \
+        -DWITH_MKL=ON \
+        -DWITH_GPU=OFF \
+        -DWITH_STATIC_LIB=OFF \
+        -DWITH_TENSORRT=OFF \
+        -DOPENCV_DIR=${OPENCV_DIR} \
+        -DCUDNN_LIB=${CUDNN_LIB_DIR} \
+        -DCUDA_LIB=${CUDA_LIB_DIR} \
+        -DTENSORRT_DIR=${TENSORRT_DIR} \
+
+    make -j
+    echo "################### build PaddleOCR demo finished ###################"
+fi
\ No newline at end of file

tests/readme.md

+
+# 介绍
+
+test.sh和params.txt文件配合使用，完成OCR轻量检测和识别模型从训练到预测的流程测试。
+
+# 安装依赖
+- 安装PaddlePaddle >= 2.0
+- 安装PaddleOCR依赖
+    ```
+    pip3 install  -r ../requirements.txt
+    ```
+- 安装autolog
+    ```
+    git clone https://github.com/LDOUBLEV/AutoLog
+    cd AutoLog
+    pip3 install -r requirements.txt
+    python3 setup.py bdist_wheel
+    pip3 install ./dist/auto_log-1.0.0-py3-none-any.whl
+    cd ../
+    ```
+
+# 目录介绍
+
+```bash
+tests/
+├── ocr_det_params.txt   # 测试OCR检测模型的参数配置文件
+├── ocr_rec_params.txt   # 测试OCR识别模型的参数配置文件
+└── prepare.sh           # 完成test.sh运行所需要的数据和模型下载
+└── test.sh              # 根据
+```
+
+# 使用方法
+test.sh包含四种运行模式，每种模式的运行数据不同，分别用于测试速度和精度，分别是：
+- 模式1 lite_train_infer，使用少量数据训练，用于快速验证训练到预测的走通流程，不验证精度和速度；
+```
+bash test/prepare.sh ./tests/ocr_det_params.txt 'lite_train_infer'
+bash tests/test.sh ./tests/ocr_det_params.txt 'lite_train_infer'
+```
+- 模式2 whole_infer，使用少量数据训练，一定量数据预测，用于验证训练后的模型执行预测，预测速度是否合理；
+```
+bash tests/prepare.sh ./tests/ocr_det_params.txt 'whole_infer'
+bash tests/test.sh ./tests/ocr_det_params.txt 'whole_infer'
+```
+
+- 模式3 infer 不训练，全量数据预测，走通开源模型评估、动转静，检查inference model预测时间和精度;
+```
+bash tests/prepare.sh ./tests/ocr_det_params.txt 'infer'
+用法1:
+bash tests/test.sh ./tests/ocr_det_params.txt 'infer'
+用法2: 指定GPU卡预测，第三个传入参数为GPU卡号
+bash tests/test.sh ./tests/ocr_det_params.txt 'infer' '1'
+```
+
+模式4: whole_train_infer , CE： 全量数据训练，全量数据预测，验证模型训练精度，预测精度，预测速度
+```
+bash tests/prepare.sh ./tests/ocr_det_params.txt 'whole_train_infer'
+bash tests/test.sh ./tests/ocr_det_params.txt 'whole_train_infer'
+```

tests/test.sh

 #!/bin/bash
 FILENAME=$1
-# MODE be one of ['lite_train_infer' 'whole_infer' 'whole_train_infer', 'infer']
+# MODE be one of ['lite_train_infer' 'whole_infer' 'whole_train_infer', 'infer', 'cpp_infer']
 MODE=$2
 
 dataline=$(cat ${FILENAME})
@@ -145,6 +145,33 @@ benchmark_value=$(func_parser_value "${lines[49]}")
 infer_key1=$(func_parser_key "${lines[50]}")
 infer_value1=$(func_parser_value "${lines[50]}")
 
+if [ ${MODE} = "cpp_infer" ]; then
+    # parser cpp inference model 
+    cpp_infer_model_dir_list=$(func_parser_value "${lines[53]}")
+    cpp_infer_is_quant=$(func_parser_value "${lines[54]}")
+    # parser cpp inference 
+    inference_cmd=$(func_parser_value "${lines[55]}")
+    cpp_use_gpu_key=$(func_parser_key "${lines[56]}")
+    cpp_use_gpu_list=$(func_parser_value "${lines[56]}")
+    cpp_use_mkldnn_key=$(func_parser_key "${lines[57]}")
+    cpp_use_mkldnn_list=$(func_parser_value "${lines[57]}")
+    cpp_cpu_threads_key=$(func_parser_key "${lines[58]}")
+    cpp_cpu_threads_list=$(func_parser_value "${lines[58]}")
+    cpp_batch_size_key=$(func_parser_key "${lines[59]}")
+    cpp_batch_size_list=$(func_parser_value "${lines[59]}")
+    cpp_use_trt_key=$(func_parser_key "${lines[60]}")
+    cpp_use_trt_list=$(func_parser_value "${lines[60]}")
+    cpp_precision_key=$(func_parser_key "${lines[61]}")
+    cpp_precision_list=$(func_parser_value "${lines[61]}")
+    cpp_infer_model_key=$(func_parser_key "${lines[62]}")
+    cpp_image_dir_key=$(func_parser_key "${lines[63]}")
+    cpp_infer_img_dir=$(func_parser_value "${lines[63]}")
+    cpp_save_log_key=$(func_parser_key "${lines[64]}")
+    cpp_benchmark_key=$(func_parser_key "${lines[65]}")
+    cpp_benchmark_value=$(func_parser_value "${lines[65]}")
+fi
+
+
 LOG_PATH="./tests/output"
 mkdir -p ${LOG_PATH}
 status_log="${LOG_PATH}/results.log"
@@ -218,6 +245,71 @@ function func_inference(){
     done
 }
 
+function func_cpp_inference(){
+    IFS='|'
+    _script=$1
+    _model_dir=$2
+    _log_path=$3
+    _img_dir=$4
+    _flag_quant=$5
+    # inference 
+    for use_gpu in ${cpp_use_gpu_list[*]}; do
+        if [ ${use_gpu} = "False" ] || [ ${use_gpu} = "cpu" ]; then
+            for use_mkldnn in ${cpp_use_mkldnn_list[*]}; do
+                if [ ${use_mkldnn} = "False" ] && [ ${_flag_quant} = "True" ]; then
+                    continue
+                fi
+                for threads in ${cpp_cpu_threads_list[*]}; do
+                    for batch_size in ${cpp_batch_size_list[*]}; do
+                        _save_log_path="${_log_path}/cpp_infer_cpu_usemkldnn_${use_mkldnn}_threads_${threads}_batchsize_${batch_size}.log"
+                        set_infer_data=$(func_set_params "${cpp_image_dir_key}" "${_img_dir}")
+                        set_benchmark=$(func_set_params "${cpp_benchmark_key}" "${cpp_benchmark_value}")
+                        set_batchsize=$(func_set_params "${cpp_batch_size_key}" "${batch_size}")
+                        set_cpu_threads=$(func_set_params "${cpp_cpu_threads_key}" "${threads}")
+                        set_model_dir=$(func_set_params "${cpp_infer_model_key}" "${_model_dir}")
+                        command="${_script} ${cpp_use_gpu_key}=${use_gpu} ${cpp_use_mkldnn_key}=${use_mkldnn} ${set_cpu_threads} ${set_model_dir} ${set_batchsize} ${set_infer_data} ${set_benchmark} > ${_save_log_path} 2>&1 "
+                        eval $command
+                        last_status=${PIPESTATUS[0]}
+                        eval "cat ${_save_log_path}"
+                        status_check $last_status "${command}" "${status_log}"
+                    done
+                done
+            done
+        elif [ ${use_gpu} = "True" ] || [ ${use_gpu} = "gpu" ]; then
+            for use_trt in ${cpp_use_trt_list[*]}; do
+                for precision in ${cpp_precision_list[*]}; do
+                    if [[ ${_flag_quant} = "False" ]] && [[ ${precision} =~ "int8" ]]; then
+                        continue
+                    fi 
+                    if [[ ${precision} =~ "fp16" || ${precision} =~ "int8" ]] && [ ${use_trt} = "False" ]; then
+                        continue
+                    fi
+                    if [[ ${use_trt} = "False" || ${precision} =~ "int8" ]] && [ ${_flag_quant} = "True" ]; then
+                        continue
+                    fi
+                    for batch_size in ${cpp_batch_size_list[*]}; do
+                        _save_log_path="${_log_path}/cpp_infer_gpu_usetrt_${use_trt}_precision_${precision}_batchsize_${batch_size}.log"
+                        set_infer_data=$(func_set_params "${cpp_image_dir_key}" "${_img_dir}")
+                        set_benchmark=$(func_set_params "${cpp_benchmark_key}" "${cpp_benchmark_value}")
+                        set_batchsize=$(func_set_params "${cpp_batch_size_key}" "${batch_size}")
+                        set_tensorrt=$(func_set_params "${cpp_use_trt_key}" "${use_trt}")
+                        set_precision=$(func_set_params "${cpp_precision_key}" "${precision}")
+                        set_model_dir=$(func_set_params "${cpp_infer_model_key}" "${_model_dir}")
+                        command="${_script} ${cpp_use_gpu_key}=${use_gpu} ${set_tensorrt} ${set_precision} ${set_model_dir} ${set_batchsize} ${set_infer_data} ${set_benchmark} > ${_save_log_path} 2>&1 "
+                        eval $command
+                        last_status=${PIPESTATUS[0]}
+                        eval "cat ${_save_log_path}"
+                        status_check $last_status "${command}" "${status_log}"
+                        
+                    done
+                done
+            done
+        else
+            echo "Does not support hardware other than CPU and GPU Currently!"
+        fi
+    done
+}
+
 if [ ${MODE} = "infer" ]; then
     GPUID=$3
     if [ ${#GPUID} -le 0 ];then
@@ -252,6 +344,25 @@ if [ ${MODE} = "infer" ]; then
         Count=$(($Count + 1))
     done
 
+elif [ ${MODE} = "cpp_infer" ]; then
+    GPUID=$3
+    if [ ${#GPUID} -le 0 ];then
+        env=" "
+    else
+        env="export CUDA_VISIBLE_DEVICES=${GPUID}"
+    fi
+    # set CUDA_VISIBLE_DEVICES
+    eval $env
+    export Count=0
+    IFS="|"
+    infer_quant_flag=(${cpp_infer_is_quant})
+    for infer_model in ${cpp_infer_model_dir_list[*]}; do
+        #run inference
+        is_quant=${infer_quant_flag[Count]}
+        func_cpp_inference "${inference_cmd}" "${infer_model}" "${LOG_PATH}" "${cpp_infer_img_dir}" ${is_quant}
+        Count=$(($Count + 1))
+    done
+
 else
     IFS="|"
     export Count=0

tools/infer/predict_det.py

@@ -101,6 +101,7 @@ class TextDetector(object):
         if args.benchmark:
             import auto_log
             pid = os.getpid()
+            gpu_id = utility.get_infer_gpuid()
             self.autolog = auto_log.AutoLogger(
                 model_name="det",
                 model_precision=args.precision,
@@ -110,7 +111,7 @@ class TextDetector(object):
                 inference_config=self.config,
                 pids=pid,
                 process_name=None,
-                gpu_ids=0,
+                gpu_ids=gpu_id if args.use_gpu else None,
                 time_keys=[
                     'preprocess_time', 'inference_time', 'postprocess_time'
                 ],

tools/infer/predict_e2e.py

@@ -74,7 +74,7 @@ class TextE2E(object):
 
         self.preprocess_op = create_operators(pre_process_list)
         self.postprocess_op = build_post_process(postprocess_params)
-        self.predictor, self.input_tensor, self.output_tensors = utility.create_predictor(
+        self.predictor, self.input_tensor, self.output_tensors, _ = utility.create_predictor(
             args, 'e2e', logger)  # paddle.jit.load(args.det_model_dir)
         # self.predictor.eval()
 

tools/infer/predict_rec.py

@@ -68,6 +68,7 @@ class TextRecognizer(object):
         if args.benchmark:
             import auto_log
             pid = os.getpid()
+            gpu_id = utility.get_infer_gpuid()
             self.autolog = auto_log.AutoLogger(
                 model_name="rec",
                 model_precision=args.precision,
@@ -77,7 +78,7 @@ class TextRecognizer(object):
                 inference_config=self.config,
                 pids=pid,
                 process_name=None,
-                gpu_ids=0 if args.use_gpu else None,
+                gpu_ids=gpu_id if args.use_gpu else None,
                 time_keys=[
                     'preprocess_time', 'inference_time', 'postprocess_time'
                 ],
@@ -87,7 +88,7 @@ class TextRecognizer(object):
     def resize_norm_img(self, img, max_wh_ratio):
         imgC, imgH, imgW = self.rec_image_shape
         assert imgC == img.shape[2]
-        if self.character_type == "ch":
+        max_wh_ratio = max(max_wh_ratio, imgW / imgH)
         imgW = int((32 * max_wh_ratio))
         h, w = img.shape[:2]
         ratio = w / float(h)
@@ -277,7 +278,7 @@ def main(args):
     if args.warmup:
         img = np.random.uniform(0, 255, [32, 320, 3]).astype(np.uint8)
         for i in range(2):
-            res = text_recognizer([img])
+            res = text_recognizer([img] * int(args.rec_batch_num))
 
     for image_file in image_file_list:
         img, flag = check_and_read_gif(image_file)

tools/infer/utility.py

@@ -159,6 +159,11 @@ def create_predictor(args, mode, logger):
         precision = inference.PrecisionType.Float32
 
     if args.use_gpu:
+        gpu_id = get_infer_gpuid()
+        if gpu_id is None:
+            raise ValueError(
+                "Not found GPU in current device. Please check your device or set args.use_gpu as False"
+            )
         config.enable_use_gpu(args.gpu_mem, 0)
         if args.use_tensorrt:
             config.enable_tensorrt_engine(
@@ -280,6 +285,20 @@ def create_predictor(args, mode, logger):
     return predictor, input_tensor, output_tensors, config
 
 
+def get_infer_gpuid():
+    cmd = "nvidia-smi"
+    res = os.popen(cmd).readlines()
+    if len(res) == 0:
+        return None
+    cmd = "env | grep CUDA_VISIBLE_DEVICES"
+    env_cuda = os.popen(cmd).readlines()
+    if len(env_cuda) == 0:
+        return 0
+    else:
+        gpu_id = env_cuda[0].strip().split("=")[1]
+        return int(gpu_id[0])
+
+
 def draw_e2e_res(dt_boxes, strs, img_path):
     src_im = cv2.imread(img_path)
     for box, str in zip(dt_boxes, strs):

tools/program.py

@@ -186,6 +186,7 @@ def train(config,
     model.train()
 
     use_srn = config['Architecture']['algorithm'] == "SRN"
+    use_nrtr = config['Architecture']['algorithm'] == "NRTR"
     use_sar = config['Architecture']['algorithm'] == 'SAR'
     try:
         model_type = config['Architecture']['model_type']
@@ -214,7 +215,7 @@ def train(config,
             images = batch[0]
             if use_srn:
                 model_average = True
-            if use_srn or model_type == 'table' or use_sar:
+            if use_srn or model_type == 'table' or use_nrtr or use_sar:
                 preds = model(images, data=batch[1:])
             else:
                 preds = model(images)
@@ -401,7 +402,11 @@ def preprocess(is_train=False):
     alg = config['Architecture']['algorithm']
     assert alg in [
         'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN',
+<<<<<<< HEAD
         'CLS', 'PGNet', 'Distillation', 'TableAttn', 'SAR'
+=======
+        'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn'
+>>>>>>> 63ed5fcab30801626ecf55a89f5dc9faf79a16d2
     ]
 
     device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu'