diff --git a/imperative/python/megengine/functional/nn.py b/imperative/python/megengine/functional/nn.py index 421defd60e6621e9ce0d17ecc84a04b0c958361e..a51a2cd0a53d63fe4215b6bf89e9f3a5108580a7 100644 --- a/imperative/python/megengine/functional/nn.py +++ b/imperative/python/megengine/functional/nn.py @@ -57,6 +57,7 @@ __all__ = [ "max_pool2d", "one_hot", "prelu", + "remap", "softmax", "softplus", "svd", @@ -892,7 +893,7 @@ def warp_perspective( border_mode: str = "REPLICATE", border_val: float = 0.0, interp_mode: str = "LINEAR", -): +) -> Tensor: r""" Applies perspective transformation to batched 2D images. @@ -907,9 +908,12 @@ def warp_perspective( :param inp: input image. :param M: `(batch, 3, 3)` transformation matrix. :param dsize: `(h, w)` size of the output image. - :param border_mode: pixel extrapolation method. Default: "REPLICATE" + :param border_mode: pixel extrapolation method. + Default: "REPLICATE". Currently also support "CONSTANT", "REFLECT", + "REFLECT_101", "WRAP". :param border_val: value used in case of a constant border. Default: 0 - :param interp_mode: interpolation methods. Default: "LINEAR" + :param interp_mode: interpolation methods. + Default: "LINEAR". Currently only support "LINEAR" mode. :return: output tensor. Note: @@ -951,6 +955,62 @@ def warp_perspective( return result +def remap( + inp: Tensor, + map_xy: Tensor, + border_mode: str = "REPLICATE", + scalar: float = 0.0, + interp_mode: str = "LINEAR", +) -> Tensor: + r""" + Applies remap transformation to batched 2D images. + + The input images are transformed to the output images by the tensor map_xy. + The output's H and W are same as map_xy's H and W. + + :param inp: input image + :param map_xy: (batch, oh, ow, 2) transformation matrix + :param border_mode: pixel extrapolation method. + Default: "REPLICATE". Currently also support "CONSTANT", "REFLECT", + "REFLECT_101", "WRAP". + :param scalar: value used in case of a constant border. Default: 0 + :param interp_mode: interpolation methods. + Default: "LINEAR". Currently only support "LINEAR" mode. + :return: output tensor. + + Examples: + + .. testcode:: + + import numpy as np + from megengine import tensor + import megengine.functional as F + inp_shape = (1, 1, 4, 4) + inp = tensor(np.arange(16, dtype=np.float32).reshape(inp_shape)) + map_xy_shape = (1, 2, 2, 2) + map_xy = tensor(np.array([[[1., 0.],[0., 1.]], + [[0., 1.],[0., 1.]]], + dtype=np.float32).reshape(map_xy_shape)) + out = F.remap(inp, map_xy) + print(out.numpy()) + + Outputs: + + .. testoutput:: + + [[[[1. 4.] + [4. 4.]]]] + + """ + + op = builtin.Remap( + imode=interp_mode, border_type=border_mode, format="NCHW", scalar=scalar + ) + assert isinstance(inp, (Tensor, megbrain_graph.VarNode)), "inp must be Tensor type" + (result,) = apply(op, inp, map_xy) + return result + + def matmul( inp1: Tensor, inp2: Tensor, @@ -1534,7 +1594,7 @@ def nms( :param boxes: tensor of shape `(N, 4)`; the boxes to perform nms on; each box is expected to be in `(x1, y1, x2, y2)` format. :param iou_thresh: IoU threshold for overlapping. :param scores: tensor of shape `(N,)`, the score of boxes. - :param max_output: the maximum number of boxes to keep; it is optional if this operator is not traced + :param max_output: the maximum number of boxes to keep; it is optional if this operator is not traced otherwise it required to be specified; if it is not specified, all boxes are kept. :return: indices of the elements that have been kept by NMS. diff --git a/imperative/python/test/unit/functional/test_functional.py b/imperative/python/test/unit/functional/test_functional.py index acdf1a51cea7ed55254b5ef0983c0886a6752efe..09968d77e1e9b42b5ff82200f4d2a4d9e73cf032 100644 --- a/imperative/python/test/unit/functional/test_functional.py +++ b/imperative/python/test/unit/functional/test_functional.py @@ -308,6 +308,37 @@ def test_one_hot(): onehot_high_dimension() +def test_warp_perspective(): + inp_shape = (1, 1, 4, 4) + x = tensor(np.arange(16, dtype=np.float32).reshape(inp_shape)) + M_shape = (1, 3, 3) + # M defines a translation: dst(1, 1, h, w) = rst(1, 1, h+1, w+1) + M = tensor( + np.array( + [[1.0, 0.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32 + ).reshape(M_shape) + ) + outp = F.warp_perspective(x, M, (2, 2)) + np.testing.assert_equal( + outp.numpy(), np.array([[[[5.0, 6.0], [9.0, 10.0]]]], dtype=np.float32) + ) + + +def test_remap(): + inp_shape = (1, 1, 4, 4) + inp = tensor(np.arange(16, dtype=np.float32).reshape(inp_shape)) + map_xy_shape = (1, 2, 2, 2) + map_xy = tensor( + np.array( + [[[1.0, 0.0], [0.0, 1.0]], [[0.0, 1.0], [0.0, 1.0]]], dtype=np.float32 + ).reshape(map_xy_shape) + ) + outp = F.remap(inp, map_xy) + np.testing.assert_equal( + outp.numpy(), np.array([[[[1.0, 4.0], [4.0, 4.0]]]], dtype=np.float32) + ) + + def test_binary_cross_entropy(): data1_shape = (2, 2) label1_shape = (2, 2)