将去扭曲模型转为onnx格式

doc_dewarp/dewarp.py

@@ -1,18 +1,21 @@
 import cv2
+import numpy as np
 import paddle
 
-from . import DEWARP
+from . import DOC_TR
 from .utils import to_tensor, to_image
 
 
 def dewarp_image(image):
-    img = cv2.resize(image, (288, 288))
-    x = to_tensor(img)
+    img = cv2.resize(image, (288, 288)).astype(np.float32)
     y = to_tensor(image)
-    bm = DEWARP(x)
+
+    img = np.transpose(img, (2, 0, 1))
+    bm = DOC_TR.run(None, {"image": img[None,]})[0]
+    bm = paddle.to_tensor(bm)
     bm = paddle.nn.functional.interpolate(
         bm, y.shape[2:], mode="bilinear", align_corners=False
     )
-    bm_nhwc = bm.transpose([0, 2, 3, 1])
+    bm_nhwc = np.transpose(bm, (0, 2, 3, 1))
     out = paddle.nn.functional.grid_sample(y, (bm_nhwc / 288 - 0.5) * 2)
     return to_image(out)