From 1cc7ad34cae92f92ab63b75a59454ac220ef2b85 Mon Sep 17 00:00:00 2001
From: WenmuZhou <572459439@qq.com>
Date: Sat, 20 Aug 2022 07:46:42 +0000
Subject: [PATCH] update layout dict in whl

---
 paddleocr.py           |  3 ++-
 tools/infer/utility.py | 19 +++++++++++++------
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/paddleocr.py b/paddleocr.py
index 8e34c4fb..f6318ad2 100644
--- a/paddleocr.py
+++ b/paddleocr.py
@@ -289,7 +289,8 @@ MODEL_URLS = {
                 'ch': {
                     'url':
                     'https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_layout_infer.tar',
-                    'dict_path': 'ppocr/utils/dict/layout_publaynet_dict.txt'
+                    'dict_path':
+                    'ppocr/utils/dict/layout_dict/layout_publaynet_dict.txt'
                 }
             }
         }
diff --git a/tools/infer/utility.py b/tools/infer/utility.py
index 1eebc73f..1355ca62 100644
--- a/tools/infer/utility.py
+++ b/tools/infer/utility.py
@@ -181,14 +181,21 @@ def create_predictor(args, mode, logger):
         return sess, sess.get_inputs()[0], None, None
 
     else:
-        model_file_path = model_dir + "/inference.pdmodel"
-        params_file_path = model_dir + "/inference.pdiparams"
+        file_names = ['model', 'inference']
+        for file_name in file_names:
+            model_file_path = '{}/{}.pdmodel'.format(model_dir, file_name)
+            params_file_path = '{}/{}.pdiparams'.format(model_dir, file_name)
+            if os.path.exists(model_file_path) and os.path.exists(
+                    params_file_path):
+                break
         if not os.path.exists(model_file_path):
-            raise ValueError("not find model file path {}".format(
-                model_file_path))
+            raise ValueError(
+                "not find model.pdmodel or inference.pdmodel in {}".format(
+                    model_dir))
         if not os.path.exists(params_file_path):
-            raise ValueError("not find params file path {}".format(
-                params_file_path))
+            raise ValueError(
+                "not find model.pdiparams or inference.pdiparams in {}".format(
+                    model_dir))
 
         config = inference.Config(model_file_path, params_file_path)
 
-- 
GitLab