From 2afa7b9336c7b3091fa0573d50802bc4a1ea5727 Mon Sep 17 00:00:00 2001
From: Ryan <44900829+DrRyanHuang@users.noreply.github.com>
Date: Tue, 20 Dec 2022 11:44:07 +0800
Subject: [PATCH]  [Dy2St] transforms.RandomRotation Support static mode
 (#49128)

---
 python/paddle/tests/test_transforms_static.py |  18 ++-
 .../vision/transforms/functional_tensor.py    | 107 +++++++++++++-----
 python/paddle/vision/transforms/transforms.py |   7 +-
 3 files changed, 101 insertions(+), 31 deletions(-)

diff --git a/python/paddle/tests/test_transforms_static.py b/python/paddle/tests/test_transforms_static.py
index 4a8018b688..756d0d76fd 100644
--- a/python/paddle/tests/test_transforms_static.py
+++ b/python/paddle/tests/test_transforms_static.py
@@ -30,7 +30,7 @@ class TestTransformUnitTestBase(unittest.TestCase):
         self.set_trans_api()
 
     def get_shape(self):
-        return (64, 64, 3)
+        return (3, 64, 64)
 
     def set_trans_api(self):
         self.api = transforms.Resize(size=16)
@@ -142,5 +142,21 @@ class TestRandomCrop_same(TestTransformUnitTestBase):
         self.api = transforms.RandomCrop(self.crop_size)
 
 
+class TestRandomRotation(TestTransformUnitTestBase):
+    def set_trans_api(self):
+        degree = np.random.uniform(-180, 180)
+        eps = 10e-5
+        degree_tuple = (degree - eps, degree + eps)
+        self.api = transforms.RandomRotation(degree_tuple)
+
+
+class TestRandomRotation_expand_True(TestTransformUnitTestBase):
+    def set_trans_api(self):
+        degree = np.random.uniform(-180, 180)
+        eps = 10e-5
+        degree_tuple = (degree - eps, degree + eps)
+        self.api = transforms.RandomRotation(degree_tuple, expand=True, fill=3)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/python/paddle/vision/transforms/functional_tensor.py b/python/paddle/vision/transforms/functional_tensor.py
index 8137a4f284..b321f62f6a 100644
--- a/python/paddle/vision/transforms/functional_tensor.py
+++ b/python/paddle/vision/transforms/functional_tensor.py
@@ -15,6 +15,8 @@
 import math
 import numbers
 
+import numpy as np
+
 import paddle
 import paddle.nn.functional as F
 
@@ -221,12 +223,23 @@ def _affine_grid(theta, w, h, ow, oh):
 
     x_grid = paddle.linspace(-ow * 0.5 + d, ow * 0.5 + d - 1, ow)
     base_grid[..., 0] = x_grid
-    y_grid = paddle.linspace(-oh * 0.5 + d, oh * 0.5 + d - 1, oh).unsqueeze_(-1)
-    base_grid[..., 1] = y_grid
 
-    scaled_theta = theta.transpose((0, 2, 1)) / paddle.to_tensor(
-        [0.5 * w, 0.5 * h]
-    )
+    if paddle.in_dynamic_mode():
+        y_grid = paddle.linspace(
+            -oh * 0.5 + d, oh * 0.5 + d - 1, oh
+        ).unsqueeze_(-1)
+        base_grid[..., 1] = y_grid
+        tmp = paddle.to_tensor([0.5 * w, 0.5 * h])
+    else:
+        # To eliminate the warning:
+        # In static mode, unsqueeze_() is the same as unsqueeze() and does not perform inplace operation.
+        y_grid = paddle.linspace(-oh * 0.5 + d, oh * 0.5 + d - 1, oh).unsqueeze(
+            -1
+        )
+        base_grid[..., 1] = y_grid
+        tmp = paddle.assign(np.array([0.5 * w, 0.5 * h], dtype="float32"))
+
+    scaled_theta = theta.transpose((0, 2, 1)) / tmp
     output_grid = base_grid.reshape((1, oh * ow, 3)).bmm(scaled_theta)
 
     return output_grid.reshape((1, oh, ow, 2))
@@ -252,11 +265,21 @@ def _grid_transform(img, grid, mode, fill):
     if fill is not None:
         mask = img[:, -1:, :, :]  # n 1 h w
         img = img[:, :-1, :, :]  # n c h w
-        mask = mask.expand_as(img)
+        mask = mask.tile([1, img.shape[1], 1, 1])
         len_fill = len(fill) if isinstance(fill, (tuple, list)) else 1
-        fill_img = (
-            paddle.to_tensor(fill).reshape((1, len_fill, 1, 1)).expand_as(img)
-        )
+
+        if paddle.in_dynamic_mode():
+            fill_img = (
+                paddle.to_tensor(fill)
+                .reshape((1, len_fill, 1, 1))
+                .astype(img.dtype)
+                .expand_as(img)
+            )
+        else:
+            fill = np.array(fill).reshape(len_fill).astype("float32")
+            fill_img = paddle.ones_like(img) * paddle.assign(fill).reshape(
+                [1, len_fill, 1, 1]
+            )
 
         if mode == 'nearest':
             mask = paddle.cast(mask < 0.5, img.dtype)
@@ -363,15 +386,29 @@ def rotate(
     else:
         rotn_center = [(p - s * 0.5) for p, s in zip(center, [w, h])]
 
-    angle = math.radians(angle)
-    matrix = [
-        math.cos(angle),
-        math.sin(angle),
-        0.0,
-        -math.sin(angle),
-        math.cos(angle),
-        0.0,
-    ]
+    if paddle.in_dynamic_mode():
+        angle = math.radians(angle)
+        matrix = [
+            math.cos(angle),
+            math.sin(angle),
+            0.0,
+            -math.sin(angle),
+            math.cos(angle),
+            0.0,
+        ]
+        matrix = paddle.to_tensor(matrix, place=img.place)
+    else:
+        angle = angle / 180 * math.pi
+        matrix = paddle.concat(
+            [
+                paddle.cos(angle),
+                paddle.sin(angle),
+                paddle.zeros([1]),
+                -paddle.sin(angle),
+                paddle.cos(angle),
+                paddle.zeros([1]),
+            ]
+        )
 
     matrix[2] += matrix[0] * (-rotn_center[0] - post_trans[0]) + matrix[1] * (
         -rotn_center[1] - post_trans[1]
@@ -383,20 +420,29 @@ def rotate(
     matrix[2] += rotn_center[0]
     matrix[5] += rotn_center[1]
 
-    matrix = paddle.to_tensor(matrix, place=img.place)
     matrix = matrix.reshape((1, 2, 3))
 
     if expand:
         # calculate output size
-        corners = paddle.to_tensor(
-            [
-                [-0.5 * w, -0.5 * h, 1.0],
-                [-0.5 * w, 0.5 * h, 1.0],
-                [0.5 * w, 0.5 * h, 1.0],
-                [0.5 * w, -0.5 * h, 1.0],
-            ],
-            place=matrix.place,
-        ).astype(matrix.dtype)
+        if paddle.in_dynamic_mode():
+            corners = paddle.to_tensor(
+                [
+                    [-0.5 * w, -0.5 * h, 1.0],
+                    [-0.5 * w, 0.5 * h, 1.0],
+                    [0.5 * w, 0.5 * h, 1.0],
+                    [0.5 * w, -0.5 * h, 1.0],
+                ],
+                place=matrix.place,
+            ).astype(matrix.dtype)
+        else:
+            corners = paddle.assign(
+                [
+                    [-0.5 * w, -0.5 * h, 1.0],
+                    [-0.5 * w, 0.5 * h, 1.0],
+                    [0.5 * w, 0.5 * h, 1.0],
+                    [0.5 * w, -0.5 * h, 1.0],
+                ],
+            ).astype(matrix.dtype)
 
         _pos = (
             corners.reshape((1, -1, 3))
@@ -410,7 +456,10 @@ def rotate(
         nw = npos[0][0]
         nh = npos[0][1]
 
-        ow, oh = int(nw.numpy()[0]), int(nh.numpy()[0])
+        if paddle.in_dynamic_mode():
+            ow, oh = int(nw.numpy()[0]), int(nh.numpy()[0])
+        else:
+            ow, oh = nw.astype("int32"), nh.astype("int32")
 
     else:
         ow, oh = w, h
diff --git a/python/paddle/vision/transforms/transforms.py b/python/paddle/vision/transforms/transforms.py
index 8c1554e00d..f68cff5e51 100644
--- a/python/paddle/vision/transforms/transforms.py
+++ b/python/paddle/vision/transforms/transforms.py
@@ -1520,7 +1520,12 @@ class RandomRotation(BaseTransform):
         self.fill = fill
 
     def _get_param(self, degrees):
-        angle = random.uniform(degrees[0], degrees[1])
+        if paddle.in_dynamic_mode():
+            angle = random.uniform(degrees[0], degrees[1])
+        else:
+            angle = paddle.uniform(
+                [1], dtype="float32", min=degrees[0], max=degrees[1]
+            )
 
         return angle
 
-- 
GitLab