From e3a3e0cd3c0dba3aeaf00922efc091ec561de7eb Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 22 Jan 2021 14:22:39 +0800
Subject: [PATCH] feat(mge/functional): add warp_affine

GitOrigin-RevId: 2b333ccd12c28e4c5a6bfd8060c43f0c8f03482a
---
 imperative/python/megengine/functional/nn.py  | 40 +++++++++++++++++++
 .../test/unit/functional/test_functional.py   | 18 +++++++++
 imperative/src/impl/ops/warp_affine.cpp       | 32 +++++++++++++++
 src/core/include/megbrain/ir/ops.td           |  2 +
 4 files changed, 92 insertions(+)
 create mode 100644 imperative/src/impl/ops/warp_affine.cpp

diff --git a/imperative/python/megengine/functional/nn.py b/imperative/python/megengine/functional/nn.py
index 91e46348c..9c8886ddf 100644
--- a/imperative/python/megengine/functional/nn.py
+++ b/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,
diff --git a/imperative/python/test/unit/functional/test_functional.py b/imperative/python/test/unit/functional/test_functional.py
index 13223a80d..0def0e4cc 100644
--- a/imperative/python/test/unit/functional/test_functional.py
+++ b/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))
diff --git a/imperative/src/impl/ops/warp_affine.cpp b/imperative/src/impl/ops/warp_affine.cpp
new file mode 100644
index 000000000..53237d694
--- /dev/null
+++ b/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
diff --git a/src/core/include/megbrain/ir/ops.td b/src/core/include/megbrain/ir/ops.td
index 6e495e748..0f87960d1 100644
--- a/src/core/include/megbrain/ir/ops.td
+++ b/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]>;
-- 
GitLab