Pretty tree print for frames

Author: LemonPi

src/pytorch_kinematics/chain.py

@@ -280,6 +280,12 @@ def get_link_names(self):
     def get_frame_indices(self, *frame_names):
         return torch.tensor([self.frame_to_idx[n] for n in frame_names], dtype=torch.long, device=self.device)
 
+    def print_link_tree(self, do_print=True):
+        tree = str(self._root)
+        if do_print:
+            print(tree)
+        return tree
+
     def forward_kinematics(self, th, frame_indices: Optional = None):
         """
         Compute forward kinematics for the given joint values.

src/pytorch_kinematics/frame.py

@@ -85,6 +85,13 @@ def __repr__(self):
                                                                                   self.axis)
 
 
+# prefix components:
+space =  '    '
+branch = '│   '
+# pointers:
+tee =    '├── '
+last =   '└── '
+
 class Frame(object):
     def __init__(self, name=None, link=None, joint=None, children=None):
         self.name = 'None' if name is None else name
@@ -93,10 +100,18 @@ def __init__(self, name=None, link=None, joint=None, children=None):
         if children is None:
             self.children = []
 
-    def __str__(self, level=0):
-        ret = " \t" * level + self.name + "\n"
-        for child in self.children:
-            ret += child.__str__(level + 1)
+    def __str__(self, prefix='', root=True):
+        pointers = [tee] * (len(self.children) - 1) + [last]
+        if root:
+            ret = prefix + self.name + "\n"
+        else:
+            ret = ""
+        for pointer, child in zip(pointers, self.children):
+            ret += prefix + pointer + child.name + "\n"
+            if child.children:
+                extension = branch if pointer == tee else space
+                # i.e. space because last, └── , above so no more |
+                ret += child.__str__(prefix=prefix + extension, root=False)
         return ret
 
     def to(self, *args, **kwargs):

tests/test_inverse_kinematics.py

@@ -213,20 +213,34 @@ def test_extract_serial_chain_from_tree():
     # device = "cpu"
     urdf = "widowx/wx250s.urdf"
     full_urdf = os.path.join(TEST_DIR, urdf)
-    chain = pk.build_serial_chain_from_urdf(open(full_urdf, mode="rb").read(), "ee_gripper_link")
-    # chain = pk.SerialChain(chain, "ee_gripper_link", "base_link")
+    chain = pk.build_chain_from_urdf(open(full_urdf, mode="rb").read())
+    # full frames
+    full_frame_expected = """
+base_link
+└── shoulder_link
+    └── upper_arm_link
+        └── upper_forearm_link
+            └── lower_forearm_link
+                └── wrist_link
+                    └── gripper_link
+                        └── ee_arm_link
+                            ├── gripper_prop_link
+                            └── gripper_bar_link
+                                └── fingers_link
+                                    ├── left_finger_link
+                                    ├── right_finger_link
+                                    └── ee_gripper_link
+    """
+    full_frame = chain.print_link_tree()
+    assert full_frame_expected.strip() == full_frame.strip()
+
+    chain = pk.SerialChain(chain, "ee_gripper_link", "base_link")
+    serial_frame = chain.print_link_tree()
     chain = chain.to(device=device)
 
     # full chain should have DOF = 8, however since we are creating just a serial chain to ee_gripper_link, should be 6
-    # TODO pretty tree print
-    """
-    /
-    └── gripper_bar_link
-        └── fingers_link
-            ├── left_finger_link
-            ├── right_finger_link
-            └── ee_gripper_link
-    """
+    dof = len(chain.get_joints(exclude_fixed=True))
+    assert dof == 6
 
     # robot frame
     pos = torch.tensor([0.0, 0.0, 0.0], device=device)