sunxianghui
93773f3887
Includes FastAPI backend, vendored step2urdf frontend, and A7 handtuned arm JSON for URDF generation.
103 lines
3.8 KiB
Python
103 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import xml.etree.ElementTree as ET
|
|
from pathlib import Path
|
|
from xml.dom import minidom
|
|
|
|
from app.models import RobotDraft
|
|
|
|
|
|
def _fmt(nums: list[float]) -> str:
|
|
return " ".join(f"{float(x):.6g}" for x in nums)
|
|
|
|
|
|
def build_urdf_xml(draft: RobotDraft, mesh_dir_uri: str = "package://robot_description/meshes") -> str:
|
|
robot = ET.Element("robot", name=draft.name)
|
|
|
|
for link in draft.links:
|
|
link_el = ET.SubElement(robot, "link", name=link.name)
|
|
visual = ET.SubElement(link_el, "visual")
|
|
ET.SubElement(visual, "origin", xyz="0 0 0", rpy="0 0 0")
|
|
geom = ET.SubElement(visual, "geometry")
|
|
mesh_name = link.mesh or f"{link.name}.stl"
|
|
ET.SubElement(geom, "mesh", filename=f"{mesh_dir_uri}/{mesh_name}", scale="0.001 0.001 0.001")
|
|
coll = ET.SubElement(link_el, "collision")
|
|
ET.SubElement(coll, "origin", xyz="0 0 0", rpy="0 0 0")
|
|
cgeom = ET.SubElement(coll, "geometry")
|
|
ET.SubElement(cgeom, "mesh", filename=f"{mesh_dir_uri}/{mesh_name}", scale="0.001 0.001 0.001")
|
|
inertial = ET.SubElement(link_el, "inertial")
|
|
ET.SubElement(inertial, "origin", xyz="0 0 0", rpy="0 0 0")
|
|
ET.SubElement(inertial, "mass", value="1.0")
|
|
ET.SubElement(
|
|
inertial,
|
|
"inertia",
|
|
ixx="0.001",
|
|
ixy="0",
|
|
ixz="0",
|
|
iyy="0.001",
|
|
iyz="0",
|
|
izz="0.001",
|
|
)
|
|
|
|
for joint in draft.joints:
|
|
j_el = ET.SubElement(robot, "joint", name=joint.name, type=joint.joint_type)
|
|
ET.SubElement(j_el, "parent", link=joint.parent)
|
|
ET.SubElement(j_el, "child", link=joint.child)
|
|
ET.SubElement(j_el, "origin", xyz=_fmt(joint.origin_xyz), rpy=_fmt(joint.origin_rpy))
|
|
if joint.joint_type in {"revolute", "prismatic", "continuous"}:
|
|
ET.SubElement(j_el, "axis", xyz=_fmt(joint.axis))
|
|
if joint.joint_type in {"revolute", "prismatic"}:
|
|
ET.SubElement(
|
|
j_el,
|
|
"limit",
|
|
lower=str(joint.lower),
|
|
upper=str(joint.upper),
|
|
effort=str(joint.effort),
|
|
velocity=str(joint.velocity),
|
|
)
|
|
|
|
rough = ET.tostring(robot, encoding="utf-8")
|
|
pretty = minidom.parseString(rough).toprettyxml(indent=" ")
|
|
# drop extra XML declaration newline quirks
|
|
lines = [ln for ln in pretty.splitlines() if ln.strip()]
|
|
return "\n".join(lines) + "\n"
|
|
|
|
|
|
def write_urdf_package(draft: RobotDraft, out_dir: Path) -> dict[str, str]:
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
meshes = out_dir / "meshes"
|
|
meshes.mkdir(exist_ok=True)
|
|
urdf_path = out_dir / f"{draft.name}.urdf"
|
|
urdf_xml = build_urdf_xml(draft, mesh_dir_uri="meshes")
|
|
urdf_path.write_text(urdf_xml, encoding="utf-8")
|
|
|
|
# Placeholder note for missing meshes
|
|
readme = out_dir / "README.txt"
|
|
readme.write_text(
|
|
"\n".join(
|
|
[
|
|
f"Robot: {draft.name}",
|
|
f"Profile: {draft.profile}",
|
|
f"Links: {len(draft.links)}",
|
|
f"Joints: {len(draft.joints)}",
|
|
"",
|
|
"Mesh files are expected under ./meshes as <link>.stl.",
|
|
"If meshes were not generated, URDF still validates structurally;",
|
|
"install OCP/cadquery or cascadio later for STEP tessellation.",
|
|
"",
|
|
"Notes:",
|
|
*[f"- {n}" for n in draft.notes],
|
|
]
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
draft_json = out_dir / "draft.json"
|
|
draft_json.write_text(draft.model_dump_json(indent=2), encoding="utf-8")
|
|
return {
|
|
"urdf": str(urdf_path),
|
|
"readme": str(readme),
|
|
"draft_json": str(draft_json),
|
|
"meshes_dir": str(meshes),
|
|
}
|