68 lines
1.5 KiB
Python
68 lines
1.5 KiB
Python
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
|