Automatic Cerberus Street download and reworked link to edge formation

requirements.txt

@@ -6,4 +6,5 @@ torch
 pandas
 torch_geometric
 rosbags
-pathlib
+pathlib
+gdown

src/grfgnn/datasets.py

@@ -9,6 +9,7 @@
 from pathlib import Path
 from itertools import islice
 import os
+from torchvision.datasets.utils import download_file_from_google_drive
 from rosbags.rosbag2 import Writer
 from rosbags.typesys import Stores, get_typestore, get_types_from_msg
 from rosbags.interfaces import ConnectionExtRosbag2
@@ -106,9 +107,9 @@ def raw_file_names(self):
         return ['street.bag']
 
     def download(self):
-        raise Exception(
-            "Please download the bag file from https://drive.google.com/drive/folders/1PdCLMVRqS97Tc9VbJY12EUl9UHpVoO_X and put it in <dataset_dir>/raw/"
-        )
+        download_file_from_google_drive("1rVQW3VPx9WwpJAh8vWKELD0eW9yI_8Vu",
+                                        Path(self.root, 'raw'),
+                                        "street.bag")
 
     @property
     def processed_file_names(self):
@@ -207,7 +208,6 @@ class HyQDataset(Dataset):
     THIS CLASS IS CURRENTLY DEFUNCT.
     """
 
-
     def __init__(self,
                  root=None,
                  transform=None,
@@ -278,8 +278,7 @@ def _find_edge_indexes_from_connections(self, edge_matrix,
                 edge_index.append(i)
         return edge_index
 
-    def _create_edge_attr_matrices(self, HyQ_CSV,
-                                   HyQ_URDF: RobotURDF,
+    def _create_edge_attr_matrices(self, HyQ_CSV, HyQ_URDF: RobotURDF,
                                    joint_to_features_dict: dict[str,
                                                                 list[str]],
                                    features_per_joint):

src/grfgnn/urdfParser.py

@@ -4,23 +4,27 @@
 from pathlib import Path
 
 
+class InvalidURDFException(Exception):
+    pass
+
+
 class RobotURDF():
 
     class Node():
         """
         Simple class that holds a node's name.
         """
 
-        def __init__(self, name):
+        def __init__(self, name: str):
             self.name = name
 
     class Edge():
         """
         Simple class that holds an edge's name, and
-        both connections.
+        the names of both connections.
         """
 
-        def __init__(self, name, parent, child):
+        def __init__(self, name: str, parent: str, child: str):
             self.name = name
             self.parent = parent
             self.child = child
@@ -83,28 +87,46 @@ def __init__(self,
             # Create joints from the links, to set as the edges
             self.edges = []
             for link in self.robot_urdf.links:
-                # Check that two joints connect to this link
+                # Get all the connections to this link
                 connections = []
                 for joint in self.robot_urdf.joints:
-                    if (joint.parent == link.name or joint.child == link.name):
-                        connections.append(joint.name)
-
-                # If two joints do connect, add it as an edge
-                if len(connections) == 2:
-                    new_edge = self.Edge(name=link.name,
-                                         parent=connections[0],
-                                         child=connections[1])
-                    self.edges.append(new_edge)
-                if len(connections
-                       ) > 2:  # TODO: Don't connect children to other children
-                    # Setup an edge for each pair of connections
-                    count = 0
-                    for i in range(0, len(connections)):
-                        for j in range(i + 1, len(connections)):
-                            self.edges.append(
-                                self.Edge(name=link.name + str(count),
-                                          parent=connections[i],
-                                          child=connections[j]))
+                    # The joint's parent is the link's child, and visa versa
+                    if (joint.parent == link.name):
+                        connections.append((joint.name, 'child'))
+                    elif (joint.child == link.name):
+                        connections.append((joint.name, 'parent'))
+
+                # If there is 1 connection, drop this link
+                if (len(connections) == 1):
+                    continue
+                if (len(connections) < 1):
+                    raise InvalidURDFException("Link connected to no joints.")
+
+                # Find the singular link parent
+                parent_connection = None
+                for conn in connections:
+                    if conn[1] == 'parent':
+                        if parent_connection is not None:
+                            raise InvalidURDFException(
+                                "Link has more than two parent joints.")
+                        else:
+                            parent_connection = conn
+                            connections.remove(conn)
+
+                # Create an edge from each link parent to each link child
+                new_edge = None
+                if (len(connections) == 1):
+                    self.edges.append(
+                        self.Edge(name=link.name,
+                                  parent=parent_connection[0],
+                                  child=connections[0][0]))
+                else:  # If there are multiple edges for this one link
+                       # give them each a unique name.
+                    for child_conn in connections:
+                        self.edges.append(
+                            self.Edge(name=link.name + "_to_" + child_conn[0],
+                                      parent=parent_connection[0],
+                                      child=child_conn[0]))
 
     def create_updated_urdf_file(self):
         """
@@ -234,10 +256,10 @@ def get_edge_name_to_connections_dict(self):
         it's connected to.
 
         Returns:
-            edge_dict (dict[str, np.array([[int, int],[int, int]])]): 
+            edge_dict (dict[str, np.array)]): 
                 This dictionary takes the name of an edge in the URDF
-                as input. It returns an np.array that contains the 
-                indices of the two nodes that this edge is connected to.
+                as input. It returns an np.array (N, 2), where N is the
+                number of connections this name supports.
         """
 
         node_dict = self.get_node_name_to_index_dict()

tests/testDatasets.py

@@ -0,0 +1,33 @@
+import unittest
+from pathlib import Path
+from grfgnn import RobotURDF, CerberusStreetDataset
+
+
+class TestCerberusStreetDataset(unittest.TestCase):
+
+    def setUp(self):
+        # Load the A1 URDF file
+        self.a1_path = Path(
+            Path(__file__).cwd(), 'urdf_files', 'A1', 'a1.urdf').absolute()
+        A1_URDF = RobotURDF(self.a1_path, 'package://a1_description/',
+                            'unitree_ros/robots/a1_description', True)
+
+        # Create the dataset
+        self.dataset_path = Path(
+            Path(__file__).cwd(), 'datasets', 'cerberus_street')
+        dataset = CerberusStreetDataset(self.dataset_path, A1_URDF)
+
+    def test_download_and_processing(self):
+        """
+        Check that the Dataset class successfully downloads and processes
+        the data.
+        """
+
+        self.bag_path = self.dataset_path / 'raw' / 'street.bag'
+        self.assertTrue(self.bag_path.is_file())
+        self.processed_txt_path = self.dataset_path / 'processed' / 'processed.txt'
+        self.assertTrue(self.processed_txt_path.is_file())
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/testUrdfParser.py

@@ -1,10 +1,11 @@
 from pathlib import Path
 import unittest
-import os
 from grfgnn import RobotURDF
+from pathlib import Path
+import copy
 import pandas as pd
 import numpy as np
-from pathlib import Path
+import os
 
 
 class TestRobotURDF(unittest.TestCase):
@@ -21,31 +22,86 @@ def setUp(self):
 
     def test_constructor(self):
         """
-        Check if self.nodes has all the name in the URDF file, and
-        check if self.edges has all the name in the URDF file.
+        Check that the constructor properly assigns all of the links and joints
+        to a node/edge.
         """
 
-        nodes_name = {
-            'world', 'base_link', 'trunk', 'lf_hipassembly', 'lh_hipassembly',
-            'rf_hipassembly', 'rh_hipassembly', 'lf_upperleg', 'lh_upperleg',
-            'rf_upperleg', 'rh_upperleg', 'lf_lowerleg', 'lh_lowerleg',
-            'rf_lowerleg', 'rh_lowerleg', 'lf_foot', 'lh_foot', 'rf_foot',
-            'rh_foot'
-        }
-
-        edges_name = {
+        node_names = [
+            'world', 'base_link', 'trunk', 'lf_hipassembly', 'lf_upperleg',
+            'lf_lowerleg', 'lf_foot', 'rf_hipassembly', 'rf_upperleg',
+            'rf_lowerleg', 'rf_foot', 'lh_hipassembly', 'lh_upperleg',
+            'lh_lowerleg', 'lh_foot', 'rh_hipassembly', 'rh_upperleg',
+            'rh_lowerleg', 'rh_foot'
+        ]
+        edge_names = [
             'floating_base_joint', 'floating_base', 'lf_haa_joint',
-            'lh_haa_joint', 'rf_haa_joint', 'rh_haa_joint', 'lf_hfe_joint',
-            'lh_hfe_joint', 'rf_hfe_joint', 'rh_hfe_joint', 'lf_kfe_joint',
-            'lh_kfe_joint', 'rf_kfe_joint', 'rh_kfe_joint', 'lf_foot_joint',
-            'lh_foot_joint', 'rf_foot_joint', 'rh_foot_joint'
-        }
-
-        for node in self.HyQ_URDF.nodes:
-            self.assertTrue(node.name in nodes_name)
-
-        for edge in self.HyQ_URDF.edges:
-            self.assertTrue(edge.name in edges_name)
+            'lf_hfe_joint', 'lf_kfe_joint', 'lf_foot_joint', 'rf_haa_joint',
+            'rf_hfe_joint', 'rf_kfe_joint', 'rf_foot_joint', 'lh_haa_joint',
+            'lh_hfe_joint', 'lh_kfe_joint', 'lh_foot_joint', 'rh_haa_joint',
+            'rh_hfe_joint', 'rh_kfe_joint', 'rh_foot_joint'
+        ]
+
+        # Check for the normal case: where links become nodes and
+        # joints become edges.
+        node_names_copy = copy.deepcopy(node_names)
+        for i, node in enumerate(self.HyQ_URDF.nodes):
+            self.assertTrue(node.name in node_names_copy)
+            node_names_copy.remove(node.name)
+        self.assertEqual(0, len(node_names_copy))
+
+        edge_names_copy = copy.deepcopy(edge_names)
+        for i, edge in enumerate(self.HyQ_URDF.edges):
+            self.assertTrue(edge.name in edge_names_copy)
+            edge_names_copy.remove(edge.name)
+        self.assertEqual(0, len(edge_names_copy))
+
+        # Check for the swapped case: where links become edges and
+        # joints become nodes.
+        edge_names_copy = copy.deepcopy(edge_names)
+        for i, node in enumerate(self.HyQ_URDF_swapped.nodes):
+            self.assertTrue(node.name in edge_names_copy)
+            edge_names_copy.remove(node.name)
+        self.assertEqual(0, len(edge_names_copy))
+
+        # Remember that links that don't connect to two or more
+        # joints get dropped, as they can't be represented as an edge.
+        # Additionally, links with multiple children joints get one
+        # edge for each child.
+        desired_edges = [
+            RobotURDF.Edge('base_link', "floating_base_joint",
+                           "floating_base"),
+            RobotURDF.Edge('trunk_to_lf_haa_joint', "floating_base",
+                           "lf_haa_joint"),
+            RobotURDF.Edge('trunk_to_lh_haa_joint', "floating_base",
+                           "lh_haa_joint"),
+            RobotURDF.Edge('trunk_to_rf_haa_joint', "floating_base",
+                           "rf_haa_joint"),
+            RobotURDF.Edge('trunk_to_rh_haa_joint', "floating_base",
+                           "rh_haa_joint"),
+            RobotURDF.Edge('lf_hipassembly', "lf_haa_joint", "lf_hfe_joint"),
+            RobotURDF.Edge('lf_upperleg', "lf_hfe_joint", "lf_kfe_joint"),
+            RobotURDF.Edge('lf_lowerleg', "lf_kfe_joint", "lf_foot_joint"),
+            RobotURDF.Edge('rf_hipassembly', "rf_haa_joint", "rf_hfe_joint"),
+            RobotURDF.Edge('rf_upperleg', "rf_hfe_joint", "rf_kfe_joint"),
+            RobotURDF.Edge('rf_lowerleg', "rf_kfe_joint", "rf_foot_joint"),
+            RobotURDF.Edge('lh_hipassembly', "lh_haa_joint", "lh_hfe_joint"),
+            RobotURDF.Edge('lh_upperleg', "lh_hfe_joint", "lh_kfe_joint"),
+            RobotURDF.Edge('lh_lowerleg', "lh_kfe_joint", "lh_foot_joint"),
+            RobotURDF.Edge('rh_hipassembly', "rh_haa_joint", "rh_hfe_joint"),
+            RobotURDF.Edge('rh_upperleg', "rh_hfe_joint", "rh_kfe_joint"),
+            RobotURDF.Edge('rh_lowerleg', "rh_kfe_joint", "rh_foot_joint")
+        ]
+        for i, edge in enumerate(self.HyQ_URDF_swapped.edges):
+            match_found = False
+            for j, desired_edge in enumerate(desired_edges):
+                if edge.name == desired_edge.name:
+                    self.assertEqual(edge.child, desired_edge.child)
+                    self.assertEqual(edge.parent, desired_edge.parent)
+                    desired_edges.remove(desired_edge)
+                    match_found = True
+                    break
+            self.assertTrue(match_found)
+        self.assertEqual(0, len(desired_edges))
 
     def test_create_updated_urdf_file(self):
         """