Add Vertex.from_g2o() classmethod (#71)

graphslam/load.py

@@ -7,15 +7,9 @@
 
 import logging
 
-import numpy as np
-
 from .edge.base_edge import BaseEdge
 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 .vertex import Vertex
 
 
@@ -70,55 +64,25 @@ def custom_edge_from_g2o(line, custom_edge_types):
 
     with open(infile) as f:
         for line in f.readlines():
-            # R^2
-            if line.startswith("VERTEX_XY "):
-                numbers = line[len("VERTEX_XY "):].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
-
-            # R^3
-            if line.startswith("VERTEX_TRACKXYZ "):
-                numbers = line[len("VERTEX_TRACKXYZ "):].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
-
-            # SE(2)
-            if line.startswith("VERTEX_SE2 "):
-                numbers = line[len("VERTEX_SE2 "):].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
-
-            # SE(3)
-            if line.startswith("VERTEX_SE3:QUAT "):
-                numbers = line[len("VERTEX_SE3:QUAT "):].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
-
-            # Custom edge types
-            custom_edge_or_none = custom_edge_from_g2o(line, custom_edge_types)
-            if custom_edge_or_none:
-                edges.append(custom_edge_or_none)
-                continue
-
-            # Odometry Edge
-            edge_or_none = EdgeOdometry.from_g2o(line)
-            if edge_or_none:
-                edges.append(edge_or_none)
-                continue
-
             if line.strip():
+                # Vertex
+                vertex_or_none = Vertex.from_g2o(line)
+                if vertex_or_none:
+                    vertices.append(vertex_or_none)
+                    continue
+
+                # Custom edge types
+                custom_edge_or_none = custom_edge_from_g2o(line, custom_edge_types)
+                if custom_edge_or_none:
+                    edges.append(custom_edge_or_none)
+                    continue
+
+                # Odometry Edge
+                edge_or_none = EdgeOdometry.from_g2o(line)
+                if edge_or_none:
+                    edges.append(edge_or_none)
+                    continue
+
                 _LOGGER.warning("Line not supported -- '%s'", line.rstrip())
 
     return Graph(edges, vertices)

graphslam/vertex.py

@@ -4,6 +4,8 @@
 
 """
 
+import numpy as np
+
 try:
     import matplotlib.pyplot as plt
 except ImportError:  # pragma: no cover
@@ -80,6 +82,52 @@ def to_g2o(self):
 
         raise NotImplementedError
 
+    @classmethod
+    def from_g2o(cls, line):
+        """Load a vertex from a line in a .g2o file.
+
+        Parameters
+        ----------
+        line : str
+            The line from the .g2o file
+
+        Returns
+        -------
+        Vertex, None
+            The instantiated vertex object, or ``None`` if ``line`` does not correspond to a vertex
+
+        """
+        # R^2
+        if line.startswith("VERTEX_XY "):
+            numbers = line[len("VERTEX_XY "):].split()  # fmt: skip
+            arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
+            p = PoseR2(arr)
+            return cls(int(numbers[0]), p)
+
+        # R^3
+        if line.startswith("VERTEX_TRACKXYZ "):
+            numbers = line[len("VERTEX_TRACKXYZ "):].split()  # fmt: skip
+            arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
+            p = PoseR3(arr)
+            return cls(int(numbers[0]), p)
+
+        # SE(2)
+        if line.startswith("VERTEX_SE2 "):
+            numbers = line[len("VERTEX_SE2 "):].split()  # fmt: skip
+            arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
+            p = PoseSE2(arr[:2], arr[2])
+            return cls(int(numbers[0]), p)
+
+        # SE(3)
+        if line.startswith("VERTEX_SE3:QUAT "):
+            numbers = line[len("VERTEX_SE3:QUAT "):].split()  # fmt: skip
+            arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
+            p = PoseSE3(arr[:3], arr[3:])
+            return cls(int(numbers[0]), p)
+
+        # This line does not correspond to a known pose type
+        return None
+
     def plot(self, color="r", marker="o", markersize=3):
         """Plot the vertex.