update infer/utility.py to support json format model

GreatV · Nov 15, 2024 · 7723cab · 7723cab
1 parent 0accd26
commit 7723cab
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/tools/infer/utility.py b/tools/infer/utility.py
@@ -234,12 +234,19 @@ def create_predictor(args, mode, logger):
         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):
+            model_json_file_path = "{}/{}.json".format(model_dir, file_name)
+            if (
+                os.path.exists(model_file_path) or os.path.exists(model_json_file_path)
+            ) and os.path.exists(params_file_path):
                 break
-        if not os.path.exists(model_file_path):
+        if not os.path.exists(model_file_path) and not os.path.exists(
+            model_json_file_path
+        ):
             raise ValueError(
                 "not find model.pdmodel or inference.pdmodel in {}".format(model_dir)
             )
+        if not os.path.exists(model_file_path) and os.path.exists(model_json_file_path):
+            model_file_path = model_json_file_path
         if not os.path.exists(params_file_path):
             raise ValueError(
                 "not find model.pdiparams or inference.pdiparams in {}".format(