merge dygraph

doc/doc_en/inference_en.md

@@ -154,12 +154,12 @@ Set as `limit_type='min', det_limit_side_len=960`, it means that the shortest si
 
 If the resolution of the input picture is relatively large and you want to use a larger resolution prediction, you can set det_limit_side_len to the desired value, such as 1216:
 ```
-python3 tools/infer/predict_det.py --image_dir="./doc/imgs/22.jpg" --det_model_dir="./inference/det_db/" --det_limit_type=max --det_limit_side_len=1216
+python3 tools/infer/predict_det.py --image_dir="./doc/imgs/1.jpg" --det_model_dir="./inference/det_db/" --det_limit_type=max --det_limit_side_len=1216
 ```
 
 If you want to use the CPU for prediction, execute the command as follows
 ```
-python3 tools/infer/predict_det.py --image_dir="./doc/imgs/22.jpg" --det_model_dir="./inference/det_db/" --use_gpu=False
+python3 tools/infer/predict_det.py --image_dir="./doc/imgs/1.jpg" --det_model_dir="./inference/det_db/" --use_gpu=False
 ```
 
 <a name="DB_DETECTION"></a>
@@ -230,7 +230,7 @@ First, convert the model saved in the SAST text detection training process into
 python3 tools/export_model.py -c configs/det/det_r50_vd_sast_totaltext.yml -o Global.pretrained_model=./det_r50_vd_sast_totaltext_v2.0_train/best_accuracy  Global.save_inference_dir=./inference/det_sast_tt
 ```
 
-**For SAST curved text detection model inference, you need to set the parameter `--det_algorithm="SAST"` and `--det_sast_polygon=True`**, run the following command:
+For SAST curved text detection model inference, you need to set the parameter `--det_algorithm="SAST"` and `--det_sast_polygon=True`, run the following command:
 
 ```
 python3 tools/infer/predict_det.py --det_algorithm="SAST" --image_dir="./doc/imgs_en/img623.jpg" --det_model_dir="./inference/det_sast_tt/" --det_sast_polygon=True

doc/doc_en/recognition_en.md

@@ -329,6 +329,7 @@ There are two ways to create the required configuration file:：
     ...
 
     ```
+Italian is made up of Latin letters, so after executing the command, you will get the rec_latin_lite_train.yml.
 
 2. Manually modify the configuration file
 
@@ -375,7 +376,9 @@ Currently, the multi-language algorithms supported by PaddleOCR are:
 
 For more supported languages, please refer to : [Multi-language model](https://github.com/PaddlePaddle/PaddleOCR/blob/release/2.1/doc/doc_en/multi_languages_en.md#4-support-languages-and-abbreviations)
 
-The multi-language model training method is the same as the Chinese model. The training data set is 100w synthetic data. A small amount of fonts and test data can be downloaded on [Baidu Netdisk](https://pan.baidu.com/s/1bS_u207Rm7YbY33wOECKDA),Extraction code:frgi.
+The multi-language model training method is the same as the Chinese model. The training data set is 100w synthetic data. A small amount of fonts and test data can be downloaded using the following two methods.
+* [Baidu Netdisk](https://pan.baidu.com/s/1bS_u207Rm7YbY33wOECKDA),Extraction code:frgi.
+* [Google drive](https://drive.google.com/file/d/18cSWX7wXSy4G0tbKJ0d9PuIaiwRLHpjA/view)
 
 If you want to finetune on the basis of the existing model effect, please refer to the following instructions to modify the configuration file:
 

doc/doc_en/update_en.md

@@ -15,8 +15,6 @@
 - 2020.6.8 Add [datasets](./datasets_en.md) and keep updating
 - 2020.6.5 Support exporting `attention` model to `inference_model`
 - 2020.6.5 Support separate prediction and recognition, output result score
-- 2020.6.5 Support exporting `attention` model to `inference_model`
-- 2020.6.5 Support separate prediction and recognition, output result score
 - 2020.5.30 Provide Lightweight Chinese OCR online experience
 - 2020.5.30 Model prediction and training support on Windows system
 - 2020.5.30 Open source general Chinese OCR model

ppocr/data/simple_dataset.py

@@ -14,7 +14,6 @@
 import numpy as np
 import os
 import random
-import traceback
 from paddle.io import Dataset
 
 from .imaug import transform, create_operators
@@ -46,7 +45,6 @@ class SimpleDataSet(Dataset):
         self.seed = seed
         logger.info("Initialize indexs of datasets:%s" % label_file_list)
         self.data_lines = self.get_image_info_list(label_file_list, ratio_list)
-        self.check_data()
         self.data_idx_order_list = list(range(len(self.data_lines)))
         if self.mode == "train" and self.do_shuffle:
             self.shuffle_data_random()
@@ -103,18 +101,25 @@ class SimpleDataSet(Dataset):
 
     def __getitem__(self, idx):
         file_idx = self.data_idx_order_list[idx]
-        data = self.data_lines[file_idx]
+        data_line = self.data_lines[file_idx]
         try:
+            data_line = data_line.decode('utf-8')
+            substr = data_line.strip("\n").split(self.delimiter)
+            file_name = substr[0]
+            label = substr[1]
+            img_path = os.path.join(self.data_dir, file_name)
+            data = {'img_path': img_path, 'label': label}
+            if not os.path.exists(img_path):
+                raise Exception("{} does not exist!".format(img_path))
             with open(data['img_path'], 'rb') as f:
                 img = f.read()
                 data['image'] = img
             data['ext_data'] = self.get_ext_data()
             outs = transform(data, self.ops)
-        except:
-            error_meg = traceback.format_exc()
+        except Exception as e:
             self.logger.error(
-                "When parsing file {} and label {}, error happened with msg: {}".format(
-                    data['img_path'],data['label'], error_meg))
+                "When parsing line {}, error happened with msg: {}".format(
+                    data_line, e))
             outs = None
         if outs is None:
             # during evaluation, we should fix the idx to get same results for many times of evaluation.
@@ -125,17 +130,3 @@ class SimpleDataSet(Dataset):
 
     def __len__(self):
         return len(self.data_idx_order_list)
-
-    def check_data(self):
-        new_data_lines = []
-        for data_line in self.data_lines:
-            data_line = data_line.decode('utf-8')
-            substr = data_line.strip("\n").strip("\r").split(self.delimiter)
-            file_name = substr[0]
-            label = substr[1]
-            img_path = os.path.join(self.data_dir, file_name)
-            if os.path.exists(img_path):
-                new_data_lines.append({'img_path': img_path, 'label': label})
-            else:
-                self.logger.info("{} does not exist!".format(img_path))
-        self.data_lines = new_data_lines
\ No newline at end of file

ppocr/losses/basic_loss.py

@@ -54,6 +54,27 @@ class CELoss(nn.Layer):
         return loss
 
 
+class KLJSLoss(object):
+    def __init__(self, mode='kl'):
+        assert mode in ['kl', 'js', 'KL', 'JS'], "mode can only be one of ['kl', 'js', 'KL', 'JS']"
+        self.mode = mode
+
+    def __call__(self, p1, p2, reduction="mean"):
+
+        loss = paddle.multiply(p2, paddle.log( (p2+1e-5)/(p1+1e-5) + 1e-5))
+
+        if self.mode.lower() == "js":
+            loss += paddle.multiply(p1, paddle.log((p1+1e-5)/(p2+1e-5) + 1e-5))
+            loss *= 0.5
+        if reduction == "mean":
+            loss = paddle.mean(loss, axis=[1,2])
+        elif reduction=="none" or reduction is None:
+            return loss 
+        else:
+            loss = paddle.sum(loss, axis=[1,2])
+
+        return loss 
+
 class DMLLoss(nn.Layer):
     """
     DMLLoss
@@ -69,17 +90,21 @@ class DMLLoss(nn.Layer):
             self.act = nn.Sigmoid()
         else:
             self.act = None
+        
+        self.jskl_loss = KLJSLoss(mode="js")
 
     def forward(self, out1, out2):
         if self.act is not None:
             out1 = self.act(out1)
             out2 = self.act(out2)
-
-        log_out1 = paddle.log(out1)
-        log_out2 = paddle.log(out2)
-        loss = (F.kl_div(
-            log_out1, out2, reduction='batchmean') + F.kl_div(
-                log_out2, out1, reduction='batchmean')) / 2.0
+        if len(out1.shape) < 2:
+            log_out1 = paddle.log(out1)
+            log_out2 = paddle.log(out2)
+            loss = (F.kl_div(
+                log_out1, out2, reduction='batchmean') + F.kl_div(
+                    log_out2, out1, reduction='batchmean')) / 2.0
+        else:
+            loss = self.jskl_loss(out1, out2)
         return loss
 
 

ppocr/losses/combined_loss.py

@@ -17,7 +17,7 @@ import paddle.nn as nn
 
 from .distillation_loss import DistillationCTCLoss
 from .distillation_loss import DistillationDMLLoss
-from .distillation_loss import DistillationDistanceLoss
+from .distillation_loss import DistillationDistanceLoss, DistillationDBLoss, DistillationDilaDBLoss
 
 
 class CombinedLoss(nn.Layer):
@@ -44,15 +44,16 @@ class CombinedLoss(nn.Layer):
 
     def forward(self, input, batch, **kargs):
         loss_dict = {}
+        loss_all = 0.
         for idx, loss_func in enumerate(self.loss_func):
             loss = loss_func(input, batch, **kargs)
             if isinstance(loss, paddle.Tensor):
                 loss = {"loss_{}_{}".format(str(loss), idx): loss}
             weight = self.loss_weight[idx]
-            loss = {
-                "{}_{}".format(key, idx): loss[key] * weight
-                for key in loss
-            }
-            loss_dict.update(loss)
-        loss_dict["loss"] = paddle.add_n(list(loss_dict.values()))
+            for key in loss.keys():
+                if key == "loss":
+                    loss_all += loss[key] * weight
+                else:
+                    loss_dict["{}_{}".format(key, idx)] = loss[key]
+        loss_dict["loss"] = loss_all
         return loss_dict

ppocr/losses/distillation_loss.py

@@ -14,23 +14,76 @@
 
 import paddle
 import paddle.nn as nn
+import numpy as np
+import cv2
 
 from .rec_ctc_loss import CTCLoss
 from .basic_loss import DMLLoss
 from .basic_loss import DistanceLoss
+from .det_db_loss import DBLoss
+from .det_basic_loss import BalanceLoss, MaskL1Loss, DiceLoss
+
+
+def _sum_loss(loss_dict):
+    if "loss" in loss_dict.keys():
+        return loss_dict
+    else:
+        loss_dict["loss"] = 0.
+        for k, value in loss_dict.items():
+            if k == "loss":
+                continue
+            else:
+                loss_dict["loss"] += value
+        return loss_dict
 
 
 class DistillationDMLLoss(DMLLoss):
     """
     """
 
-    def __init__(self, model_name_pairs=[], act=None, key=None,
-                 name="loss_dml"):
+    def __init__(self,
+                 model_name_pairs=[],
+                 act=None,
+                 key=None,
+                 maps_name=None,
+                 name="dml"):
         super().__init__(act=act)
         assert isinstance(model_name_pairs, list)
         self.key = key
-        self.model_name_pairs = model_name_pairs
+        self.model_name_pairs = self._check_model_name_pairs(model_name_pairs)
         self.name = name
+        self.maps_name = self._check_maps_name(maps_name)
+    
+    def _check_model_name_pairs(self, model_name_pairs):
+        if not isinstance(model_name_pairs, list):
+            return []
+        elif isinstance(model_name_pairs[0], list) and isinstance(model_name_pairs[0][0], str):
+            return model_name_pairs
+        else:
+            return [model_name_pairs]
+
+    def _check_maps_name(self, maps_name):
+        if maps_name is None:
+            return None
+        elif type(maps_name) == str:
+            return [maps_name]
+        elif type(maps_name) == list:
+            return [maps_name]
+        else:
+            return None
+
+    def _slice_out(self, outs):
+        new_outs = {}
+        for k in self.maps_name:
+            if k == "thrink_maps":
+                new_outs[k] = outs[:, 0, :, :]
+            elif k == "threshold_maps":
+                new_outs[k] = outs[:, 1, :, :]
+            elif k == "binary_maps":
+                new_outs[k] = outs[:, 2, :, :]
+            else:
+                continue
+        return new_outs
 
     def forward(self, predicts, batch):
         loss_dict = dict()
@@ -40,13 +93,30 @@ class DistillationDMLLoss(DMLLoss):
             if self.key is not None:
                 out1 = out1[self.key]
                 out2 = out2[self.key]
-            loss = super().forward(out1, out2)
-            if isinstance(loss, dict):
-                for key in loss:
-                    loss_dict["{}_{}_{}_{}".format(key, pair[0], pair[1],
-                                                   idx)] = loss[key]
+
+            if self.maps_name is None:
+                loss = super().forward(out1, out2)
+                if isinstance(loss, dict):
+                    for key in loss:
+                        loss_dict["{}_{}_{}_{}".format(key, pair[0], pair[1],
+                                                       idx)] = loss[key]
+                else:
+                    loss_dict["{}_{}".format(self.name, idx)] = loss
             else:
-                loss_dict["{}_{}".format(self.name, idx)] = loss
+                outs1 = self._slice_out(out1)
+                outs2 = self._slice_out(out2)
+                for _c, k in enumerate(outs1.keys()):
+                    loss = super().forward(outs1[k], outs2[k])
+                    if isinstance(loss, dict):
+                        for key in loss:
+                            loss_dict["{}_{}_{}_{}_{}".format(key, pair[
+                                0], pair[1], map_name, idx)] = loss[key]
+                    else:
+                        loss_dict["{}_{}_{}".format(self.name, self.maps_name[_c],
+                                                    idx)] = loss
+        
+        loss_dict = _sum_loss(loss_dict)
+
         return loss_dict
 
 
@@ -73,6 +143,98 @@ class DistillationCTCLoss(CTCLoss):
         return loss_dict
 
 
+class DistillationDBLoss(DBLoss):
+    def __init__(self,
+                 model_name_list=[],
+                 balance_loss=True,
+                 main_loss_type='DiceLoss',
+                 alpha=5,
+                 beta=10,
+                 ohem_ratio=3,
+                 eps=1e-6,
+                 name="db",
+                 **kwargs):
+        super().__init__()
+        self.model_name_list = model_name_list
+        self.name = name
+        self.key = None
+
+    def forward(self, predicts, batch):
+        loss_dict = {}
+        for idx, model_name in enumerate(self.model_name_list):
+            out = predicts[model_name]
+            if self.key is not None:
+                out = out[self.key]
+            loss = super().forward(out, batch)
+
+            if isinstance(loss, dict):
+                for key in loss.keys():
+                    if key == "loss":
+                        continue
+                    name = "{}_{}_{}".format(self.name, model_name, key)
+                    loss_dict[name] = loss[key]
+            else:
+                loss_dict["{}_{}".format(self.name, model_name)] = loss
+
+        loss_dict = _sum_loss(loss_dict)
+        return loss_dict
+
+
+class DistillationDilaDBLoss(DBLoss):
+    def __init__(self,
+                 model_name_pairs=[],
+                 key=None,
+                 balance_loss=True,
+                 main_loss_type='DiceLoss',
+                 alpha=5,
+                 beta=10,
+                 ohem_ratio=3,
+                 eps=1e-6,
+                 name="dila_dbloss"):
+        super().__init__()
+        self.model_name_pairs = model_name_pairs
+        self.name = name
+        self.key = key
+
+    def forward(self, predicts, batch):
+        loss_dict = dict()
+        for idx, pair in enumerate(self.model_name_pairs):
+            stu_outs = predicts[pair[0]]
+            tch_outs = predicts[pair[1]]
+            if self.key is not None:
+                stu_preds = stu_outs[self.key]
+                tch_preds = tch_outs[self.key]
+
+            stu_shrink_maps = stu_preds[:, 0, :, :]
+            stu_binary_maps = stu_preds[:, 2, :, :]
+
+            # dilation to teacher prediction
+            dilation_w = np.array([[1, 1], [1, 1]])
+            th_shrink_maps = tch_preds[:, 0, :, :]
+            th_shrink_maps = th_shrink_maps.numpy() > 0.3  # thresh = 0.3 
+            dilate_maps = np.zeros_like(th_shrink_maps).astype(np.float32)
+            for i in range(th_shrink_maps.shape[0]):
+                dilate_maps[i] = cv2.dilate(
+                    th_shrink_maps[i, :, :].astype(np.uint8), dilation_w)
+            th_shrink_maps = paddle.to_tensor(dilate_maps)
+
+            label_threshold_map, label_threshold_mask, label_shrink_map, label_shrink_mask = batch[
+                1:]
+
+            # calculate the shrink map loss
+            bce_loss = self.alpha * self.bce_loss(
+                stu_shrink_maps, th_shrink_maps, label_shrink_mask)
+            loss_binary_maps = self.dice_loss(stu_binary_maps, th_shrink_maps,
+                                              label_shrink_mask)
+
+            # k = f"{self.name}_{pair[0]}_{pair[1]}"
+            k = "{}_{}_{}".format(self.name, pair[0], pair[1])
+            loss_dict[k] = bce_loss + loss_binary_maps
+
+        loss_dict = _sum_loss(loss_dict)
+        return loss_dict
+
+
 class DistillationDistanceLoss(DistanceLoss):
     """
     """

ppocr/metrics/det_metric.py

@@ -55,6 +55,7 @@ class DetMetric(object):
             result = self.evaluator.evaluate_image(gt_info_list, det_info_list)
             self.results.append(result)
 
+
     def get_metric(self):
         """
         return metrics {

ppocr/metrics/distillation_metric.py

@@ -24,8 +24,8 @@ from .cls_metric import ClsMetric
 class DistillationMetric(object):
     def __init__(self,
                  key=None,
-                 base_metric_name="RecMetric",
-                 main_indicator='acc',
+                 base_metric_name=None,
+                 main_indicator=None,
                  **kwargs):
         self.main_indicator = main_indicator
         self.key = key
@@ -42,16 +42,13 @@ class DistillationMetric(object):
                 main_indicator=self.main_indicator, **self.kwargs)
             self.metrics[key].reset()
 
-    def __call__(self, preds, *args, **kwargs):
+    def __call__(self, preds, batch, **kwargs):
         assert isinstance(preds, dict)
         if self.metrics is None:
             self._init_metrcis(preds)
         output = dict()
         for key in preds:
-            metric = self.metrics[key].__call__(preds[key], *args, **kwargs)
-            for sub_key in metric:
-                output["{}_{}".format(key, sub_key)] = metric[sub_key]
-        return output
+            self.metrics[key].__call__(preds[key], batch, **kwargs)
 
     def get_metric(self):
         """

ppocr/modeling/architectures/base_model.py

@@ -79,7 +79,10 @@ class BaseModel(nn.Layer):
             x = self.neck(x)
         y["neck_out"] = x
         x = self.head(x, targets=data)
-        y["head_out"] = x
+        if isinstance(x, dict):
+            y.update(x)
+        else:
+            y["head_out"] = x
         if self.return_all_feats:
             return y
         else:

ppocr/modeling/architectures/distillation_model.py

@@ -21,7 +21,7 @@ from ppocr.modeling.backbones import build_backbone
 from ppocr.modeling.necks import build_neck
 from ppocr.modeling.heads import build_head
 from .base_model import BaseModel
-from ppocr.utils.save_load import init_model
+from ppocr.utils.save_load import init_model, load_pretrained_params
 
 __all__ = ['DistillationModel']
 
@@ -46,7 +46,7 @@ class DistillationModel(nn.Layer):
                 pretrained = model_config.pop("pretrained")
             model = BaseModel(model_config)
             if pretrained is not None:
-                init_model(model, path=pretrained)
+                load_pretrained_params(model, pretrained)
             if freeze_params:
                 for param in model.parameters():
                     param.trainable = False

ppocr/postprocess/__init__.py

@@ -21,7 +21,7 @@ import copy
 
 __all__ = ['build_post_process']
 
-from .db_postprocess import DBPostProcess
+from .db_postprocess import DBPostProcess, DistillationDBPostProcess
 from .east_postprocess import EASTPostProcess
 from .sast_postprocess import SASTPostProcess
 from .rec_postprocess import CTCLabelDecode, AttnLabelDecode, SRNLabelDecode, DistillationCTCLabelDecode, \
@@ -33,9 +33,10 @@ from .pse_postprocess import PSEPostProcess
 
 def build_post_process(config, global_config=None):
     support_dict = [
-        'DBPostProcess', 'EASTPostProcess', 'SASTPostProcess', 'CTCLabelDecode',
+        'DBPostProcess','PSEPostProcess', 'EASTPostProcess', 'SASTPostProcess', 'CTCLabelDecode',
         'AttnLabelDecode', 'ClsPostProcess', 'SRNLabelDecode', 'PGPostProcess',
-        'DistillationCTCLabelDecode', 'TableLabelDecode', 'PSEPostProcess'
+        'DistillationCTCLabelDecode', 'TableLabelDecode',
+        'DistillationDBPostProcess'
     ]
 
     config = copy.deepcopy(config)

ppocr/postprocess/db_postprocess.py

@@ -187,3 +187,29 @@ class DBPostProcess(object):
 
             boxes_batch.append({'points': boxes})
         return boxes_batch
+
+
+class DistillationDBPostProcess(object):
+    def __init__(self, model_name=["student"],
+                 key=None,
+                 thresh=0.3,
+                 box_thresh=0.6,
+                 max_candidates=1000,
+                 unclip_ratio=1.5,
+                 use_dilation=False,
+                 score_mode="fast",
+                 **kwargs):
+        self.model_name = model_name
+        self.key = key
+        self.post_process = DBPostProcess(thresh=thresh,
+                 box_thresh=box_thresh,
+                 max_candidates=max_candidates,
+                 unclip_ratio=unclip_ratio,
+                 use_dilation=use_dilation,
+                 score_mode=score_mode)
+
+    def __call__(self, predicts, shape_list):
+        results = {}
+        for k in self.model_name:
+            results[k] = self.post_process(predicts[k], shape_list=shape_list)
+        return results

ppocr/utils/gen_label.py

-#copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
 #
-#Licensed under the Apache License, Version 2.0 (the "License");
-#you may not use this file except in compliance with the License.
-#You may obtain a copy of the License at
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
 #
 #    http://www.apache.org/licenses/LICENSE-2.0
 #
-#Unless required by applicable law or agreed to in writing, software
-#distributed under the License is distributed on an "AS IS" BASIS,
-#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#See the License for the specific language governing permissions and
-#limitations under the License.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import os
 import argparse
 import json
@@ -31,7 +31,9 @@ def gen_det_label(root_path, input_dir, out_label):
         for label_file in os.listdir(input_dir):
             img_path = root_path + label_file[3:-4] + ".jpg"
             label = []
-            with open(os.path.join(input_dir, label_file), 'r') as f:
+            with open(
+                    os.path.join(input_dir, label_file), 'r',
+                    encoding='utf-8-sig') as f:
                 for line in f.readlines():
                     tmp = line.strip("\n\r").replace("\xef\xbb\xbf",
                                                      "").split(',')

ppocr/utils/save_load.py

@@ -116,6 +116,27 @@ def load_dygraph_params(config, model, logger, optimizer):
         logger.info(f"loaded pretrained_model successful from {pm}")
         return {}
 
+def load_pretrained_params(model, path):
+    if path is None:
+        return False
+    if not os.path.exists(path) and not os.path.exists(path + ".pdparams"):
+        print(f"The pretrained_model {path} does not exists!")
+        return False
+
+    path = path if path.endswith('.pdparams') else path + '.pdparams'
+    params = paddle.load(path)
+    state_dict = model.state_dict()
+    new_state_dict = {}
+    for k1, k2 in zip(state_dict.keys(), params.keys()):
+        if list(state_dict[k1].shape) == list(params[k2].shape):
+            new_state_dict[k1] = params[k2]
+        else:
+            print(
+                f"The shape of model params {k1} {state_dict[k1].shape} not matched with loaded params {k2} {params[k2].shape} !"
+            )
+    model.set_state_dict(new_state_dict)
+    print(f"load pretrain successful from {path}")
+    return model
 
 def save_model(model,
                optimizer,

test/ocr_det_params.txt

 model_name:ocr_det
 python:python3.7
 gpu_list:0|0,1
-Global.auto_cast:False
+Global.auto_cast:null
 Global.epoch_num:10
 Global.save_model_dir:./output/
-Global.save_inference_dir:./output/
 Train.loader.batch_size_per_card:
-Global.use_gpu
-Global.pretrained_model
+Global.use_gpu:
+Global.pretrained_model:null
 
 trainer:norm|pact
 norm_train:tools/train.py -c configs/det/det_mv3_db.yml -o Global.pretrained_model=./pretrain_models/MobileNetV3_large_x0_5_pretrained
@@ -17,6 +16,8 @@ distill_train:null
 
 eval:tools/eval.py -c configs/det/det_mv3_db.yml -o 
 
+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
@@ -29,7 +30,6 @@ inference:tools/infer/predict_det.py
 --rec_batch_num:1
 --use_tensorrt:True|False
 --precision:fp32|fp16|int8
---det_model_dir
---image_dir
---save_log_path
-
+--det_model_dir:./inference/ch_ppocr_mobile_v2.0_det_infer/
+--image_dir:./inference/ch_det_data_50/all-sum-510/
+--save_log_path:./test/output/

test/ocr_rec_params.txt

+model_name:ocr_rec
+python:python
+gpu_list:0|0,1
+Global.auto_cast:null
+Global.epoch_num:10
+Global.save_model_dir:./output/
+Train.loader.batch_size_per_card:
+Global.use_gpu:
+Global.pretrained_model:null
+
+trainer:norm|pact
+norm_train:tools/train.py -c configs/rec/rec_mv3_none_bilstm_ctc.yml
+quant_train:deploy/slim/quantization/quant.py -c configs/rec/rec_mv3_none_bilstm_ctc.yml
+fpgm_train:null
+distill_train:null
+
+eval:tools/eval.py -c configs/rec/rec_mv3_none_bilstm_ctc.yml -o
+
+Global.save_inference_dir:./output/
+Global.pretrained_model:
+norm_export:tools/export_model.py -c configs/rec/rec_mv3_none_bilstm_ctc.yml -o
+quant_export:deploy/slim/quantization/export_model.py -c configs/rec/rec_mv3_none_bilstm_ctc.yml -o
+fpgm_export:null
+distill_export:null
+
+inference:tools/infer/predict_rec.py
+--use_gpu:True|False
+--enable_mkldnn:True|False
+--cpu_threads:1|6
+--rec_batch_num:1
+--use_tensorrt:True|False
+--precision:fp32|fp16|int8
+--rec_model_dir:./inference/ch_ppocr_mobile_v2.0_rec_infer/
+--image_dir:./inference/rec_inference
+--save_log_path:./test/output/
\ No newline at end of file

test/prepare.sh

@@ -26,20 +26,24 @@ IFS=$'\n'
 # The training params
 model_name=$(func_parser_value "${lines[0]}")
 train_model_list=$(func_parser_value "${lines[0]}")
+
 trainer_list=$(func_parser_value "${lines[10]}")
 
 # MODE be one of ['lite_train_infer' 'whole_infer' 'whole_train_infer']
 MODE=$2
-# prepare pretrained weights and dataset 
-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 ../
-
+# prepare pretrained weights and dataset
+if [ ${train_model_list[*]} = "ocr_det" ]; then
+  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 ../
+  fi
 if [ ${MODE} = "lite_train_infer" ];then
     # pretrain lite train data
     rm -rf ./train_data/icdar2015
     wget -nc -P ./train_data/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/icdar2015_lite.tar
-    cd ./train_data/ && tar xf icdar2015_lite.tar
+    wget -nc -P ./train_data/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/ic15_data.tar # todo change to bcebos
+
+    cd ./train_data/ && tar xf icdar2015_lite.tar && tar xf ic15_data.tar
     ln -s ./icdar2015_lite ./icdar2015
     cd ../
     epoch=10
@@ -47,13 +51,15 @@ if [ ${MODE} = "lite_train_infer" ];then
 elif [ ${MODE} = "whole_train_infer" ];then
     rm -rf ./train_data/icdar2015
     wget -nc -P ./train_data/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/icdar2015.tar
-    cd ./train_data/ && tar xf icdar2015.tar && cd ../
+    wget -nc -P ./train_data/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/ic15_data.tar
+    cd ./train_data/ && tar xf icdar2015.tar && tar xf ic15_data.tar && cd ../
     epoch=500
     eval_batch_step=200
 elif [ ${MODE} = "whole_infer" ];then
     rm -rf ./train_data/icdar2015
     wget -nc -P ./train_data/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/icdar2015_infer.tar
-    cd ./train_data/ && tar xf icdar2015_infer.tar
+    wget -nc -P ./train_data/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/ic15_data.tar
+    cd ./train_data/ && tar xf icdar2015_infer.tar && tar xf ic15_data.tar
     ln -s ./icdar2015_infer ./icdar2015
     cd ../
     epoch=10
@@ -62,8 +68,8 @@ else
     rm -rf ./train_data/icdar2015
     wget -nc -P ./train_data https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/ch_det_data_50.tar
     if [ ${model_name} = "ocr_det" ]; then
-        eval_model_name="ch_ppocr_mobile_v2.0_det_train"
-        wget -nc  -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_train.tar
+        eval_model_name="ch_ppocr_mobile_v2.0_det_infer"
+        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 && cd ../
     else 
         eval_model_name="ch_ppocr_mobile_v2.0_rec_train"
@@ -86,15 +92,17 @@ for train_model in ${train_model_list[*]}; do
     elif [ ${train_model} = "ocr_rec" ];then
         model_name="ocr_rec"
         yml_file="configs/rec/rec_mv3_none_bilstm_ctc.yml"
-        wget -nc -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/ch_rec_data_200.tar 
-        cd ./inference && tar xf ch_rec_data_200.tar  && cd ../
-        img_dir="./inference/ch_rec_data_200/"
+        wget -nc -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/test/rec_inference.tar
+        cd ./inference && tar xf rec_inference.tar  && cd ../
+        img_dir="./inference/rec_inference/"
+        data_dir=./inference/rec_inference
+        data_label_file=[./inference/rec_inference/rec_gt_test.txt]
     fi
 
     # eval 
     for slim_trainer in ${trainer_list[*]}; do 
         if [ ${slim_trainer} = "norm" ]; then
-            if [ ${model_name} = "ocr_det" ]; then
+            if [ ${model_name} = "det" ]; then
                 eval_model_name="ch_ppocr_mobile_v2.0_det_train"
                 wget -nc  -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_train.tar
                 cd ./inference && tar xf ${eval_model_name}.tar && cd ../
@@ -104,7 +112,7 @@ for train_model in ${train_model_list[*]}; do
                 cd ./inference && tar xf ${eval_model_name}.tar && cd ../
             fi 
         elif [ ${slim_trainer} = "pact" ]; then
-            if [ ${model_name} = "ocr_det" ]; then
+            if [ ${model_name} = "det" ]; then
                 eval_model_name="ch_ppocr_mobile_v2.0_det_quant_train"
                 wget -nc  -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/slim/ch_ppocr_mobile_v2.0_det_quant_train.tar
                 cd ./inference && tar xf ${eval_model_name}.tar && cd ../
@@ -114,7 +122,7 @@ for train_model in ${train_model_list[*]}; do
                 cd ./inference && tar xf ${eval_model_name}.tar && cd ../
             fi
         elif [ ${slim_trainer} = "distill" ]; then
-            if [ ${model_name} = "ocr_det" ]; then
+            if [ ${model_name} = "det" ]; then
                 eval_model_name="ch_ppocr_mobile_v2.0_det_distill_train"
                 wget -nc  -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/slim/ch_ppocr_mobile_v2.0_det_distill_train.tar
                 cd ./inference && tar xf ${eval_model_name}.tar && cd ../
@@ -124,7 +132,7 @@ for train_model in ${train_model_list[*]}; do
                 cd ./inference && tar xf ${eval_model_name}.tar && cd ../
             fi 
         elif [ ${slim_trainer} = "fpgm" ]; then
-            if [ ${model_name} = "ocr_det" ]; then
+            if [ ${model_name} = "det" ]; then
                 eval_model_name="ch_ppocr_mobile_v2.0_det_prune_train"
                 wget -nc  -P ./inference https://paddleocr.bj.bcebos.com/dygraph_v2.0/slim/ch_ppocr_mobile_v2.0_det_prune_train.tar
                 cd ./inference && tar xf ${eval_model_name}.tar && cd ../

test/test.sh

@@ -41,59 +41,51 @@ gpu_list=$(func_parser_value "${lines[2]}")
 autocast_list=$(func_parser_value "${lines[3]}")
 autocast_key=$(func_parser_key "${lines[3]}")
 epoch_key=$(func_parser_key "${lines[4]}")
+epoch_num=$(func_parser_value "${lines[4]}")
 save_model_key=$(func_parser_key "${lines[5]}")
-save_infer_key=$(func_parser_key "${lines[6]}")
-train_batch_key=$(func_parser_key "${lines[7]}")
-train_use_gpu_key=$(func_parser_key "${lines[8]}")
-pretrain_model_key=$(func_parser_key "${lines[9]}")
-
-trainer_list=$(func_parser_value "${lines[10]}")
-norm_trainer=$(func_parser_value "${lines[11]}")
-pact_trainer=$(func_parser_value "${lines[12]}")
-fpgm_trainer=$(func_parser_value "${lines[13]}")
-distill_trainer=$(func_parser_value "${lines[14]}")
-
-eval_py=$(func_parser_value "${lines[15]}")
-norm_export=$(func_parser_value "${lines[16]}")
-pact_export=$(func_parser_value "${lines[17]}")
-fpgm_export=$(func_parser_value "${lines[18]}")
-distill_export=$(func_parser_value "${lines[19]}")
-
-inference_py=$(func_parser_value "${lines[20]}")
-use_gpu_key=$(func_parser_key "${lines[21]}")
-use_gpu_list=$(func_parser_value "${lines[21]}")
-use_mkldnn_key=$(func_parser_key "${lines[22]}")
-use_mkldnn_list=$(func_parser_value "${lines[22]}")
-cpu_threads_key=$(func_parser_key "${lines[23]}")
-cpu_threads_list=$(func_parser_value "${lines[23]}")
-batch_size_key=$(func_parser_key "${lines[24]}")
-batch_size_list=$(func_parser_value "${lines[24]}")
-use_trt_key=$(func_parser_key "${lines[25]}")
-use_trt_list=$(func_parser_value "${lines[25]}")
-precision_key=$(func_parser_key "${lines[26]}")
-precision_list=$(func_parser_value "${lines[26]}")
-model_dir_key=$(func_parser_key "${lines[27]}")
-image_dir_key=$(func_parser_key "${lines[28]}")
-save_log_key=$(func_parser_key "${lines[29]}")
+train_batch_key=$(func_parser_key "${lines[6]}")
+train_use_gpu_key=$(func_parser_key "${lines[7]}")
+pretrain_model_key=$(func_parser_key "${lines[8]}")
+pretrain_model_value=$(func_parser_value "${lines[8]}")
+
+trainer_list=$(func_parser_value "${lines[9]}")
+norm_trainer=$(func_parser_value "${lines[10]}")
+pact_trainer=$(func_parser_value "${lines[11]}")
+fpgm_trainer=$(func_parser_value "${lines[12]}")
+distill_trainer=$(func_parser_value "${lines[13]}")
+
+eval_py=$(func_parser_value "${lines[14]}")
+
+save_infer_key=$(func_parser_key "${lines[15]}")
+export_weight=$(func_parser_key "${lines[16]}")
+norm_export=$(func_parser_value "${lines[17]}")
+pact_export=$(func_parser_value "${lines[18]}")
+fpgm_export=$(func_parser_value "${lines[19]}")
+distill_export=$(func_parser_value "${lines[20]}")
+
+inference_py=$(func_parser_value "${lines[21]}")
+use_gpu_key=$(func_parser_key "${lines[22]}")
+use_gpu_list=$(func_parser_value "${lines[22]}")
+use_mkldnn_key=$(func_parser_key "${lines[23]}")
+use_mkldnn_list=$(func_parser_value "${lines[23]}")
+cpu_threads_key=$(func_parser_key "${lines[24]}")
+cpu_threads_list=$(func_parser_value "${lines[24]}")
+batch_size_key=$(func_parser_key "${lines[25]}")
+batch_size_list=$(func_parser_value "${lines[25]}")
+use_trt_key=$(func_parser_key "${lines[26]}")
+use_trt_list=$(func_parser_value "${lines[26]}")
+precision_key=$(func_parser_key "${lines[27]}")
+precision_list=$(func_parser_value "${lines[27]}")
+infer_model_key=$(func_parser_key "${lines[28]}")
+infer_model=$(func_parser_value "${lines[28]}")
+image_dir_key=$(func_parser_key "${lines[29]}")
+infer_img_dir=$(func_parser_value "${lines[29]}")
+save_log_key=$(func_parser_key "${lines[30]}")
 
 LOG_PATH="./test/output"
 mkdir -p ${LOG_PATH}
 status_log="${LOG_PATH}/results.log"
 
-if [ ${MODE} = "lite_train_infer" ]; then
-    export infer_img_dir="./train_data/icdar2015/text_localization/ch4_test_images/"
-    export epoch_num=10
-elif [ ${MODE} = "whole_infer" ]; then
-    export infer_img_dir="./train_data/icdar2015/text_localization/ch4_test_images/"
-    export epoch_num=10
-elif [ ${MODE} = "whole_train_infer" ]; then
-    export infer_img_dir="./train_data/icdar2015/text_localization/ch4_test_images/"
-    export epoch_num=300
-else
-    export infer_img_dir="./inference/ch_det_data_50/all-sum-510"
-    export infer_model_dir="./inference/ch_ppocr_mobile_v2.0_det_train/best_accuracy"
-fi
-
 
 function func_inference(){
     IFS='|'
@@ -109,8 +101,8 @@ function func_inference(){
             for use_mkldnn in ${use_mkldnn_list[*]}; do
                 for threads in ${cpu_threads_list[*]}; do
                     for batch_size in ${batch_size_list[*]}; do
-                        _save_log_path="${_log_path}/infer_cpu_usemkldnn_${use_mkldnn}_threads_${threads}_batchsize_${batch_size}"
-                        command="${_python} ${_script} ${use_gpu_key}=${use_gpu} ${use_mkldnn_key}=${use_mkldnn} ${cpu_threads_key}=${threads} ${model_dir_key}=${_model_dir} ${batch_size_key}=${batch_size} ${image_dir_key}=${_img_dir}  ${save_log_key}=${_save_log_path} --benchmark=True"
+                        _save_log_path="${_log_path}/infer_cpu_usemkldnn_${use_mkldnn}_threads_${threads}_batchsize_${batch_size}.log"
+                        command="${_python} ${_script} ${use_gpu_key}=${use_gpu} ${use_mkldnn_key}=${use_mkldnn} ${cpu_threads_key}=${threads} ${infer_model_key}=${_model_dir} ${batch_size_key}=${batch_size} ${image_dir_key}=${_img_dir}  ${save_log_key}=${_save_log_path} --benchmark=True"
                         eval $command
                         status_check $? "${command}" "${status_log}"
                     done
@@ -123,8 +115,8 @@ function func_inference(){
                         continue
                     fi
                     for batch_size in ${batch_size_list[*]}; do
-                        _save_log_path="${_log_path}/infer_gpu_usetrt_${use_trt}_precision_${precision}_batchsize_${batch_size}"
-                        command="${_python} ${_script} ${use_gpu_key}=${use_gpu} ${use_trt_key}=${use_trt} ${precision_key}=${precision} ${model_dir_key}=${_model_dir} ${batch_size_key}=${batch_size} ${image_dir_key}=${_img_dir}  ${save_log_key}=${_save_log_path} --benchmark=True"
+                        _save_log_path="${_log_path}/infer_gpu_usetrt_${use_trt}_precision_${precision}_batchsize_${batch_size}.log"
+                        command="${_python} ${_script} ${use_gpu_key}=${use_gpu} ${use_trt_key}=${use_trt} ${precision_key}=${precision} ${infer_model_key}=${_model_dir} ${batch_size_key}=${batch_size} ${image_dir_key}=${_img_dir}  ${save_log_key}=${_save_log_path}  --benchmark=True"
                         eval $command
                         status_check $? "${command}" "${status_log}"
                     done
@@ -138,12 +130,13 @@ if [ ${MODE} != "infer" ]; then
 
 IFS="|"
 for gpu in ${gpu_list[*]}; do
-    train_use_gpu=True
+    use_gpu=True
     if [ ${gpu} = "-1" ];then
-        train_use_gpu=False
+        use_gpu=False
         env=""
     elif [ ${#gpu} -le 1 ];then
         env="export CUDA_VISIBLE_DEVICES=${gpu}"
+        eval ${env}
     elif [ ${#gpu} -le 15 ];then
         IFS=","
         array=(${gpu})
@@ -155,6 +148,7 @@ for gpu in ${gpu_list[*]}; do
         ips=${array[0]}
         gpu=${array[1]}
         IFS="|"
+        env=" "
     fi
     for autocast in ${autocast_list[*]}; do 
         for trainer in ${trainer_list[*]}; do 
@@ -179,13 +173,32 @@ for gpu in ${gpu_list[*]}; do
                 continue
             fi
 
-            save_log="${LOG_PATH}/${trainer}_gpus_${gpu}_autocast_${autocast}"
-            if [ ${#gpu} -le 2 ];then  # epoch_num #TODO
-                cmd="${python} ${run_train} ${train_use_gpu_key}=${train_use_gpu} ${autocast_key}=${autocast} ${epoch_key}=${epoch_num} ${save_model_key}=${save_log} "
-            elif [ ${#gpu} -le 15 ];then
-                cmd="${python} -m paddle.distributed.launch --gpus=${gpu} ${run_train} ${autocast_key}=${autocast} ${epoch_key}=${epoch_num}  ${save_model_key}=${save_log}"
+            # not set autocast when autocast is null
+            if [ ${autocast} = "null" ]; then
+                set_autocast=" "
+            else
+                set_autocast="${autocast_key}=${autocast}"
+            fi
+            # not set epoch when whole_train_infer
+            if [ ${MODE} != "whole_train_infer" ]; then
+                set_epoch="${epoch_key}=${epoch_num}"
+            else
+                set_epoch=" "
+            fi
+            # set pretrain
+            if [ ${pretrain_model_value} != "null" ]; then
+                set_pretrain="${pretrain_model_key}=${pretrain_model_value}"
             else
-                cmd="${python} -m paddle.distributed.launch --ips=${ips} --gpus=${gpu} ${run_train} ${autocast_key}=${autocast} ${epoch_key}=${epoch_num} ${save_model_key}=${save_log}"
+                set_pretrain=" "
+            fi
+
+            save_log="${LOG_PATH}/${trainer}_gpus_${gpu}_autocast_${autocast}"
+            if [ ${#gpu} -le 2 ];then  # train with cpu or single gpu
+                cmd="${python} ${run_train} ${train_use_gpu_key}=${use_gpu}  ${save_model_key}=${save_log} ${set_epoch} ${set_pretrain} ${set_autocast}"
+            elif [ ${#gpu} -le 15 ];then  # train with multi-gpu
+                cmd="${python} -m paddle.distributed.launch --gpus=${gpu} ${run_train} ${save_model_key}=${save_log}  ${set_epoch} ${set_pretrain} ${set_autocast}"
+            else     # train with multi-machine
+                cmd="${python} -m paddle.distributed.launch --ips=${ips} --gpus=${gpu} ${run_train} ${save_model_key}=${save_log} ${set_pretrain} ${set_epoch} ${set_autocast}"
             fi
             # run train
             eval $cmd
@@ -198,24 +211,27 @@ for gpu in ${gpu_list[*]}; do
 
             # run export model
             save_infer_path="${save_log}"
-            export_cmd="${python} ${run_export} ${save_model_key}=${save_log} ${pretrain_model_key}=${save_log}/latest ${save_infer_key}=${save_infer_path}"
+            export_cmd="${python} ${run_export} ${save_model_key}=${save_log} ${export_weight}=${save_log}/latest ${save_infer_key}=${save_infer_path}"
             eval $export_cmd
             status_check $? "${export_cmd}" "${status_log}"
 
             #run inference
+            eval $env
             save_infer_path="${save_log}"
             func_inference "${python}" "${inference_py}" "${save_infer_path}" "${LOG_PATH}" "${infer_img_dir}"
+            eval "unset CUDA_VISIBLE_DEVICES"
         done
     done
 done
 
 else
-    save_infer_path="${LOG_PATH}/${MODE}"
-    run_export=${norm_export}
-    export_cmd="${python} ${run_export} ${save_model_key}=${save_infer_path} ${pretrain_model_key}=${infer_model_dir} ${save_infer_key}=${save_infer_path}"
-    eval $export_cmd
-    status_check $? "${export_cmd}" "${status_log}"
-
+    GPUID=$3
+    if [ ${#GPUID} -le 0 ];then
+        env=" "
+    else
+        env="export CUDA_VISIBLE_DEVICES=${GPUID}"
+    fi
+    echo $env
     #run inference
-    func_inference "${python}" "${inference_py}" "${save_infer_path}" "${LOG_PATH}" "${infer_img_dir}"
+    func_inference "${python}" "${inference_py}" "${infer_model}" "${LOG_PATH}" "${infer_img_dir}"
 fi