22 lines
592 B
Python
22 lines
592 B
Python
import cv2
|
|
import numpy as np
|
|
import paddle
|
|
|
|
from . import DOC_TR
|
|
from .utils import to_tensor, to_image
|
|
|
|
|
|
def dewarp_image(image):
|
|
img = cv2.resize(image, (288, 288)).astype(np.float32)
|
|
y = to_tensor(image)
|
|
|
|
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 = np.transpose(bm, (0, 2, 3, 1))
|
|
out = paddle.nn.functional.grid_sample(y, (bm_nhwc / 288 - 0.5) * 2)
|
|
return to_image(out)
|