Support loading custom edge types (#70)

graphslam/edge/base_edge.py

@@ -212,6 +212,9 @@ def equals(self, other, tol=1e-6):
             Whether the two edges are equal
 
         """
+        if not type(self) is type(other):
+            return False
+
         if len(self.vertex_ids) != len(other.vertex_ids):
             return False
 

graphslam/graph.py

@@ -508,7 +508,9 @@ def to_g2o(self, outfile):
                 f.write(v.to_g2o())
 
             for e in self._edges:
-                f.write(e.to_g2o())
+                edge_str_or_none = e.to_g2o()
+                if edge_str_or_none:
+                    f.write(edge_str_or_none)
 
     def plot(self, vertex_color="r", vertex_marker="o", vertex_markersize=3, edge_color="b", title=None):
         """Plot the graph.

graphslam/load.py

@@ -9,6 +9,7 @@
 
 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
@@ -21,13 +22,15 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-def load_g2o(infile):
+def load_g2o(infile, custom_edge_types=None):
     r"""Load a graph from a .g2o file.
 
     Parameters
     ----------
     infile : str
         The path to the .g2o file
+    custom_edge_types : list[type], None
+        A list of custom edge types, which must be subclasses of ``BaseEdge``
 
     Returns
     -------
@@ -38,6 +41,33 @@ def load_g2o(infile):
     edges = []
     vertices = []
 
+    custom_edge_types = custom_edge_types or []
+    for edge_type in custom_edge_types:
+        assert issubclass(edge_type, BaseEdge)
+
+    def custom_edge_from_g2o(line, custom_edge_types):
+        """Load a custom edge from a .g2o line.
+
+        Parameters
+        ----------
+        line : str
+            A line from a .g2o file
+        custom_edge_types : list[type]
+            A list of custom edge types, which must be subclasses of ``BaseEdge``
+
+        Returns
+        -------
+        BaseEdge, None
+            The instantiated edge object, or ``None`` if the line does not correspond to any of the custom edge types
+
+        """
+        for custom_edge_type in custom_edge_types:
+            edge_or_none = custom_edge_type.from_g2o(line)
+            if edge_or_none:
+                return edge_or_none
+
+        return None
+
     with open(infile) as f:
         for line in f.readlines():
             # R^2
@@ -76,6 +106,12 @@ def load_g2o(infile):
                 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:

tests/test_base_edge.py

@@ -94,9 +94,12 @@ def test_equals(self):
         v1 = Vertex(1, p1)
         v2 = Vertex(2, p2)
 
+        e_simple = SimpleEdge(0, 1, 0)
         e1 = EdgeOdometry([1, 2], np.eye(2), estimate, [v1, v2])
         e2 = EdgeOdometry([1, 2], np.eye(2), estimate, [v1, v2])
 
+        self.assertFalse(e_simple.equals(e1))
+
         self.assertTrue(e1.equals(e2))
 
         e2.estimate = 123

tests/test_load.py

@@ -9,11 +9,61 @@
 import unittest
 from unittest import mock
 
-from graphslam.load import load_g2o_r2, load_g2o_r3, load_g2o_se2, load_g2o_se3
+import numpy as np
 
+from graphslam.graph import Graph
+from graphslam.load import load_g2o, load_g2o_r2, load_g2o_r3, load_g2o_se2, load_g2o_se3
+from graphslam.pose.r2 import PoseR2
+from graphslam.util import upper_triangular_matrix_to_full_matrix
+from graphslam.vertex import Vertex
+
+from .edge_types import BaseEdgeForTests
 from .patchers import FAKE_FILE, open_fake_file
 
 
+class EdgeWithoutToG2OWithoutFromG2O(BaseEdgeForTests):
+    """An edge class without ``to_g2o`` and ``from_g2o`` support.
+
+    This class is only compatible with ``PoseR2`` poses.
+
+    """
+
+    def calc_error(self):
+        """Return an error vector."""
+        return np.array([1.0, 2.0])
+
+
+class EdgeWithToG2OWithoutFromG2O(EdgeWithoutToG2OWithoutFromG2O):
+    """An edge class with a ``to_g2o`` method but not a ``from_g2o`` method."""
+
+    def to_g2o(self):
+        """Write to g2o format."""
+        # fmt: off
+        return "TestEdge {} {} {} ".format(self.vertex_ids[0], self.estimate[0], self.estimate[1]) + " ".join([str(x) for x in self.information[np.triu_indices(2, 0)]]) + "\n"
+        # fmt: on
+
+
+class EdgeWithoutToG2OWithFromG2O(EdgeWithoutToG2OWithoutFromG2O):
+    """An edge class with a ``from_g2o`` method but not a ``to_g2o`` method."""
+
+    @classmethod
+    def from_g2o(cls, line):
+        """Write to g2o format."""
+        if line.startswith("TestEdge "):
+            numbers = line[len("TestEdge "):].split()  # fmt: skip
+            arr = np.array([float(number) for number in numbers[1:]], dtype=np.float64)
+            vertex_ids = [int(numbers[0])]
+            estimate = arr[:2]
+            information = upper_triangular_matrix_to_full_matrix(arr[2:], 2)
+            return cls(vertex_ids, information, estimate)
+
+        return None
+
+
+class EdgeWithToG2OWithFromG2O(EdgeWithToG2OWithoutFromG2O, EdgeWithoutToG2OWithFromG2O):
+    """An edge class with ``to_g2o`` and ``from_g2o`` methods."""
+
+
 class TestLoad(unittest.TestCase):
     """Tests for the ``load`` functions."""
 
@@ -82,3 +132,67 @@ def test_load_g2o_se3(self):
         with mock.patch("graphslam.load.open", open_fake_file):
             g2 = load_g2o_se3("test.g2o")
             self.assertAlmostEqual(chi2, g2.calc_chi2())
+
+    def test_load_custom_edge_without_to_g2o_without_from_g2o(self):
+        """Test that loading a graph with an edge type that does not have ``to_g2o`` and ``from_g2o`` methods works as expected."""
+        p0 = PoseR2([1.0, 2.0])
+        v0 = Vertex(0, p0)
+
+        e = EdgeWithoutToG2OWithoutFromG2O([0], np.eye(2), np.ones(2))
+
+        g = Graph([e], [v0])
+
+        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("test.g2o")
+            self.assertFalse(g.equals(g2))
+
+    def test_load_custom_edge_with_to_g2o_without_from_g2o(self):
+        """Test that loading a graph with an edge type that has a ``to_g2o`` method but not a ``from_g2o`` method works as expected."""
+        p0 = PoseR2([1.0, 2.0])
+        v0 = Vertex(0, p0)
+
+        e = EdgeWithToG2OWithoutFromG2O([0], np.eye(2), np.ones(2))
+
+        g = Graph([e], [v0])
+
+        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("test.g2o", [EdgeWithToG2OWithoutFromG2O])
+            self.assertFalse(g.equals(g2))
+
+    def test_load_custom_edge_without_to_g2o_with_from_g2o(self):
+        """Test that loading a graph with an edge type that has a ``from_g2o`` method but not a ``to_g2o`` method works as expected."""
+        p0 = PoseR2([1.0, 2.0])
+        v0 = Vertex(0, p0)
+
+        e = EdgeWithoutToG2OWithFromG2O([0], np.eye(2), np.ones(2))
+
+        g = Graph([e], [v0])
+
+        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("test.g2o", [EdgeWithoutToG2OWithFromG2O])
+            self.assertFalse(g.equals(g2))
+
+    def test_load_custom_edge_with_to_g2o_with_from_g2o(self):
+        """Test that loading a graph with an edge type that has ``to_g2o`` and ``from_g2o`` methods works as expected."""
+        p0 = PoseR2([1.0, 2.0])
+        v0 = Vertex(0, p0)
+
+        e = EdgeWithToG2OWithFromG2O([0], np.eye(2), np.ones(2))
+
+        g = Graph([e], [v0])
+
+        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("test.g2o", [EdgeWithToG2OWithFromG2O])
+            self.assertTrue(g.equals(g2))