Merge pull request #3851 from tink2123/upload_seed

Add seed for ocr_rec

requirements.txt

@@ -12,3 +12,4 @@ cython
 lxml
 premailer
 openpyxl
+fasttext==0.9.1
\ No newline at end of file

tools/eval.py

@@ -54,8 +54,7 @@ def main():
             config['Architecture']["Head"]['out_channels'] = char_num
 
     model = build_model(config['Architecture'])
-    use_srn = config['Architecture']['algorithm'] == "SRN"
-    use_sar = config['Architecture']['algorithm'] == "SAR"
+    extra_input = config['Architecture']['algorithm'] in ["SRN", "SAR"]
     if "model_type" in config['Architecture'].keys():
         model_type = config['Architecture']['model_type']
     else:
@@ -72,7 +71,7 @@ def main():
 
     # start eval
     metric = program.eval(model, valid_dataloader, post_process_class,
-                        eval_class, model_type, use_srn, use_sar)
+                          eval_class, model_type, extra_input)
     logger.info('metric eval ***************')
     for k, v in metric.items():
         logger.info('{}:{}'.format(k, v))

tools/program.py

@@ -186,12 +186,13 @@ def train(config,
     model.train()
 
     use_srn = config['Architecture']['algorithm'] == "SRN"
-    use_nrtr = config['Architecture']['algorithm'] == "NRTR"
-    use_sar = config['Architecture']['algorithm'] == 'SAR'
+    extra_input = config['Architecture'][
+        'algorithm'] in ["SRN", "NRTR", "SAR", "SEED"]
     try:
         model_type = config['Architecture']['model_type']
     except:
         model_type = None
+    algorithm = config['Architecture']['algorithm']
 
     if 'start_epoch' in best_model_dict:
         start_epoch = best_model_dict['start_epoch']
@@ -215,7 +216,7 @@ def train(config,
             images = batch[0]
             if use_srn:
                 model_average = True
-            if use_srn or model_type == 'table' or use_nrtr or use_sar:
+            if model_type == 'table' or extra_input:
                 preds = model(images, data=batch[1:])
             else:
                 preds = model(images)
@@ -279,8 +280,7 @@ def train(config,
                     post_process_class,
                     eval_class,
                     model_type,
-                    use_srn=use_srn,
-                    use_sar=use_sar)
+                    extra_input=extra_input)
                 cur_metric_str = 'cur metric, {}'.format(', '.join(
                     ['{}: {}'.format(k, v) for k, v in cur_metric.items()]))
                 logger.info(cur_metric_str)
@@ -352,8 +352,7 @@ def eval(model,
          post_process_class,
          eval_class,
          model_type=None,
-         use_srn=False,
-         use_sar=False):
+         extra_input=False):
     model.eval()
     with paddle.no_grad():
         total_frame = 0.0
@@ -366,7 +365,7 @@ def eval(model,
                 break
             images = batch[0]
             start = time.time()
-            if use_srn or model_type == 'table' or use_sar:
+            if model_type == 'table' or extra_input:
                 preds = model(images, data=batch[1:])
             else:
                 preds = model(images)
@@ -402,7 +401,8 @@ def preprocess(is_train=False):
     alg = config['Architecture']['algorithm']
     assert alg in [
         'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN',
-        'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE'
+        'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE',
+        'ASTER'
     ]
 
     device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu'