移动doc_dewarp

services/paddle_services/doc_dewarp/utils.py

@@ -0,0 +1,67 @@
+import numpy as np
+import paddle
+from paddle.nn import functional as F
+
+
+def to_tensor(img: np.ndarray):
+    """
+    Converts a numpy array image (HWC) to a Paddle tensor (NCHW).
+
+    Args:
+        img (numpy.ndarray): The input image as a numpy array.
+
+    Returns:
+        out (paddle.Tensor): The output tensor.
+    """
+    img = img[:, :, ::-1]
+    img = img.astype("float32") / 255.0
+    img = img.transpose(2, 0, 1)
+    out: paddle.Tensor = paddle.to_tensor(img)
+    out = paddle.unsqueeze(out, axis=0)
+
+    return out
+
+
+def to_image(x: paddle.Tensor):
+    """
+    Converts a Paddle tensor (NCHW) to a numpy array image (HWC).
+
+    Args:
+        x (paddle.Tensor): The input tensor.
+
+    Returns:
+        out (numpy.ndarray): The output image as a numpy array.
+    """
+    out: np.ndarray = x.squeeze().numpy()
+    out = out.transpose(1, 2, 0)
+    out = out * 255.0
+    out = out.astype("uint8")
+    out = out[:, :, ::-1]
+
+    return out
+
+
+def unwarp(img, bm, bm_data_format="NCHW"):
+    """
+    Unwarp an image using a flow field.
+
+    Args:
+        img (paddle.Tensor): The input image.
+        bm (paddle.Tensor): The flow field.
+
+    Returns:
+        out (paddle.Tensor): The output image.
+    """
+    _, _, h, w = img.shape
+
+    if bm_data_format == "NHWC":
+        bm = bm.transpose([0, 3, 1, 2])
+
+    # NCHW
+    bm = F.upsample(bm, size=(h, w), mode="bilinear", align_corners=True)
+    # NHWC
+    bm = bm.transpose([0, 2, 3, 1])
+    # NCHW
+    out = F.grid_sample(img, bm)
+
+    return out