import cv2
import paddle

from . import DEWARP
from .utils import to_tensor, to_image


def dewarp_image(image):
    img = cv2.resize(image, (288, 288))
    x = to_tensor(img)
    y = to_tensor(image)
    bm = DEWARP(x)
    bm = paddle.nn.functional.interpolate(
        bm, y.shape[2:], mode="bilinear", align_corners=False
    )
    bm_nhwc = bm.transpose([0, 2, 3, 1])
    out = paddle.nn.functional.grid_sample(y, (bm_nhwc / 288 - 0.5) * 2)
    return to_image(out)