add Sim(2)

argoverse/utils/sim2.py

@@ -0,0 +1,95 @@
+"""
+Utility for 2d rigid body transformations with scaling.
+
+Refs:
+    http://ethaneade.com/lie_groups.pdf
+    https://github.com/borglab/gtsam/blob/develop/gtsam_unstable/geometry/Similarity3.h
+"""
+
+from typing import Union
+
+import numpy as np
+
+from argoverse.utils.helpers import assert_np_array_shape
+
+
+class Sim2:
+    """ Implements the Similarity(2) class."""
+
+    def __init__(self, R: np.ndarray, t: np.ndarray, s: Union[int, float]) -> None:
+        """Initialize from rotation R, translation t, and scale s.
+
+        Args:
+            R: array of shape (2x2) representing 2d rotation matrix
+            t: array of shape (2,) representing 2d translation
+            s: scaling factor
+        """
+        assert_np_array_shape(R, (2, 2))
+        assert_np_array_shape(t, (2,))
+        assert isinstance(s, float) or isinstance(s, int)
+        self.R_ = R.astype(np.float32)
+        self.t_ = t.astype(np.float32)
+        self.s_ = float(s)
+
+    def __eq__(self, other: object) -> bool:
+        """Check for equality with other Sim(2) object"""
+        if not np.isclose(self.scale(), other.scale()):
+            return False
+
+        if not np.allclose(self.rotation(), other.rotation()):
+            return False
+
+        if not np.allclose(self.translation(), other.translation()):
+            return False
+
+        return True
+
+    def rotation(self) -> np.ndarray:
+        """Return the 2x2 rotation matrix"""
+        return self.R_
+
+    def translation(self) -> np.ndarray:
+        """Return the (2,) translation vector"""
+        return self.t_
+
+    def scale(self) -> float:
+        """Return the scale."""
+        return self.s_
+
+    def matrix(self) -> np.ndarray:
+        """Calculate 3*3 matrix group equivalent"""
+        T = np.zeros((3, 3))
+        T[:2, :2] = self.R_
+        T[:2, 2] = self.t_
+        T[2, 2] = 1 / self.s_
+        return T
+
+    def compose(self, S: "Sim2") -> "Sim2":
+        """Composition with another Sim2."""
+        return Sim2(self.R_ * S.R_, ((1.0 / S.s_) * self.t_) + self.R_ @ S.t_, self.s_ * S.s_)
+
+    def inverse(self) -> "Sim2":
+        """Return the inverse."""
+        Rt = self.R_.T
+        sRt = -Rt @ (self.s_ * self.t_)
+        return Sim2(Rt, sRt, 1.0 / self.s_)
+
+    def transformFrom(self, point_cloud: np.ndarray) -> np.ndarray:
+        """Transform point cloud such that if they are in frame A,
+        and our Sim(3) transform is defines as bSa, then we get points
+        back in frame B:
+            p_b = bSa * p_a
+        Action on a point p is s*(R*p+t).
+
+        Args:
+            point_cloud: Nx2 array representing 2d points in frame A
+
+        Returns:
+            transformed_point_cloud: Nx2 array representing 2d points in frame B
+        """
+        assert_np_array_shape(point_cloud, (None, 2))
+        # (2,2) x (2,N) + (2,1) = (2,N) -> transpose
+        transformed_point_cloud = (self.R_ @ point_cloud.T + self.t_.reshape(2, 1)).T
+
+        # now scale points
+        return transformed_point_cloud * self.s_

tests/test_sim2.py

@@ -0,0 +1,136 @@
+import numpy as np
+
+from argoverse.utils.sim2 import Sim2
+
+
+def test_constructor() -> None:
+    """Sim(2) to perform p_b = bSa * p_a"""
+    bRa = np.eye(2)
+    bta = np.array([1, 2])
+    bsa = 3.0
+    bSa = Sim2(R=bRa, t=bta, s=bsa)
+    assert isinstance(bSa, Sim2)
+    assert np.allclose(bSa.R_, bRa)
+    assert np.allclose(bSa.t_, bta)
+    assert np.allclose(bSa.s_, bsa)
+
+
+def test_is_eq() -> None:
+    """ Ensure object equality works properly (are equal). """
+    bSa = Sim2(R=np.eye(2), t=np.array([1, 2]), s=3.0)
+    bSa_ = Sim2(R=np.eye(2), t=np.array([1.0, 2.0]), s=3)
+    assert bSa == bSa_
+
+
+def test_not_eq() -> None:
+    """ Ensure object equality works properly (not equal). """
+    bSa = Sim2(R=np.eye(2), t=np.array([2, 1]), s=3.0)
+    bSa_ = Sim2(R=np.eye(2), t=np.array([1.0, 2.0]), s=3)
+    assert bSa != bSa_
+
+
+def test_rotation() -> None:
+    """ Ensure rotation component is returned properly. """
+    R = np.array([[0, -1], [1, 0]])
+    t = np.array([1, 2])
+    bSa = Sim2(R=R, t=t, s=3.0)
+
+    expected_R = np.array([[0, -1], [1, 0]])
+    assert np.allclose(expected_R, bSa.rotation())
+
+
+def test_translation() -> None:
+    """ Ensure translation component is returned properly. """
+    R = np.array([[0, -1], [1, 0]])
+    t = np.array([1, 2])
+    bSa = Sim2(R=R, t=t, s=3.0)
+
+    expected_t = np.array([1, 2])
+    assert np.allclose(expected_t, bSa.translation())
+
+
+def test_scale() -> None:
+    """ Ensure the scale factor is returned properly. """
+    bRa = np.eye(2)
+    bta = np.array([1, 2])
+    bsa = 3.0
+    bSa = Sim2(R=bRa, t=bta, s=bsa)
+    assert bSa.scale() == 3.0
+
+
+def test_compose():
+    """ Ensure we can compose two Sim(2) transforms together. """
+    scale = 2.0
+    imgSw = Sim2(R=np.eye(2), t=np.array([1.0, 3.0]), s=scale)
+
+    scale = 0.5
+    wSimg = Sim2(R=np.eye(2), t=np.array([-2.0, -6.0]), s=scale)
+
+    # identity
+    wSw = Sim2(R=np.eye(2), t=np.zeros((2,)), s=1.0)
+    assert wSw == imgSw.compose(wSimg)
+
+
+def test_inverse():
+    """ """
+    scale = 2.0
+    imgSw = Sim2(R=np.eye(2), t=np.array([1.0, 3.0]), s=scale)
+
+    scale = 0.5
+    wSimg = Sim2(R=np.eye(2), t=np.array([-2.0, -6.0]), s=scale)
+
+    assert imgSw == wSimg.inverse()
+    assert wSimg == imgSw.inverse()
+
+
+def test_matrix() -> None:
+    """ Ensure 3x3 matrix is formed correctly"""
+    bRa = np.array([[0, -1], [1, 0]])
+    bta = np.array([1, 2])
+    bsa = 3.0
+    bSa = Sim2(R=bRa, t=bta, s=bsa)
+
+    bSa_expected = np.array([[0, -1, 1], [1, 0, 2], [0, 0, 1 / 3]])
+    assert np.allclose(bSa_expected, bSa.matrix())
+
+
+def test_matrix_homogenous_transform() -> None:
+    """ Ensure 3x3 matrix transforms homogenous points as expected."""
+    expected_img_pts = np.array([[6, 4], [4, 6], [0, 0], [1, 7]])
+
+    world_pts = np.array([[2, -1], [1, 0], [-1, -3], [-0.5, 0.5]])
+    scale = 2.0
+    imgSw = Sim2(R=np.eye(2), t=np.array([1.0, 3.0]), s=scale)
+
+    # convert to homogeneous
+    world_pts_h = np.hstack([world_pts, np.ones((4, 1))])
+
+    # multiply each (3,1) homogeneous point vector w/ transform matrix
+    img_pts_h = (imgSw.matrix() @ world_pts_h.T).T
+    # divide (x,y,s) by s
+    img_pts = img_pts_h[:, :2] / img_pts_h[:, 2].reshape(-1, 1)
+    assert np.allclose(expected_img_pts, img_pts)
+
+
+def test_transformFrom_forwards() -> None:
+    """ """
+    expected_img_pts = np.array([[6, 4], [4, 6], [0, 0], [1, 7]])
+
+    world_pts = np.array([[2, -1], [1, 0], [-1, -3], [-0.5, 0.5]])
+    scale = 2.0
+    imgSw = Sim2(R=np.eye(2), t=np.array([1.0, 3.0]), s=scale)
+
+    img_pts = imgSw.transformFrom(world_pts)
+    assert np.allclose(expected_img_pts, img_pts)
+
+
+def test_transformFrom_backwards() -> None:
+    """ """
+    img_pts = np.array([[6, 4], [4, 6], [0, 0], [1, 7]])
+
+    expected_world_pts = np.array([[2, -1], [1, 0], [-1, -3], [-0.5, 0.5]])
+    scale = 0.5
+    wSimg = Sim2(R=np.eye(2), t=np.array([-2.0, -6.0]), s=scale)
+
+    world_pts = wSimg.transformFrom(img_pts)
+    assert np.allclose(expected_world_pts, world_pts)