Added graphs and a buggy way to generate a urdf.

.gitignore

@@ -166,3 +166,5 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 poetry.lock
+examples/*/*.urdf
+examples/*/*.stl

examples/bike/main.py

@@ -1,6 +1,17 @@
+import matplotlib.pyplot as plt
+import networkx as nx
+
 import onshape_api as osa
-from onshape_api.graph import create_graph
-from onshape_api.parse import get_instances, get_mates, get_occurences, get_parts, get_subassemblies
+from onshape_api.graph import create_graph, get_urdf_components, show_graph
+from onshape_api.models.robot import Robot
+from onshape_api.parse import (
+    get_instances,
+    get_mass_properties,
+    get_mates,
+    get_occurences,
+    get_parts,
+    get_subassemblies,
+)
 
 # Initialize the client with the constructed path
 client = osa.Client()
@@ -22,6 +33,13 @@
 instances = get_instances(assembly)
 subassemblies = get_subassemblies(assembly, instances)
 parts = get_parts(assembly, instances)
+mass_properties = get_mass_properties(parts, doc.wid, client)
 mates = get_mates(assembly)
 
-create_graph(occurences=occurences, instances=instances, subassemblies=subassemblies, parts=parts, mates=mates)
+graph = create_graph(occurences=occurences, instances=instances, parts=parts, mates=mates, directed=False)
+# show_graph(graph)
+
+links, joints = get_urdf_components(graph, doc.wid, parts, mass_properties, mates, client)
+
+robot = Robot(name="my_robot", links=links, joints=joints)
+robot.save("bike.urdf")

onshape_api/connect.py

@@ -6,6 +6,7 @@
 import secrets
 import string
 from enum import Enum
+from typing import BinaryIO
 from urllib.parse import parse_qs, urlencode, urlparse
 
 import requests
@@ -251,7 +252,7 @@ def get_assembly(self, did, wtype, wid, eid, configuration="default"):
     def get_parts(self, did, wid, eid):
         pass
 
-    def download_stl(self, did, wid, eid, partID, save_path="./part.stl"):
+    def download_stl(self, did, wid, eid, partID, buffer: BinaryIO):
         """
         Exports STL export from a part studio and saves it to a file.
 
@@ -275,13 +276,9 @@ def download_stl(self, did, wid, eid, partID, save_path="./part.stl"):
             log_response=False,
         )
         if response.status_code == 200:
-            with open(save_path, "wb") as file:
-                file.write(response.content)
-            LOGGER.info(f"STL file saved to {save_path}")
-            return f"STL file saved to {save_path}"
+            buffer.write(response.content)
         else:
             LOGGER.info(f"Failed to download STL file: {response.status_code} - {response.text}")
-            return f"Failed to download STL file: {response.status_code} - {response.text}"
 
     def get_mass_properties(self, did, wid, eid, partID):
         """

onshape_api/graph.py

@@ -1,27 +1,57 @@
+import io
+from typing import Union
+
 import matplotlib.pyplot as plt
 import networkx as nx
+import numpy as np
+import stl
 
+from onshape_api.connect import Client
 from onshape_api.models.assembly import (
+    MATETYPE,
     Instance,
     InstanceType,
     MateFeatureData,
     Occurrence,
     Part,
-    SubAssembly,
 )
+from onshape_api.models.geometry import MeshGeometry
+from onshape_api.models.joint import FixedJoint, JointLimits, RevoluteJoint
+from onshape_api.models.link import (
+    COLORS,
+    Axis,
+    CollisionLink,
+    Inertia,
+    InertialLink,
+    Link,
+    Material,
+    Origin,
+    VisualLink,
+)
+from onshape_api.models.mass import MassModel
 from onshape_api.parse import MATE_JOINER, SUBASSEMBLY_JOINER
-from onshape_api.utilities.helpers import print_dict
 from onshape_api.utilities.logging import LOGGER
+from onshape_api.utilities.mesh import transform_mesh
+
+
+def show_graph(graph: nx.Graph):
+    nx.draw_circular(graph, with_labels=True)
+    plt.show()
 
+def convert_to_digraph(graph: nx.Graph) -> nx.DiGraph:
+    _centrality = nx.closeness_centrality(graph)
+    _root_node = max(_centrality, key=_centrality.get)
+    _graph = nx.bfs_tree(graph, _root_node)
+    return _graph, _root_node
 
 def create_graph(
     occurences: dict[str, Occurrence],
     instances: dict[str, Instance],
-    subassemblies: dict[str, SubAssembly],
     parts: dict[str, Part],
     mates: dict[str, MateFeatureData],
+    directed: bool = True,
 ):
-    graph = nx.DiGraph()
+    graph = nx.Graph()
 
     for occurence in occurences:
         if instances[occurence].type == InstanceType.PART:
@@ -33,19 +63,134 @@ def create_graph(
             except KeyError:
                 LOGGER.warning(f"Part {occurence} not found")
 
-        # elif instances[occurence].type == InstanceType.ASSEMBLY:
-        #     try:
-        #         graph.add_node(occurence, **subassemblies[occurence].model_dump())
-        #     except KeyError:
-        #         LOGGER.warning(f"SubAssembly {occurence} not found")
-
     for mate in mates:
         try:
             child, parent = mate.split(MATE_JOINER)
-            print(child, parent)
             graph.add_edge(child, parent, **mates[mate].model_dump())
         except KeyError:
             LOGGER.warning(f"Mate {mate} not found")
 
-    nx.draw_circular(graph, with_labels=True)
-    plt.show()
+    if directed:
+        graph = convert_to_digraph(graph)
+
+    return graph
+
+def flip_tuple(t: tuple) -> tuple:
+    return t[::-1]
+
+def download_stl_mesh(did, wid, eid, partID, client: Client, path: str):
+    buffer = io.BytesIO()
+    client.download_stl(did, wid, eid, partID, buffer)
+    buffer.seek(0)
+
+    raw_mesh = stl.mesh.Mesh.from_file(None, fh=buffer)
+    transformed_mesh = transform_mesh(raw_mesh, np.eye(4))
+    transformed_mesh.save(path)
+
+    return path
+
+def get_urdf_components(
+    graph: Union[nx.Graph, nx.DiGraph],
+    workspaceId: str,
+    parts: dict[str, Part],
+    mass_properties: dict[str, MassModel],
+    mates: dict[str, MateFeatureData],
+    client: Client,
+):
+    if not isinstance(graph, nx.DiGraph):
+        graph, root_node = convert_to_digraph(graph)
+
+    joints = []
+    links = []
+
+    _child_transforms = {}
+    _sorted_nodes = list(nx.topological_sort(graph))
+
+    for node in _sorted_nodes:
+        for edge in graph.edges(node):
+            # check if the mate is between root assembly and subassembly
+            if "".join(edge).count(SUBASSEMBLY_JOINER) == 1:
+                _mate_key = MATE_JOINER.join(edge)
+            else:
+                _mate_key = MATE_JOINER.join(flip_tuple(edge))
+
+            if _mate_key in mates:
+                if mates[_mate_key].mateType == MATETYPE.REVOLUTE:
+                    _child_transforms[node] = mates[_mate_key].matedEntities[0].matedCS.part_to_mate_transform
+                    joints.append(
+                        RevoluteJoint(
+                            name=f"joint_{edge[0]}_{edge[1]}",
+                            parent=edge[0],
+                            child=edge[1],
+                            origin=Origin.from_matrix(mates[_mate_key].matedEntities[1].matedCS.part_to_mate_transform),
+                            limits=JointLimits(effort=1.0, velocity=1.0, lower=-1.0, upper=1.0),
+                            axis=Axis((1, 0, 0)),
+                    ))
+
+                elif mates[_mate_key].mateType == MATETYPE.FASTENED:
+                    _child_transforms[node] = mates[_mate_key].matedEntities[0].matedCS.part_to_mate_transform
+                    joints.append(
+                        FixedJoint(
+                            name=f"joint_{edge[0]}_{edge[1]}",
+                            parent=edge[0],
+                            child=edge[1],
+                            origin=Origin.from_matrix(mates[_mate_key].matedEntities[1].matedCS.part_to_mate_transform),
+                        )
+                    )
+
+    for node in _sorted_nodes:
+        _link_to_stl_transform = np.eye(4)
+        _mass_property = mass_properties.get(node)
+        _mass = _mass_property.mass[0]
+        _default_origin = Origin.zero_origin()
+
+        if node in _child_transforms:
+            _link_to_stl_transform = _child_transforms[node]
+        else:
+            _link_to_stl_transform[:3, 3] = np.array(_mass_property.center_of_mass).reshape(3)
+
+        _stl_to_link_transform = np.matrix(np.linalg.inv(_link_to_stl_transform))
+
+        _center_of_mass = _mass_property.center_of_mass_wrt(_stl_to_link_transform)
+        _inertia_matrix = _mass_property.inertia_wrt(np.matrix(_stl_to_link_transform[:3, :3]))
+        _default_principal_axes = (0, 0, 0)
+
+        stl_path = download_stl_mesh(
+            parts[node].documentId,
+            workspaceId,
+            parts[node].elementId,
+            parts[node].partId,
+            client,
+            f"{node}.stl"
+        )
+
+        links.append(
+            Link(
+                name=node,
+                visual=VisualLink(
+                    origin=_default_origin,
+                    geometry=MeshGeometry(stl_path),
+                    material=Material.from_color(name=f"{node}_material", color=COLORS.RED)
+                ),
+                inertial=InertialLink(
+                    origin=Origin(_center_of_mass, _default_principal_axes),
+                    mass=_mass,
+                    inertia=Inertia(
+                        ixx=_inertia_matrix[0, 0],
+                        iyy=_inertia_matrix[1, 1],
+                        izz=_inertia_matrix[2, 2],
+                        ixy=_inertia_matrix[0, 1],
+                        ixz=_inertia_matrix[0, 2],
+                        iyz=_inertia_matrix[1, 2],
+                    )
+                ),
+                collision=CollisionLink(
+                    origin=_default_origin,
+                    geometry=MeshGeometry(stl_path)
+                )
+            )
+        )
+
+
+    return links, joints
+

onshape_api/models/assembly.py

@@ -4,14 +4,15 @@
 from pydantic import BaseModel, field_validator
 
 from onshape_api.utilities.helpers import generate_uid
+import numpy as np
 
 
 class InstanceType(str, Enum):
     PART = "Part"
     ASSEMBLY = "Assembly"
 
 
-class MateType(str, Enum):
+class MATETYPE(str, Enum):
     SLIDER = "SLIDER"
     CYLINDRICAL = "CYLINDRICAL"
     REVOLUTE = "REVOLUTE"
@@ -222,6 +223,15 @@ def check_vectors(cls, v):
 
         return v
 
+    @property
+    def part_to_mate_transform(self) -> np.matrix:
+        rotation_matrix = np.array([self.xAxis, self.yAxis, self.zAxis]).T
+        translation_vector = np.array(self.origin)
+        part_to_mate_tf = np.eye(4)
+        part_to_mate_tf[:3, :3] = rotation_matrix
+        part_to_mate_tf[:3, 3] = translation_vector
+        return np.matrix(part_to_mate_tf)
+
 
 class MatedEntity(BaseModel):
     """
@@ -274,7 +284,7 @@ class MateFeatureData(BaseModel):
     """
 
     matedEntities: list[MatedEntity]
-    mateType: MateType
+    mateType: MATETYPE
     name: str
 
 

onshape_api/models/link.py

@@ -102,8 +102,8 @@ def to_xml(self, root: ET.Element | None = None) -> ET.Element:
         return material
 
     @classmethod
-    def from_color(cls, color: COLORS) -> "Material":
-        return cls(color, color)
+    def from_color(cls, name: str, color: COLORS) -> "Material":
+        return cls(name, color)
 
 
 @dataclass
@@ -149,7 +149,7 @@ def to_xml(self, root: ET.Element | None = None) -> ET.Element:
         if self.inertial is not None:
             self.inertial.to_xml(link)
         return link
-    
+
     @classmethod
     def from_part(cls, part):
         _cls = cls(name=part.partId)

onshape_api/models/mass.py

@@ -65,7 +65,7 @@ def principal_inertia(self) -> np.ndarray:
 
     @property
     def center_of_mass(self) -> np.ndarray:
-        return np.array(self.centroid)
+        return np.array(self.centroid[:3])
 
     @property
     def inertia_matrix(self) -> np.matrix: