Support VERTEX_XY and VERTEX_TRACKXYZ vertices in .g2o files (#47)

* Support VERTEX_XY and VERTEX_TRACKXYZ vertices in .g2o files

* Fix test

* Coverage

graphslam/load.py

@@ -11,6 +11,8 @@
 
 from .edge.edge_odometry import EdgeOdometry
 from .graph import Graph
+from .pose.r2 import PoseR2
+from .pose.r3 import PoseR3
 from .pose.se2 import PoseSE2
 from .pose.se3 import PoseSE3
 from .util import upper_triangular_matrix_to_full_matrix
@@ -20,6 +22,72 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+def load_g2o_r2(infile):
+    r"""Load an :math:`\mathbb{R}^2` graph from a .g2o file.
+
+    Parameters
+    ----------
+    infile : str
+        The path to the .g2o file
+
+    Returns
+    -------
+    Graph
+        The loaded graph
+
+    """
+    edges = []
+    vertices = []
+
+    with open(infile) as f:
+        for line in f.readlines():
+            if line.startswith("VERTEX_XY"):
+                numbers = line[len("VERTEX_XY") + 1:].split()  # fmt: skip
+                arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
+                p = PoseR2(arr)
+                v = Vertex(int(numbers[0]), p)
+                vertices.append(v)
+                continue
+
+            if line.strip():
+                _LOGGER.warning("Line not supported -- '%s'", line.rstrip())
+
+    return Graph(edges, vertices)
+
+
+def load_g2o_r3(infile):
+    r"""Load an :math:`\mathbb{R}^3` graph from a .g2o file.
+
+    Parameters
+    ----------
+    infile : str
+        The path to the .g2o file
+
+    Returns
+    -------
+    Graph
+        The loaded graph
+
+    """
+    edges = []
+    vertices = []
+
+    with open(infile) as f:
+        for line in f.readlines():
+            if line.startswith("VERTEX_TRACKXYZ"):
+                numbers = line[len("VERTEX_TRACKXYZ") + 1:].split()  # fmt: skip
+                arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
+                p = PoseR3(arr)
+                v = Vertex(int(numbers[0]), p)
+                vertices.append(v)
+                continue
+
+            if line.strip():
+                _LOGGER.warning("Line not supported -- '%s'", line.rstrip())
+
+    return Graph(edges, vertices)
+
+
 def load_g2o_se2(infile):
     """Load an :math:`SE(2)` graph from a .g2o file.
 
@@ -40,15 +108,23 @@ def load_g2o_se2(infile):
     with open(infile) as f:
         for line in f.readlines():
             if line.startswith("VERTEX_SE2"):
-                numbers = line[10:].split()
+                numbers = line[len("VERTEX_SE2") + 1:].split()  # fmt: skip
                 arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
                 p = PoseSE2(arr[:2], arr[2])
                 v = Vertex(int(numbers[0]), p)
                 vertices.append(v)
                 continue
 
+            if line.startswith("VERTEX_XY"):
+                numbers = line[len("VERTEX_XY") + 1:].split()  # fmt: skip
+                arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
+                p = PoseR2(arr)
+                v = Vertex(int(numbers[0]), p)
+                vertices.append(v)
+                continue
+
             if line.startswith("EDGE_SE2"):
-                numbers = line[9:].split()
+                numbers = line[len("EDGE_SE2") + 1:].split()  # fmt: skip
                 arr = np.array([float(number) for number in numbers[2:]], dtype=np.float64)
                 vertex_ids = [int(numbers[0]), int(numbers[1])]
                 estimate = PoseSE2(arr[:2], arr[2])
@@ -83,15 +159,23 @@ def load_g2o_se3(infile):
     with open(infile) as f:
         for line in f.readlines():
             if line.startswith("VERTEX_SE3:QUAT"):
-                numbers = line[16:].split()
+                numbers = line[len("VERTEX_SE3:QUAT") + 1:].split()  # fmt: skip
                 arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
                 p = PoseSE3(arr[:3], arr[3:])
                 v = Vertex(int(numbers[0]), p)
                 vertices.append(v)
                 continue
 
+            if line.startswith("VERTEX_TRACKXYZ"):
+                numbers = line[len("VERTEX_TRACKXYZ") + 1:].split()  # fmt: skip
+                arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
+                p = PoseR3(arr)
+                v = Vertex(int(numbers[0]), p)
+                vertices.append(v)
+                continue
+
             if line.startswith("EDGE_SE3:QUAT"):
-                numbers = line[14:].split()
+                numbers = line[len("EDGE_SE3:QUAT") + 1:].split()  # fmt: skip
                 arr = np.array([float(number) for number in numbers[2:]], dtype=np.float64)
                 vertex_ids = [int(numbers[0]), int(numbers[1])]
                 estimate = PoseSE3(arr[:3], arr[3:7])

graphslam/vertex.py

@@ -72,6 +72,12 @@ def to_g2o(self):
                 self.pose[6],
             )
 
+        if isinstance(self.pose, PoseR2):
+            return "VERTEX_XY {} {} {}\n".format(self.id, self.pose[0], self.pose[1])
+
+        if isinstance(self.pose, PoseR3):
+            return "VERTEX_TRACKXYZ {} {} {} {}\n".format(self.id, self.pose[0], self.pose[1], self.pose[2])
+
         raise NotImplementedError
 
     def plot(self, color="r", marker="o", markersize=3):

tests/test_graph.py

@@ -100,19 +100,16 @@ def test_optimize_fix_12(self):
     # pylint: disable=protected-access
     def test_to_g2o(self):
         """Test that the ``to_g2o`` method is implemented correctly, or raises ``NotImplementedError``."""
-        # Supported types
+        # Fully supported types
         if isinstance(self.g._vertices[0].pose, (PoseSE2, PoseSE3)):
             print(self.g._vertices[0].to_g2o())
             print(self.g._edges[0].to_g2o())
 
             with patch("graphslam.graph.open", mock_open()):
                 self.g.to_g2o("test.g2o")
 
-        # Unsupported types
-        else:
-            with self.assertRaises(NotImplementedError):
-                print(self.g._vertices[0].to_g2o())
-
+        # Unsupported edges
+        if isinstance(self.g._vertices[0].pose, (PoseR2, PoseR3)):
             with self.assertRaises(NotImplementedError):
                 print(self.g._edges[0].to_g2o())
 

tests/test_load.py

@@ -9,14 +9,48 @@
 import unittest
 from unittest import mock
 
-from graphslam.load import load_g2o_se2, load_g2o_se3
+from graphslam.load import load_g2o_r2, load_g2o_r3, load_g2o_se2, load_g2o_se3
 
 from .patchers import open_fake_file
 
 
 class TestLoad(unittest.TestCase):
     """Tests for the ``load`` functions."""
 
+    def test_load_g2o_r2(self):
+        """Test the ``load_g2o_r2()`` function."""
+        infile = os.path.join(os.path.dirname(__file__), "test_r2.g2o")
+        g = load_g2o_r2(infile)
+        chi2 = g.calc_chi2()
+
+        # There are currently no edges in this graph
+        self.assertEqual(chi2, 0.0)
+
+        with mock.patch("graphslam.graph.open", open_fake_file):
+            g.to_g2o("test.g2o")
+
+        with mock.patch("graphslam.load.open", open_fake_file):
+            g2 = load_g2o_r2("test.g2o")
+            self.assertAlmostEqual(chi2, g2.calc_chi2())
+
+    def test_load_g2o_r3(self):
+        """Test the ``load_g2o_r3()`` function."""
+        infile = os.path.join(os.path.dirname(__file__), "test_r3.g2o")
+        g = load_g2o_r3(infile)
+        chi2 = g.calc_chi2()
+
+        # There are currently no edges in this graph
+        self.assertEqual(chi2, 0.0)
+
+        with mock.patch("graphslam.graph.open", open_fake_file):
+            g.to_g2o("test.g2o")
+
+        with mock.patch("graphslam.load.open", open_fake_file):
+            g2 = load_g2o_r3("test.g2o")
+            self.assertAlmostEqual(chi2, g2.calc_chi2())
+
+    #########################
+
     def test_load_g2o_se2(self):
         """Test the ``load_g2o_se2()`` function."""
         infile = os.path.join(os.path.dirname(__file__), "test_se2.g2o")

tests/test_r2.g2o

@@ -0,0 +1,3 @@
+VERTEX_XY 0 0.000000 0.000000
+VERTEX_XY 1 0.000000 0.000000
+This line is not supported

tests/test_r3.g2o

@@ -0,0 +1,6 @@
+VERTEX_SE3:QUAT 0 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 1.000000
+VERTEX_SE3:QUAT 1 0.000000 1.000000 2.000000 0.500000 0.500000 0.500000 0.500000
+VERTEX_TRACKXYZ 2 0.000000 0.000000 0.000000
+VERTEX_TRACKXYZ 3 0.000000 1.000000 2.000000
+EDGE_SE3:QUAT 0 1 3.000000 4.000000 5.000000 0.500000 0.500000 0.500000 0.500000 1.000000 2.000000 3.000000 4.000000 5.000000 6.000000 7.000000 8.000000 9.000000 10.000000 11.000000 12.000000 13.000000 14.000000 15.000000 16.000000 17.000000 18.000000 19.000000 20.000000 21.000000
+This line is not supported

tests/test_se2.g2o

@@ -1,4 +1,6 @@
 VERTEX_SE2 0 0.000000 0.000000 0.000000
 VERTEX_SE2 1 0.000000 0.000000 1.000000
+VERTEX_XY 2 0.000000 0.000000
+VERTEX_XY 3 0.000000 0.000000
 EDGE_SE2 0 1 0.000000 0.000000 0.000000 1.000000 2.000000 3.000000 4.000000 5.000000 6.000000
 This line is not supported

tests/test_se3.g2o

@@ -1,4 +1,6 @@
 VERTEX_SE3:QUAT 0 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 1.000000
 VERTEX_SE3:QUAT 1 0.000000 1.000000 2.000000 0.500000 0.500000 0.500000 0.500000
+VERTEX_XY 2 0.000000 0.000000
+VERTEX_XY 3 0.000000 0.000000
 EDGE_SE3:QUAT 0 1 3.000000 4.000000 5.000000 0.500000 0.500000 0.500000 0.500000 1.000000 2.000000 3.000000 4.000000 5.000000 6.000000 7.000000 8.000000 9.000000 10.000000 11.000000 12.000000 13.000000 14.000000 15.000000 16.000000 17.000000 18.000000 19.000000 20.000000 21.000000
 This line is not supported

tests/test_vertex.py

@@ -45,6 +45,14 @@ def test_plot(self):
                 fig.add_subplot(111, projection="3d")
             v.plot()
 
+    def test_to_g2o(self):
+        """Test that an unsupported pose type cannot be written to a .g2o file."""
+        # Use `None` in lieue of an actual unsupported pose
+        v = Vertex(0, None)
+
+        with self.assertRaises(NotImplementedError):
+            v.to_g2o()
+
 
 if __name__ == "__main__":
     unittest.main()