feat(mge/functional): add warp_affine

GitOrigin-RevId: 2b333ccd12

imperative/python/megengine/functional/nn.py

@@ -62,6 +62,7 @@ __all__ = [
     "softmax",
     "softplus",
     "svd",
+    "warp_affine",
     "warp_perspective",
     "conv1d",
 ]
@@ -914,6 +915,45 @@ def resize(
     return result
 
 
+def warp_affine(
+    inp: Tensor,
+    weight: Tensor,
+    out_shape,
+    border_mode="REPLICATE",
+    border_val=0,
+    format="NHWC",
+    imode="LINEAR",
+):
+    """
+    Batched affine transform on 2D images.
+
+    :param inp: input image.
+    :param weight: weight tensor.
+    :param out_shape: output tensor shape.
+    :param border_mode: pixel extrapolation method.
+        Default: "WRAP". Currently "CONSTANT", "REFLECT",
+        "REFLECT_101", "ISOLATED", "WRAP", "REPLICATE", "TRANSPARENT" are supported.
+    :param border_val: value used in case of a constant border. Default: 0
+    :param format: "NHWC" as default based on historical concerns,
+        "NCHW" is also supported. Default: "NCHW".
+    :param imode: interpolation methods. Could be "LINEAR", "NEAREST", "CUBIC", "AREA".
+        Default: "LINEAR".
+    :return: output tensor.
+
+    .. note::
+
+    Here all available options for params are listed,
+    however it does not mean that you can use all the combinations.
+    On different platforms, different combinations are supported.
+    """
+    op = builtin.WarpAffine(
+        border_mode=border_mode, border_val=border_val, format=format, imode=imode
+    )
+    out_shape = utils.astensor1d(out_shape, inp, dtype="int32", device=inp.device)
+    (result,) = apply(op, inp, weight, out_shape)
+    return result
+
+
 def warp_perspective(
     inp: Tensor,
     M: Tensor,

imperative/python/test/unit/functional/test_functional.py

@@ -369,6 +369,24 @@ def test_warp_perspective():
     )
 
 
+def test_warp_affine():
+    inp_shape = (1, 3, 3, 3)
+    x = tensor(np.arange(27, dtype=np.float32).reshape(inp_shape))
+    weightv = [[[1.26666667, 0.6, -83.33333333], [-0.33333333, 1, 66.66666667]]]
+    outp = F.warp_affine(x, tensor(weightv), (2, 2), border_mode="WRAP")
+    res = np.array(
+        [
+            [
+                [[7.875, 8.875, 9.875], [8.90625, 9.90625, 10.90625]],
+                [[18.75, 19.75, 20.75], [14.90625, 15.90625, 16.90625]],
+            ]
+        ],
+        dtype=np.float32,
+    )
+    if not is_cuda_available():
+        np.testing.assert_almost_equal(outp.numpy(), res, 5)
+
+
 def test_remap():
     inp_shape = (1, 1, 4, 4)
     inp = tensor(np.arange(16, dtype=np.float32).reshape(inp_shape))

imperative/src/impl/ops/warp_affine.cpp

@@ -0,0 +1,32 @@
+/**
+ * \file imperative/src/impl/ops/warp_affine.cpp
+ * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
+ *
+ * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ */
+#include "../op_trait.h"
+#include "megbrain/imperative/ops/autogen.h"
+
+#include "megbrain/opr/imgproc.h"
+
+namespace mgb::imperative {
+
+namespace { namespace warp_affine {
+    auto apply_on_var_node(
+            const OpDef& def,
+            const VarNodeArray& inputs) {
+        mgb_assert(inputs.size() == 3);
+        auto&& op = static_cast<const WarpAffine&>(def);
+        return opr::WarpAffine::make(inputs[0], inputs[1], inputs[2], op.param());
+    }
+
+OP_TRAIT_REG(WarpAffine, WarpAffine)
+    .apply_on_var_node(apply_on_var_node)
+    .fallback();
+}} // warp_affine
+
+} // namespace mgb::imperative

src/core/include/megbrain/ir/ops.td

@@ -74,6 +74,8 @@ def ROIAlign: MgbHashableOp<"ROIAlign", [ROIAlignParam]>;
 
 def WarpPerspective: MgbHashableOp<"WarpPerspective", [WarpPerspectiveParam]>;
 
+def WarpAffine: MgbHashableOp<"WarpAffine", [WarpAffineParam]>;
+
 def Remap: MgbHashableOp<"Remap", [RemapParam]>;
 
 def Resize: MgbHashableOp<"Resize", [ResizeParam]>;