将去扭曲模型转为onnx格式

doc_dewarp/__init__.py

@@ -1,10 +1,4 @@
-import paddle
+from onnxruntime import InferenceSession
 
-from doc_dewarp.GeoTr import GeoTr
-
-model_path = "model/dewarp_model/best.ckpt"
-checkpoint = paddle.load(model_path)
-state_dict = checkpoint["model"]
-DEWARP = GeoTr()
-DEWARP.set_state_dict(state_dict)
-DEWARP.eval()
+DOC_TR = InferenceSession("model/dewarp_model/doc_tr_pp.onnx",
+                          providers=["CUDAExecutionProvider"], provider_options=[{"device_id": 0}])