feature/update-json - Updated json file support

Author: duccuong.le

paddleocr_convert/main.py

@@ -33,11 +33,11 @@ def __call__(
 
         model_path = self.get_file_path(model_path, save_dir)
 
-        model_dir = unzip_file(model_path, save_dir, is_del_raw=is_del_raw)
+        unzip_result = unzip_file(model_path, save_dir, is_del_raw=is_del_raw)
 
-        save_onnx_path = model_dir / f"{Path(model_path).stem}.onnx"
+        save_onnx_path = unzip_result.get('model_dir', '') / f"{Path(model_path).stem}.onnx"
         try:
-            self.convert_to_onnx(model_dir, save_onnx_path)
+            self.convert_to_onnx(unzip_result, save_onnx_path)
         except ConvertError as e:
             raise e
 
@@ -80,16 +80,20 @@ def get_file_path(file_path: str, save_dir: InputType) -> str:
             raise FileExistsError(f"{file_path} does not exist.")
         return file_path
 
-    def convert_to_onnx(self, model_dir: str, save_onnx_path: str) -> None:
+    def convert_to_onnx(self, unzip_result: object, save_onnx_path: str) -> None:
         """借助 :code:`paddle2onnx` 工具转换模型为onnx格式
 
         Args:
             model_dir (str): 保存paddle格式模型所在目录
             save_onnx_path (str): 保存的onnx全路径
         """
+        model_dir = unzip_result.get("model_dir", "")
+        my_files = unzip_result.get("my_files", {})
+        model_filename = 'inference.json' if "inference.json" in my_files else 'inference.pdmodel'
+
         shell_str = (
             f"paddle2onnx --model_dir {model_dir} "
-            "--model_filename inference.pdmodel "
+            f"--model_filename {model_filename} "
             "--params_filename inference.pdiparams "
             f"--opset_version {self.opset} "
             f"--save_file {save_onnx_path}"

paddleocr_convert/utils.py

@@ -75,14 +75,17 @@ def unzip_file(file_path: str, save_dir: InputType, is_del_raw: bool = True) ->
     model_dir = Path(save_dir) / Path(file_path).stem
     mkdir(model_dir)
 
-    tar_file_name_list = [".pdiparams", ".pdiparams.info", ".pdmodel"]
+    tar_file_name_list = [".pdiparams", ".pdiparams.info", ".pdmodel", ".json"]
+    my_files = set()
     with tarfile.open(file_path, "r") as tarObj:
+        
         for member in tarObj.getmembers():
             filename = None
 
             for tar_file_name in tar_file_name_list:
                 if member.name.endswith(tar_file_name):
                     filename = "inference" + tar_file_name
+                    my_files.add(filename)
 
             if filename is None:
                 continue
@@ -94,7 +97,7 @@ def unzip_file(file_path: str, save_dir: InputType, is_del_raw: bool = True) ->
     if is_del_raw:
         Path(file_path).unlink()
         print(f"The {file_path} has been deleted.")
-    return model_dir
+    return {"model_dir": model_dir, "my_files": my_files}
 
 
 def is_http_url(s: InputType) -> bool: