Permalink
Browse files

implement get_rotation_matrix2d

  • Loading branch information...
edgarriba committed Oct 14, 2018
1 parent cd0ad41 commit 876b2c601d4d5028ab84cc2735fa3fbb018b4c1b
Showing with 147 additions and 4 deletions.
  1. +56 −0 test/test_imgwarp.py
  2. +91 −4 torchgeometry/imgwarp.py
@@ -64,6 +64,62 @@ def test_warp_perspective_gradcheck(self):
raise_exception=True)
self.assertTrue(res)

def test_rotation_matrix2d(self):
# generate input data
batch_size = 1
center_base = torch.zeros(batch_size, 2)
angle_base = torch.ones(batch_size, 1)
scale_base = torch.ones(batch_size, 1)

# 90 deg rotation
center = center_base
angle = 90. * angle_base
scale = scale_base
M = tgm.get_rotation_matrix2d(center, angle, scale)

self.assertAlmostEqual(M[0, 0, 0].item(), 0.0)
self.assertAlmostEqual(M[0, 0, 1].item(), 1.0)
self.assertAlmostEqual(M[0, 1, 0].item(), -1.0)
self.assertAlmostEqual(M[0, 1, 1].item(), 0.0)

# 90 deg rotation + 2x scale
center = center_base
angle = 90. * angle_base
scale = 2. * scale_base
M = tgm.get_rotation_matrix2d(center, angle, scale)

self.assertAlmostEqual(M[0, 0, 0].item(), 0.0, 4)
self.assertAlmostEqual(M[0, 0, 1].item(), 2.0, 4)
self.assertAlmostEqual(M[0, 1, 0].item(), -2.0, 4)
self.assertAlmostEqual(M[0, 1, 1].item(), 0.0, 4)

# 45 deg rotation
center = center_base
angle = 45. * angle_base
scale = scale_base
M = tgm.get_rotation_matrix2d(center, angle, scale)

self.assertAlmostEqual(M[0, 0, 0].item(), 0.7071, 4)
self.assertAlmostEqual(M[0, 0, 1].item(), 0.7071, 4)
self.assertAlmostEqual(M[0, 1, 0].item(), -0.7071, 4)
self.assertAlmostEqual(M[0, 1, 1].item(), 0.7071, 4)

def test_get_rotation_matrix2d_gradcheck(self):
# generate input data
batch_size = 1
center = torch.zeros(batch_size, 2)
angle = torch.ones(batch_size, 1)
scale = torch.ones(batch_size, 1)

center = utils.tensor_to_gradcheck_var(center) # to var
angle = utils.tensor_to_gradcheck_var(angle) # to var
scale = utils.tensor_to_gradcheck_var(scale) # to var

# evaluate function gradient
res = gradcheck(tgm.get_rotation_matrix2d, (center, angle, scale),
raise_exception=True)
self.assertTrue(res)


if __name__ == '__main__':
unittest.main()
@@ -2,10 +2,12 @@

from .utils import inverse
from .homography_warper import homography_warp
from .conversions import deg2rad


__all__ = [
"warp_perspective",
"get_rotation_matrix2d",
]


@@ -47,16 +49,16 @@ def normalize_transform_to_pix(transform, height, width):

def warp_perspective(src, M, dsize, flags='bilinear', border_mode=None,
border_value=0):
"""Applies a perspective transformation to an image.
r"""Applies a perspective transformation to an image.
The function warpPerspective transforms the source image using
the specified matrix:
.. math::
dst(x, y) = src \left(
\\frac{M_{11} x + M_{12} y + M_{33}}{M_{31} x + M_{32} y + M_{33}} ,
\\frac{M_{21} x + M_{22} y + M_{23}}{M_{31} x + M_{32} y + M_{33}}
\\right)
\frac{M_{11} x + M_{12} y + M_{33}}{M_{31} x + M_{32} y + M_{33}} ,
\frac{M_{21} x + M_{22} y + M_{23}}{M_{31} x + M_{32} y + M_{33}}
\right)
Args:
src (Tensor): input image.
@@ -92,3 +94,88 @@ def warp_perspective(src, M, dsize, flags='bilinear', border_mode=None,
M_new = normalize_transform_to_pix(M_new, height, width)
# warp and return
return homography_warp(src, M_new, dsize)


def get_rotation_matrix2d(center, angle, scale):
r"""Calculates an affine matrix of 2D rotation.
The function calculates the following matrix:
.. math::
\begin{bmatrix}
\alpha & \beta & (1 - \alpha) \cdot \text{x}
- \beta \cdot \text{y} \\
-\beta & \alpha & \beta \cdot \text{x}
+ (1 - \alpha) \cdot \text{y}
\end{bmatrix}
where
.. math::
\alpha = \text{scale} \cdot cos(\text{angle}) \\
\beta = \text{scale} \cdot sin(\text{angle})
The transformation maps the rotation center to itself
If this is not the target, adjust the shift.
Args:
center (Tensor): center of the rotation in the source image.
angle (Tensor): rotation angle in degrees. Positive values mean
counter-clockwise rotation (the coordinate origin is assumed to
be the top-left corner).
scale (Tensor): isotropic scale factor.
Returns:
Tensor: the affine matrix of 2D rotation.
Shape:
- Input: :math:`(B, 2)`, :math:`(B, 1)` and :math:`(B, 1)`
- Output: :math:`(B, 2, 3)`
Example:
>>> center = torch.zeros(1, 2)
>>> scale = torch.ones(1, 1)
>>> angle = 45. * torch.ones(1, 1)
>>> M = tgm.get_rotation_matrix2d(center, angle, scale)
tensor([[[ 0.7071, 0.7071, 0.0000],
[-0.7071, 0.7071, 0.0000]]])
"""
if not torch.is_tensor(center):
raise TypeError("Input center type is not a torch.Tensor. Got {}"
.format(type(center)))
if not torch.is_tensor(angle):
raise TypeError("Input angle type is not a torch.Tensor. Got {}"
.format(type(angle)))
if not torch.is_tensor(scale):
raise TypeError("Input scale type is not a torch.Tensor. Got {}"
.format(type(scale)))
if not (len(center.shape) == 2 and center.shape[1] == 2):
raise ValueError("Input center must be a Bx2 tensor. Got {}"
.format(center.shape))
if not (len(angle.shape) == 2 and angle.shape[1] == 1):
raise ValueError("Input angle must be a Bx1 tensor. Got {}"
.format(angle.shape))
if not (len(scale.shape) == 2 and scale.shape[1] == 1):
raise ValueError("Input scale must be a Bx1 tensor. Got {}"
.format(scale.shape))
if not (center.shape[0] == angle.shape[0] == scale.shape[0]):
raise ValueError("Inputs must have same batch size dimension. Got {}"
.format(center.shape, angle.shape, scale.shape))
# convert angle and apply scale
angle_rad = deg2rad(angle)
alpha = torch.cos(angle_rad) * scale
beta = torch.sin(angle_rad) * scale

# unpack the center to x, y coordinates
x, y = torch.chunk(center, chunks=2, dim=1)

# create output tensor
batch_size, _ = center.shape
M = torch.zeros(batch_size, 2, 3, device=center.device, dtype=center.dtype)
M[..., 0, 0] = alpha
M[..., 0, 1] = beta
M[..., 0, 2] = (1. - alpha) * x - beta * y
M[..., 1, 0] = -beta
M[..., 1, 1] = alpha
M[..., 1, 2] = beta * x + (1. - alpha) * y
return M

0 comments on commit 876b2c6

Please sign in to comment.