diff --git a/benchmark/connectors.py b/benchmark/connectors.py index c46f877..ffb2d16 100644 --- a/benchmark/connectors.py +++ b/benchmark/connectors.py @@ -8,9 +8,11 @@ if __name__ == "__main__": client = opa.Client() - # connectors: https://cad.onshape.com/documents/8df4a8934dea6cc8a51a6f85/w/5473d6310f998d61ffe1045e/e/dcf8d3ded2234a1bea3856ba + # connectors witch cycle: https://cad.onshape.com/documents/8df4a8934dea6cc8a51a6f85/w/5473d6310f998d61ffe1045e/e/dcf8d3ded2234a1bea3856ba + # simple connectors: https://cad.onshape.com/documents/2e35965a561baafef08b14bc/w/967b9a19ea54f3cb2703ba2d/e/ebe351dafa0a477efb87912e + # T joint: https://cad.onshape.com/documents/7481cfb61a62765f39b1f8b6/w/7c3a25b7fe7637d71f04f632/e/fbf64c254c6cde3d16c91250 document = Document.from_url( - "https://cad.onshape.com/documents/a6286c1e44b5c2b74af3271f/w/0f4019a5df5bd9f85d8d8c44/e/79a6f37f589bddff34769e6a" + "https://cad.onshape.com/documents/7481cfb61a62765f39b1f8b6/w/7c3a25b7fe7637d71f04f632/e/fbf64c254c6cde3d16c91250" ) assembly, _ = client.get_assembly( did=document.did, diff --git a/benchmark/robots.py b/benchmark/robots.py new file mode 100644 index 0000000..2dd9433 --- /dev/null +++ b/benchmark/robots.py @@ -0,0 +1,52 @@ +import os + +import onshape_api as opa +from onshape_api.models.document import Document + +SCRIPT_DIRECTORY = os.path.dirname(os.path.realpath(__file__)) + +if __name__ == "__main__": + client = opa.Client() + # robot = https://cad.onshape.com/documents/a8f62e825e766a6512320ceb/w/b9099bcbdc92e6d6c810f0b7/e/f5b0475edd5ad0193d280fc4 + + document = Document.from_url( + "https://cad.onshape.com/documents/cf6b852d2c88d661ac2e17e8/w/c842455c29cc878dc48bdc68/e/b5e293d409dd0b88596181ef" + ) + assembly, _ = client.get_assembly( + did=document.did, + wtype=document.wtype, + wid=document.wid, + eid=document.eid, + with_meta_data=True, + ) + + opa.LOGGER.info(assembly.document.url) + assembly_robot_name = f"{assembly.document.name + '-' + assembly.name}" + + instances, id_to_name_map = opa.get_instances(assembly) + occurences = opa.get_occurences(assembly, id_to_name_map) + parts = opa.get_parts(assembly, client, instances) + subassemblies = opa.get_subassemblies(assembly, instances) + mates, relations = opa.get_mates_and_relations(assembly, subassemblies, id_to_name_map) + + graph, root_node = opa.create_graph( + occurences=occurences, + instances=instances, + parts=parts, + mates=mates, + use_user_defined_root=False, + ) + opa.save_graph(graph, f"{assembly_robot_name}.png") + + links, joints = opa.get_urdf_components( + assembly=assembly, + graph=graph, + root_node=root_node, + parts=parts, + mates=mates, + relations=relations, + client=client, + ) + + robot = opa.Robot(name=assembly_robot_name, links=links, joints=joints) + robot.save(f"{assembly_robot_name}.urdf") diff --git a/onshape_api/connect.py b/onshape_api/connect.py index 2d6f59a..118259e 100644 --- a/onshape_api/connect.py +++ b/onshape_api/connect.py @@ -29,7 +29,7 @@ from onshape_api.log import LOG_LEVEL, LOGGER from onshape_api.models.assembly import Assembly -from onshape_api.models.document import Document, DocumentMetaData +from onshape_api.models.document import Document, DocumentMetaData, WorkspaceType, generate_url from onshape_api.models.element import Element from onshape_api.models.mass import MassProperties from onshape_api.models.variable import Variable @@ -467,7 +467,16 @@ def get_assembly( return _assembly, _assembly_json - def download_stl(self, did: str, wid: str, eid: str, partID: str, buffer: BinaryIO) -> BinaryIO: + def download_stl( + self, + did: str, + wid: str, + eid: str, + partID: str, + buffer: BinaryIO, + wtype: str = WorkspaceType.W.value, + vid: Optional[str] = None, + ) -> BinaryIO: """ Download an STL file from a part studio. The file is written to the buffer. @@ -477,6 +486,8 @@ def download_stl(self, did: str, wid: str, eid: str, partID: str, buffer: Binary eid: The unique identifier of the element. partID: The unique identifier of the part. buffer: BinaryIO object to write the STL file to. + wtype: The type of workspace. + vid: The unique identifier of the version workspace. Returns: BinaryIO: BinaryIO object containing the STL file @@ -488,7 +499,9 @@ def download_stl(self, did: str, wid: str, eid: str, partID: str, buffer: Binary ... "0d17b8ebb2a4c76be9fff3c7", ... "a86aaf34d2f4353288df8812", ... "0b0c209535554345432581fe", - ... buffer + ... buffer, + ... "w", + ... "0d17b8ebb2a4c76be9fff3c7" ... ) >>> buffer.seek(0) >>> raw_mesh = stl.mesh.Mesh.from_file(None, fh=buffer) @@ -496,7 +509,9 @@ def download_stl(self, did: str, wid: str, eid: str, partID: str, buffer: Binary """ req_headers = {"Accept": "application/vnd.onshape.v1+octet-stream"} - _request_path = f"/api/parts/d/{did}/w/{wid}/e/{eid}/partid/{partID}/stl" + _request_path = ( + f"/api/parts/d/{did}/{wtype}/" f"{wid if wtype == WorkspaceType.W else vid}/e/{eid}/partid/{partID}/stl" + ) _query = { "mode": "binary", "grouping": True, @@ -511,12 +526,38 @@ def download_stl(self, did: str, wid: str, eid: str, partID: str, buffer: Binary ) if response.status_code == 200: buffer.write(response.content) + elif response.status_code == 404 or response.status_code == 400: + if vid and wtype == WorkspaceType.W: + return self.download_stl(did, wid, eid, partID, buffer, WorkspaceType.V.value, vid) + else: + LOGGER.info(f"{ + generate_url( + did=did, + wtype="w", + wid=wid, + eid=eid, + ) + }") + LOGGER.info( + f"No version ID provided, failed to download STL file: {response.status_code} - {response.text}" + ) + else: + LOGGER.info(f"{ + generate_url( + did=did, + wtype="w", + wid=wid, + eid=eid, + ) + }") LOGGER.info(f"Failed to download STL file: {response.status_code} - {response.text}") return buffer - def get_mass_property(self, did: str, wid: str, eid: str, partID: str) -> MassProperties: + def get_mass_property( + self, did: str, wid: str, eid: str, partID: str, vid: Optional[str], wtype: str = WorkspaceType.W.value + ) -> MassProperties: """ Get mass properties of a part in a part studio. @@ -525,6 +566,8 @@ def get_mass_property(self, did: str, wid: str, eid: str, partID: str) -> MassPr wid: The unique identifier of the workspace. eid: The unique identifier of the element. partID: The identifier of the part. + vid: The unique identifier of the document version. + wtype: The type of workspace. Returns: MassProperties object containing the mass properties of the part. @@ -535,6 +578,8 @@ def get_mass_property(self, did: str, wid: str, eid: str, partID: str) -> MassPr ... wid="0d17b8ebb2a4c76be9fff3c7", ... eid="a86aaf34d2f4353288df8812", ... partID="0b0c209535554345432581fe" + ... vid="0d17bae7b2a4c76be9fff3c7", + ... wtype="w" ... ) >>> print(mass_properties) MassProperties( @@ -546,12 +591,27 @@ def get_mass_property(self, did: str, wid: str, eid: str, partID: str) -> MassPr principalAxes=[...] ) """ - _request_path = "/api/parts/d/" + did + "/w/" + wid + "/e/" + eid + "/partid/" + partID + "/massproperties" + _request_path = ( + f"/api/parts/d/{did}/{wtype}/" + f"{wid if wtype == WorkspaceType.W else vid}/e/{eid}/partid/{partID}/massproperties" + ) res = self.request(HTTP.GET, _request_path, {"useMassPropertiesOverrides": True}) if res.status_code == 404: # TODO: There doesn't seem to be a way to assign material to a part currently - raise ValueError(f"Part: {partID} does not have a material assigned") + # It is possible that the workspace got deleted + if vid and wtype == WorkspaceType.W: + print("Trying to get mass properties from a version workspace") + return self.get_mass_property(did, wid, eid, partID, vid, WorkspaceType.V.value) + + raise ValueError(f"Part: { + generate_url( + did=did, + wtype="w", + wid=wid, + eid=eid, + ) + } does not have a material assigned or the part is not found") _resonse_json = res.json() diff --git a/onshape_api/graph.py b/onshape_api/graph.py index 4eabc7e..4bd42dc 100644 --- a/onshape_api/graph.py +++ b/onshape_api/graph.py @@ -98,6 +98,7 @@ def convert_to_digraph(graph: nx.Graph, user_defined_root: Union[str, None] = No >>> convert_to_digraph(graph) (digraph, root_node) """ + centrality = nx.closeness_centrality(graph) root_node = user_defined_root if user_defined_root else max(centrality, key=centrality.get) @@ -166,13 +167,50 @@ def create_graph( >>> create_graph(occurences, instances, parts, mates, directed=True) """ - graph: nx.Graph = nx.Graph() - root_node = None - user_defined_root = None + graph = nx.Graph() + user_defined_root = add_nodes_to_graph(graph, occurences, instances, parts, use_user_defined_root) + add_edges_to_graph(graph, mates) + + cur_graph = remove_unconnected_subgraphs(graph) + + if directed: + output_graph, root_node = convert_to_digraph(cur_graph, user_defined_root) + else: + output_graph = cur_graph + root_node = None + + LOGGER.info( + f"Graph created with {len(output_graph.nodes)} nodes and " + f"{len(output_graph.edges)} edges with root node: {root_node}" + ) + + return output_graph, root_node + + +def add_nodes_to_graph( + graph: nx.Graph, + occurences: dict[str, Occurrence], + instances: dict[str, Union[PartInstance, AssemblyInstance]], + parts: dict[str, Part], + use_user_defined_root: bool, +) -> str: + """ + Add nodes to the graph. + + Args: + graph: The graph to add nodes to. + occurences: Dictionary of occurrences in the assembly. + instances: Dictionary of instances in the assembly. + parts: Dictionary of parts in the assembly. + use_user_defined_root: Whether to use the user defined root node. - if len(mates) == 0: - raise ValueError("No mates found in assembly") + Returns: + The user defined root node if it exists. + Examples: + >>> add_nodes_to_graph(graph, occurences, instances, parts, use_user_defined_root=True) + """ + user_defined_root = None for occurence in occurences: if use_user_defined_root and occurences[occurence].fixed: user_defined_root = occurence @@ -185,7 +223,20 @@ def create_graph( graph.add_node(occurence, **parts[occurence].model_dump()) except KeyError: LOGGER.warning(f"Part {occurence} not found") + return user_defined_root + + +def add_edges_to_graph(graph: nx.Graph, mates: dict[str, MateFeatureData]) -> None: + """ + Add edges to the graph. + + Args: + graph: The graph to add edges to. + mates: Dictionary of mates in the assembly. + Examples: + >>> add_edges_to_graph(graph, mates) + """ for mate in mates: try: child, parent = mate.split(MATE_JOINER) @@ -193,9 +244,20 @@ def create_graph( except KeyError: LOGGER.warning(f"Mate {mate} not found") - if directed: - graph, root_node = convert_to_digraph(graph, user_defined_root) - LOGGER.info(f"Graph created with {len(graph.nodes)} nodes and {len(graph.edges)} edges with root node: {root_node}") +def remove_unconnected_subgraphs(graph: nx.Graph) -> nx.Graph: + """ + Remove unconnected subgraphs from the graph. - return graph, root_node + Args: + graph: The graph to remove unconnected subgraphs from. + + Returns: + The main connected subgraph of the graph, which is the largest connected subgraph. + """ + if not nx.is_connected(graph): + sub_graphs = list(nx.connected_components(graph)) + main_graph_nodes = max(sub_graphs, key=len) + main_graph = graph.subgraph(main_graph_nodes).copy() + return main_graph + return graph diff --git a/onshape_api/models/assembly.py b/onshape_api/models/assembly.py index 959399b..d0004c8 100644 --- a/onshape_api/models/assembly.py +++ b/onshape_api/models/assembly.py @@ -331,6 +331,7 @@ class Part(IDBase): isStandardContent: bool = Field(..., description="Indicates if the part is standard content.") partId: str = Field(..., description="The unique identifier of the part.") bodyType: str = Field(..., description="The type of the body (e.g., solid, surface).") + documentVersion: str = Field(None, description="The version of the document.") MassProperty: Union[MassProperties, None] = Field( None, description="The mass properties of the part, this is a retrieved via a separate API call." ) @@ -403,6 +404,7 @@ class PartInstance(IDBase): isStandardContent: bool = Field(..., description="Indicates if the part is standard content.") type: InstanceType = Field(..., description="The type of the instance, must be 'Part'.") + documentVersion: str = Field(None, description="The version of the document.") id: str = Field(..., description="The unique identifier for the part instance.") name: str = Field(..., description="The name of the part instance.") suppressed: bool = Field(..., description="Indicates if the part instance is suppressed.") diff --git a/onshape_api/parse.py b/onshape_api/parse.py index 53a6432..f4592e1 100644 --- a/onshape_api/parse.py +++ b/onshape_api/parse.py @@ -231,7 +231,11 @@ def get_parts( if part.uid in part_instance_map: for key in part_instance_map[part.uid]: part.MassProperty = client.get_mass_property( - part.documentId, assembly.document.wid, part.elementId, part.partId + did=part.documentId, + wid=assembly.document.wid, + eid=part.elementId, + partID=part.partId, + vid=part.documentVersion, ) part_map[key] = part @@ -344,14 +348,11 @@ def traverse_assembly( # TODO: Verify mate relation convention child_joint_id = feature.featureData.mates[RELATION_CHILD].featureId - # if feature.featureData.relationType == RelationType.GEAR: - # parent_joint_id = feature.featureData.mates[RELATION_PARENT].featureId - # _relations_map[parent_joint_id] = feature.featureData - # _relations_map[parent_joint_id].relationRatio = 1 / feature.featureData.relationRatio - _relations_map[child_joint_id] = feature.featureData elif feature.featureType == AssemblyFeatureType.MATECONNECTOR: + # Mate connectors' MatedCS data is already included in the MateFeatureData + # TODO: This might not be true for all cases? pass return _mates_map, _relations_map diff --git a/onshape_api/urdf.py b/onshape_api/urdf.py index 223f938..f0f25ca 100644 --- a/onshape_api/urdf.py +++ b/onshape_api/urdf.py @@ -52,7 +52,7 @@ def download_stl_mesh( - did: str, wid: str, eid: str, partID: str, client: Client, transform: np.ndarray, file_name: str + did: str, wid: str, eid: str, partID: str, vid: str, client: Client, transform: np.ndarray, file_name: str ) -> str: """ Download an STL mesh from an Onshape part studio, transform it, and save it to a file. @@ -62,6 +62,7 @@ def download_stl_mesh( wid: The unique identifier of the workspace. eid: The unique identifier of the element. partID: The unique identifier of the part. + vid: The unique identifier of the version of the document. client: The Onshape client object to use for sending API requests. transform: The transformation matrix to apply to the mesh. file_name: The name of the file to save the mesh to. @@ -73,15 +74,23 @@ def download_stl_mesh( Exception: If an error occurs while downloading the mesh. Examples: - >>> download_stl_mesh("0b0c209535554345432581fe", "0b0c209535554345432581fe", "0b0c209535554345432581fe", - ... "0b0c209535554345432581fe", client, np.eye(4), "part.stl") + >>> download_stl_mesh( + ... did="document_id", + ... wid="workspace_id", + ... eid="element_id", + ... partID="part_id", + ... vid="version_id", + ... client=client, + ... transform=np.eye(4), + ... file_name="part.stl" + ... ) "meshes/part.stl" """ try: with io.BytesIO() as buffer: LOGGER.info(f"Downloading mesh for {file_name}...") - client.download_stl(did, wid, eid, partID, buffer) + client.download_stl(did=did, wid=wid, eid=eid, partID=partID, buffer=buffer, vid=vid) buffer.seek(0) raw_mesh = stl.mesh.Mesh.from_file(None, fh=buffer) @@ -154,13 +163,14 @@ def get_robot_link( LOGGER.info(f"Creating robot link for {name}") _mesh_path = download_stl_mesh( - part.documentId, - wid, - part.elementId, - part.partId, - client, - _stl_to_link_tf, - f"{name}.stl", + did=part.documentId, + wid=wid, + eid=part.elementId, + partID=part.partId, + vid=part.documentVersion, + client=client, + transform=_stl_to_link_tf, + file_name=f"{name}.stl", ) _link = Link( @@ -374,7 +384,9 @@ def get_urdf_components( LOGGER.info(f"Processing root node: {root_node}") - root_link, stl_to_root_tf = get_robot_link(root_node, parts[root_node], assembly.document.wid, client, None) + root_link, stl_to_root_tf = get_robot_link( + name=root_node, part=parts[root_node], wid=assembly.document.wid, client=client, mate=None + ) links.append(root_link) stl_to_link_tf_map[root_node] = stl_to_root_tf @@ -387,7 +399,12 @@ def get_urdf_components( mate_key = f"{parent}{MATE_JOINER}{child}" LOGGER.info(f"Processing edge: {parent} -> {child}") - parent_tf = stl_to_link_tf_map[parent] + try: + parent_tf = stl_to_link_tf_map[parent] + except KeyError: + LOGGER.warning(f"Parent {parent} not found in stl_to_link_tf_map") + LOGGER.info(f"stl_to_link_tf_map keys: {stl_to_link_tf_map.keys()}") + exit(1) if parent and child not in parts: LOGGER.warning(f"Part {parent} or {child} not found in parts")