Initial commit: step2urdf tool with handtuned JSON import.
Includes FastAPI backend, vendored step2urdf frontend, and A7 handtuned arm JSON for URDF generation.
This commit is contained in:
sunxianghui
@@ -0,0 +1,120 @@
|
||||
"""A7 / ARM7 gold kinematics from lkls73_i1_arm_description/urdf/ARM7_urdf.urdf.
|
||||
|
||||
STEP Viewer / step2urdf UI stores linear joint origins in **millimetres**;
|
||||
export applies ×0.001 → metres. Values below are therefore in mm so they match
|
||||
the interactive UI and still export to the gold URDF metre numbers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.models import JointDef, LinkDef, RobotDraft
|
||||
|
||||
# UI / CAD units (mm). Gold URDF metres × 1000.
|
||||
A7_LINK_ORDER = [
|
||||
"base_link",
|
||||
"A1_Link",
|
||||
"A2_Link",
|
||||
"A3_Link",
|
||||
"A4_Link",
|
||||
"A5_Link",
|
||||
"A6_Link",
|
||||
"A7_Link",
|
||||
"A8_Link",
|
||||
]
|
||||
|
||||
# (name, parent, child, type, xyz_mm, axis, lower, upper, effort, velocity)
|
||||
A7_GOLD_JOINTS_MM: list[tuple] = [
|
||||
("A1_joint", "base_link", "A1_Link", "revolute", [0.0, 0.0, 9.0], [0.0, 0.0, -1.0], -2.9, 1.0, 80.0, 10.47),
|
||||
("A2_joint", "A1_Link", "A2_Link", "revolute", [0.0, 0.0, 59.5], [1.0, 0.0, 0.0], -0.15, 3.14, 80.0, 10.47),
|
||||
("A3_joint", "A2_Link", "A3_Link", "revolute", [0.0, 0.0, 123.0], [0.0, 0.0, -1.0], -2.35, 2.35, 36.0, 9.1),
|
||||
("A4_joint", "A3_Link", "A4_Link", "revolute", [0.0, -0.10871, 187.1], [1.0, 0.0, 0.0], 0.0, 2.2, 36.0, 9.1),
|
||||
("A5_joint", "A4_Link", "A5_Link", "revolute", [0.0, -0.13906, 124.0], [0.0, 0.0, -1.0], -2.35, 2.35, 36.0, 9.1),
|
||||
("A6_joint", "A5_Link", "A6_Link", "revolute", [0.0, -0.20755, 88.099], [1.0, 0.0, 0.0], -1.57, 1.57, 36.0, 9.1),
|
||||
("A7_joint", "A6_Link", "A7_Link", "revolute", [0.60033, 0.45533, 65.501], [0.0, 1.0, 0.0], -1.57, 1.57, 36.0, 9.1),
|
||||
("A8_joint", "A7_Link", "A8_Link", "fixed", [-0.041388, -1.1787, 44.999], [0.0, 0.0, 0.0], 0.0, 0.0, 0.0, 0.0),
|
||||
]
|
||||
|
||||
|
||||
def gold_joint_defs() -> list[JointDef]:
|
||||
out: list[JointDef] = []
|
||||
for name, parent, child, jtype, xyz, axis, lo, hi, eff, vel in A7_GOLD_JOINTS_MM:
|
||||
out.append(
|
||||
JointDef(
|
||||
name=name,
|
||||
joint_type=jtype, # type: ignore[arg-type]
|
||||
parent=parent,
|
||||
child=child,
|
||||
origin_xyz=list(xyz),
|
||||
origin_rpy=[0.0, 0.0, 0.0],
|
||||
axis=list(axis),
|
||||
lower=lo,
|
||||
upper=hi,
|
||||
effort=eff,
|
||||
velocity=vel,
|
||||
rationale="ARM7_urdf.urdf gold kinematics (origins in mm for step2urdf UI).",
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def empty_gold_links() -> list[LinkDef]:
|
||||
return [LinkDef(name=n, part_names=[]) for n in A7_LINK_ORDER]
|
||||
|
||||
|
||||
def enforce_a7_gold_kinematics(draft: RobotDraft) -> RobotDraft:
|
||||
"""Keep part_names from draft when link names match; always replace joints with gold.
|
||||
|
||||
Also remaps common aliases (link1 → A1_Link, etc.) so LLM/profile seeds still work.
|
||||
"""
|
||||
alias = {
|
||||
"link1": "A1_Link",
|
||||
"link2": "A2_Link",
|
||||
"link3": "A3_Link",
|
||||
"link4": "A4_Link",
|
||||
"link5": "A5_Link",
|
||||
"link6": "A6_Link",
|
||||
"link7": "A7_Link",
|
||||
"link_1": "A1_Link",
|
||||
"link_2": "A2_Link",
|
||||
"link_3": "A3_Link",
|
||||
"link_4": "A4_Link",
|
||||
"link_5": "A5_Link",
|
||||
"link_6": "A6_Link",
|
||||
"link_7": "A7_Link",
|
||||
"a1_link": "A1_Link",
|
||||
"a2_link": "A2_Link",
|
||||
"a3_link": "A3_Link",
|
||||
"a4_link": "A4_Link",
|
||||
"a5_link": "A5_Link",
|
||||
"a6_link": "A6_Link",
|
||||
"a7_link": "A7_Link",
|
||||
"a8_link": "A8_Link",
|
||||
}
|
||||
|
||||
buckets: dict[str, list[str]] = {n: [] for n in A7_LINK_ORDER}
|
||||
for link in draft.links:
|
||||
key = alias.get(link.name, alias.get(link.name.lower(), link.name))
|
||||
if key not in buckets:
|
||||
# Unknown link → fold into base
|
||||
key = "base_link"
|
||||
for pn in link.part_names:
|
||||
if pn not in buckets[key]:
|
||||
buckets[key].append(pn)
|
||||
|
||||
links = [LinkDef(name=n, part_names=buckets[n]) for n in A7_LINK_ORDER]
|
||||
joints = gold_joint_defs()
|
||||
notes = list(draft.notes) + [
|
||||
"A7 关节轴/原点已强制替换为 ARM7 金标（原点单位 mm，导出时 ×0.001 → m）。",
|
||||
"智谱仅用于零件归组；请勿依赖模型猜测的 axis/rpy。",
|
||||
"轴规律: A1/A3/A5 → [0,0,-1]; A2/A4/A6 → [1,0,0]; A7 → [0,1,0]; A8 fixed。",
|
||||
]
|
||||
return RobotDraft(
|
||||
name=draft.name or "ARM7_urdf",
|
||||
profile="a7",
|
||||
links=links,
|
||||
joints=joints,
|
||||
notes=notes,
|
||||
raw_parts=draft.raw_parts,
|
||||
axis_candidates=draft.axis_candidates,
|
||||
llm_raw=draft.llm_raw,
|
||||
)
|
||||
@@ -0,0 +1,74 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
def geometry_backend_status() -> dict[str, Any]:
|
||||
status = {"ocp": False, "cascadio": False, "available": False, "detail": []}
|
||||
try:
|
||||
import OCP # noqa: F401
|
||||
|
||||
status["ocp"] = True
|
||||
status["detail"].append("OCP import OK")
|
||||
except Exception as e: # noqa: BLE001
|
||||
status["detail"].append(f"OCP unavailable: {e}")
|
||||
|
||||
try:
|
||||
import cascadio # noqa: F401
|
||||
|
||||
status["cascadio"] = True
|
||||
status["detail"].append("cascadio import OK")
|
||||
except Exception as e: # noqa: BLE001
|
||||
status["detail"].append(f"cascadio unavailable: {e}")
|
||||
|
||||
status["available"] = bool(status["ocp"] or status["cascadio"])
|
||||
return status
|
||||
|
||||
|
||||
def export_meshes_stub(step_path: Path, meshes_dir: Path, link_names: list[str]) -> dict[str, Any]:
|
||||
"""MVP mesh export.
|
||||
|
||||
Full per-link tessellation needs OCP/pythonocc. Until then we record
|
||||
intended mesh filenames and leave placeholders.
|
||||
"""
|
||||
meshes_dir.mkdir(parents=True, exist_ok=True)
|
||||
status = geometry_backend_status()
|
||||
created: list[str] = []
|
||||
pending: list[str] = []
|
||||
|
||||
if status["cascadio"]:
|
||||
try:
|
||||
import cascadio
|
||||
|
||||
# Whole-assembly export as a starting point (not per-link).
|
||||
out = meshes_dir / "assembly.glb"
|
||||
# cascadio API: step_to_glb(in, out) in recent versions
|
||||
if hasattr(cascadio, "step_to_glb"):
|
||||
cascadio.step_to_glb(str(step_path), str(out))
|
||||
created.append(str(out))
|
||||
else:
|
||||
pending.append("cascadio installed but step_to_glb API not found")
|
||||
except Exception as e: # noqa: BLE001
|
||||
pending.append(f"cascadio export failed: {e}")
|
||||
|
||||
for name in link_names:
|
||||
target = meshes_dir / f"{name}.stl"
|
||||
if not target.exists():
|
||||
# Tiny valid-ish ASCII STL placeholder so packages are non-empty.
|
||||
target.write_text(
|
||||
"solid placeholder\n"
|
||||
" facet normal 0 0 1\n"
|
||||
" outer loop\n"
|
||||
" vertex 0 0 0\n"
|
||||
" vertex 1 0 0\n"
|
||||
" vertex 0 1 0\n"
|
||||
" endloop\n"
|
||||
" endfacet\n"
|
||||
"endsolid placeholder\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
created.append(str(target))
|
||||
pending.append(f"{name}.stl is a placeholder triangle (replace after OCCT tessellation)")
|
||||
|
||||
return {"backend": status, "created": created, "pending": pending}
|
||||
@@ -0,0 +1,119 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from app.config import settings
|
||||
from app.models import ParseResult, RobotDraft
|
||||
from app.services.a7_gold import enforce_a7_gold_kinematics
|
||||
from app.services.profiles import apply_profile
|
||||
from app.services.step_text import parse_step_products
|
||||
from app.services.urdf_builder import write_urdf_package
|
||||
from app.services.zhipu import propose_joints_with_zhipu
|
||||
from app.services.mesh_export import export_meshes_stub, geometry_backend_status
|
||||
|
||||
|
||||
def _job_dir(job_id: str) -> Path:
|
||||
return settings.jobs_dir / job_id
|
||||
|
||||
|
||||
def create_job_from_upload(src: Path, filename: str) -> ParseResult:
|
||||
job_id = uuid.uuid4().hex[:12]
|
||||
jdir = _job_dir(job_id)
|
||||
jdir.mkdir(parents=True, exist_ok=True)
|
||||
dest = jdir / filename
|
||||
shutil.copy2(src, dest)
|
||||
|
||||
parsed = parse_step_products(dest)
|
||||
result = ParseResult(
|
||||
job_id=job_id,
|
||||
filename=filename,
|
||||
root_name=parsed["root_name"],
|
||||
parts=parsed["parts"],
|
||||
product_names=parsed["product_names"],
|
||||
stats={**parsed["stats"], "geometry": geometry_backend_status()},
|
||||
)
|
||||
(jdir / "parse.json").write_text(result.model_dump_json(indent=2), encoding="utf-8")
|
||||
(jdir / "step_path.txt").write_text(str(dest), encoding="utf-8")
|
||||
return result
|
||||
|
||||
|
||||
def load_parse(job_id: str) -> ParseResult:
|
||||
path = _job_dir(job_id) / "parse.json"
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"job not found: {job_id}")
|
||||
return ParseResult.model_validate_json(path.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def step_path_for_job(job_id: str) -> Path:
|
||||
txt = _job_dir(job_id) / "step_path.txt"
|
||||
return Path(txt.read_text(encoding="utf-8").strip())
|
||||
|
||||
|
||||
async def propose_draft(
|
||||
job_id: str | None,
|
||||
*,
|
||||
profile: str,
|
||||
robot_name: str,
|
||||
extra_hint: str = "",
|
||||
part_names: list[str] | None = None,
|
||||
solids_geom: list[dict[str, Any]] | None = None,
|
||||
) -> RobotDraft:
|
||||
"""A7 → ARM7 gold URDF (no LLM). generic → Zhipu full propose."""
|
||||
stats: dict[str, Any] = {}
|
||||
names: list[str] = []
|
||||
|
||||
if part_names:
|
||||
names = [n.strip() for n in part_names if n and str(n).strip()]
|
||||
stats = {"source": "client", "part_count": len(names)}
|
||||
elif job_id:
|
||||
parsed = load_parse(job_id)
|
||||
names = list(parsed.product_names)
|
||||
stats = dict(parsed.stats)
|
||||
else:
|
||||
raise ValueError("Provide part_names (preferred) or job_id")
|
||||
|
||||
if not names:
|
||||
raise ValueError("No part names available for proposal")
|
||||
|
||||
seed = apply_profile(profile, names, robot_name=robot_name)
|
||||
|
||||
# A7: always return gold kinematics (skip LLM — avoids JSON truncation).
|
||||
if profile == "a7":
|
||||
draft = enforce_a7_gold_kinematics(seed)
|
||||
draft.notes = list(draft.notes) + [
|
||||
"已使用 ARM7_urdf.urdf 金标关节（未调用智谱）。前端会按模型最长轴摆放。"
|
||||
]
|
||||
else:
|
||||
draft = await propose_joints_with_zhipu(
|
||||
seed,
|
||||
extra_hint=extra_hint,
|
||||
stats=stats,
|
||||
solids_geom=solids_geom,
|
||||
)
|
||||
|
||||
if job_id:
|
||||
out = _job_dir(job_id) / "draft.json"
|
||||
out.write_text(draft.model_dump_json(indent=2), encoding="utf-8")
|
||||
return draft
|
||||
|
||||
|
||||
def export_package(job_id: str, draft: RobotDraft, include_meshes: bool = True) -> dict[str, Any]:
|
||||
jdir = _job_dir(job_id)
|
||||
out_dir = jdir / "export" / draft.name
|
||||
if out_dir.exists():
|
||||
shutil.rmtree(out_dir)
|
||||
paths = write_urdf_package(draft, out_dir)
|
||||
mesh_info: dict[str, Any] = {}
|
||||
if include_meshes:
|
||||
mesh_info = export_meshes_stub(
|
||||
step_path_for_job(job_id),
|
||||
Path(paths["meshes_dir"]),
|
||||
[l.name for l in draft.links],
|
||||
)
|
||||
meta = {"files": paths, "meshes": mesh_info, "job_id": job_id}
|
||||
(out_dir / "export_meta.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
return meta
|
||||
@@ -0,0 +1,54 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from app.models import LinkDef, RobotDraft
|
||||
from app.services.a7_gold import A7_LINK_ORDER, enforce_a7_gold_kinematics, gold_joint_defs
|
||||
|
||||
|
||||
def _norm(name: str) -> str:
|
||||
return name.strip()
|
||||
|
||||
|
||||
def group_links_a7(product_names: list[str], robot_name: str = "ARM7_urdf") -> RobotDraft:
|
||||
"""ARM7 gold kinematics; part_names filled later (Solid_N clustered on client)."""
|
||||
names = [_norm(n) for n in product_names if _norm(n)]
|
||||
# Spread Solid_* across gold links by index so apply has a starting bind.
|
||||
links = [LinkDef(name=n, part_names=[]) for n in A7_LINK_ORDER]
|
||||
solids = [n for n in names if n.lower().startswith("solid")]
|
||||
other = [n for n in names if not n.lower().startswith("solid")]
|
||||
if solids:
|
||||
n_link = len(links)
|
||||
for i, sn in enumerate(solids):
|
||||
links[min(i * n_link // max(len(solids), 1), n_link - 1)].part_names.append(sn)
|
||||
if other:
|
||||
links[0].part_names.extend(other)
|
||||
|
||||
draft = RobotDraft(
|
||||
name=robot_name or "ARM7_urdf",
|
||||
profile="a7",
|
||||
links=links,
|
||||
joints=gold_joint_defs(),
|
||||
notes=[
|
||||
"A7 profile: joints from ARM7_urdf.urdf gold standard.",
|
||||
"Axes: A1/3/5=[0,0,-1], A2/4/6=[1,0,0], A7=[0,1,0], A8=fixed.",
|
||||
],
|
||||
raw_parts=names,
|
||||
)
|
||||
return enforce_a7_gold_kinematics(draft)
|
||||
|
||||
|
||||
def seed_generic(product_names: list[str], robot_name: str = "robot") -> RobotDraft:
|
||||
names = [_norm(n) for n in product_names if _norm(n)]
|
||||
return RobotDraft(
|
||||
name=robot_name or "robot",
|
||||
profile="generic",
|
||||
links=[LinkDef(name="base_link", part_names=[])],
|
||||
joints=[],
|
||||
notes=["Generic seed; joints come from LLM."],
|
||||
raw_parts=names,
|
||||
)
|
||||
|
||||
|
||||
def apply_profile(profile: str, product_names: list[str], robot_name: str) -> RobotDraft:
|
||||
if profile == "a7":
|
||||
return group_links_a7(product_names, robot_name=robot_name)
|
||||
return seed_generic(product_names, robot_name=robot_name)
|
||||
@@ -0,0 +1,149 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
from app.models import PartNode
|
||||
|
||||
|
||||
_X2_RE = re.compile(r"\\X2\\([0-9A-Fa-f]+)\\X0\\")
|
||||
_PRODUCT_RE = re.compile(
|
||||
r"#(\d+)\s*=\s*PRODUCT\s*\(\s*'([^']*)'\s*,\s*'([^']*)'\s*,\s*'([^']*)'",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
# PRODUCT_DEFINITION_FORMATION(... , #product, ...)
|
||||
_PDF_RE = re.compile(
|
||||
r"#(\d+)\s*=\s*PRODUCT_DEFINITION_FORMATION\s*\([^;]*?#(\d+)\s*\)\s*;",
|
||||
re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
# PRODUCT_DEFINITION(..., #pdf, ...)
|
||||
_PD_RE = re.compile(
|
||||
r"#(\d+)\s*=\s*PRODUCT_DEFINITION\s*\([^;]*?#(\d+)\s*\)\s*;",
|
||||
re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
# NEXT_ASSEMBLY_USAGE_OCCURRENCE('name', ..., #parent_pd, #child_pd, ...)
|
||||
_NAUO_RE = re.compile(
|
||||
r"#(\d+)\s*=\s*NEXT_ASSEMBLY_USAGE_OCCURRENCE\s*\(\s*'([^']*)'\s*,"
|
||||
r"[^;]*?#(\d+)\s*,\s*#(\d+)\s*",
|
||||
re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
|
||||
|
||||
def decode_step_string(s: str) -> str:
|
||||
def repl(m: re.Match[str]) -> str:
|
||||
hexpart = m.group(1)
|
||||
chars: list[str] = []
|
||||
for i in range(0, len(hexpart), 4):
|
||||
chars.append(chr(int(hexpart[i : i + 4], 16)))
|
||||
return "".join(chars)
|
||||
|
||||
return _X2_RE.sub(repl, s)
|
||||
|
||||
|
||||
def parse_step_products(step_path: Path) -> dict:
|
||||
text = step_path.read_text(encoding="utf-8", errors="replace")
|
||||
|
||||
products: dict[str, str] = {}
|
||||
for m in _PRODUCT_RE.finditer(text):
|
||||
eid, name, *_ = m.groups()
|
||||
products[eid] = decode_step_string(name)
|
||||
|
||||
# pdf_id -> product_id
|
||||
pdf_to_product: dict[str, str] = {}
|
||||
for m in _PDF_RE.finditer(text):
|
||||
pdf_id, product_id = m.groups()
|
||||
pdf_to_product[pdf_id] = product_id
|
||||
|
||||
# pd_id -> pdf_id
|
||||
pd_to_pdf: dict[str, str] = {}
|
||||
for m in _PD_RE.finditer(text):
|
||||
pd_id, pdf_id = m.groups()
|
||||
pd_to_pdf[pd_id] = pdf_id
|
||||
|
||||
def pd_name(pd_id: str) -> str:
|
||||
pdf = pd_to_pdf.get(pd_id)
|
||||
if not pdf:
|
||||
return f"pd_{pd_id}"
|
||||
prod = pdf_to_product.get(pdf)
|
||||
if not prod:
|
||||
return f"pdf_{pdf}"
|
||||
return products.get(prod, f"product_{prod}")
|
||||
|
||||
children: dict[str, list[str]] = defaultdict(list)
|
||||
parents: dict[str, str] = {}
|
||||
edge_names: dict[tuple[str, str], str] = {}
|
||||
|
||||
for m in _NAUO_RE.finditer(text):
|
||||
_eid, edge_name, parent_pd, child_pd = m.groups()
|
||||
children[parent_pd].append(child_pd)
|
||||
parents[child_pd] = parent_pd
|
||||
edge_names[(parent_pd, child_pd)] = decode_step_string(edge_name)
|
||||
|
||||
# Roots = PDs that appear as parents or products but have no parent
|
||||
all_pd = set(pd_to_pdf) | set(children) | set(parents)
|
||||
roots = [pd for pd in all_pd if pd not in parents]
|
||||
if not roots and products:
|
||||
# fallback: flat product list
|
||||
nodes = [
|
||||
PartNode(id=pid, name=name, parent_id=None)
|
||||
for pid, name in products.items()
|
||||
]
|
||||
return {
|
||||
"root_name": nodes[0].name if nodes else "assembly",
|
||||
"parts": nodes,
|
||||
"product_names": sorted({p.name for p in nodes}),
|
||||
"stats": {
|
||||
"products": len(products),
|
||||
"nauo": 0,
|
||||
"axis2_placement": text.upper().count("AXIS2_PLACEMENT_3D"),
|
||||
"circles": len(re.findall(r"\bCIRCLE\b", text, flags=re.I)),
|
||||
"cylinders": text.upper().count("CYLINDRICAL_SURFACE"),
|
||||
"mode": "flat_products",
|
||||
},
|
||||
}
|
||||
|
||||
# Prefer the largest tree root as assembly root
|
||||
def subtree_size(pd: str, seen: set[str] | None = None) -> int:
|
||||
seen = seen or set()
|
||||
if pd in seen:
|
||||
return 0
|
||||
seen.add(pd)
|
||||
return 1 + sum(subtree_size(c, seen) for c in children.get(pd, []))
|
||||
|
||||
roots_sorted = sorted(roots, key=subtree_size, reverse=True)
|
||||
root_pd = roots_sorted[0] if roots_sorted else None
|
||||
|
||||
nodes: list[PartNode] = []
|
||||
for pd in sorted(all_pd, key=lambda x: int(x) if x.isdigit() else 0):
|
||||
nodes.append(
|
||||
PartNode(
|
||||
id=pd,
|
||||
name=pd_name(pd),
|
||||
parent_id=parents.get(pd),
|
||||
children=list(children.get(pd, [])),
|
||||
)
|
||||
)
|
||||
|
||||
product_names = sorted({decode_step_string(n) for n in products.values()})
|
||||
root_name = pd_name(root_pd) if root_pd else ""
|
||||
if not root_name or root_name.startswith(("pd_", "pdf_", "product_")):
|
||||
preferred = next((n for n in product_names if "装配体" in n or "assembly" in n.lower()), None)
|
||||
root_name = preferred or (product_names[0] if product_names else "assembly")
|
||||
|
||||
return {
|
||||
"root_name": root_name,
|
||||
"parts": nodes,
|
||||
"product_names": product_names,
|
||||
"stats": {
|
||||
"products": len(products),
|
||||
"product_definitions": len(pd_to_pdf),
|
||||
"nauo": len(edge_names),
|
||||
"roots": len(roots),
|
||||
"axis2_placement": text.upper().count("AXIS2_PLACEMENT_3D"),
|
||||
"circles": len(re.findall(r"\bCIRCLE\b", text, flags=re.I)),
|
||||
"cylinders": text.upper().count("CYLINDRICAL_SURFACE"),
|
||||
"mode": "assembly_tree",
|
||||
"geometry_backend": "text_only",
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
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),
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import settings
|
||||
from app.models import JointDef, LinkDef, RobotDraft
|
||||
|
||||
|
||||
SYSTEM_PROMPT = """你是机器人 URDF 运动学专家。根据零件列表与包围盒，直接设计 link/joint 树。
|
||||
|
||||
约定：
|
||||
- 单位：origin_xyz 用毫米；origin_rpy 用弧度。
|
||||
- 第一个关节相对世界系，origin 靠近对应零件 center，不要写成 [0,0,0]。
|
||||
- part_names 只能用输入里的名字（如 Solid_0），禁止发明新名字。
|
||||
- 每个非 base_link 恰好一个 parent joint。
|
||||
- axis 为单位向量；不要把所有轴都写成 [0,0,1]。
|
||||
- rationale 最多 20 字。
|
||||
- 只输出一个合法 JSON 对象，不要 Markdown、不要注释、不要尾逗号。
|
||||
|
||||
JSON schema:
|
||||
{"robot_name":"string","links":[{"name":"string","part_names":["string"]}],"joints":[{"name":"string","joint_type":"revolute|prismatic|fixed","parent":"string","child":"string","origin_xyz":[0,0,0],"origin_rpy":[0,0,0],"axis":[0,0,1],"lower":-3.14,"upper":3.14,"effort":100,"velocity":1,"rationale":"string"}],"notes":["string"]}
|
||||
"""
|
||||
|
||||
|
||||
def _strip_fences(text: str) -> str:
|
||||
text = text.strip()
|
||||
if text.startswith("```"):
|
||||
text = re.sub(r"^```(?:json)?\s*", "", text)
|
||||
text = re.sub(r"\s*```$", "", text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _repair_json_text(text: str) -> str:
|
||||
"""Best-effort cleanup for common LLM JSON issues / truncation."""
|
||||
text = _strip_fences(text)
|
||||
# Chinese quotes → ASCII
|
||||
text = text.replace("“", '"').replace("”", '"').replace("‘", "'").replace("’", "'")
|
||||
# Remove // comments
|
||||
text = re.sub(r"(?m)^\s*//.*?$", "", text)
|
||||
# Trailing commas before } or ]
|
||||
text = re.sub(r",\s*([}\]])", r"\1", text)
|
||||
|
||||
# If truncated, close open braces/brackets
|
||||
start = text.find("{")
|
||||
if start < 0:
|
||||
return text
|
||||
text = text[start:]
|
||||
# Drop incomplete trailing string if odd number of unescaped quotes in last line
|
||||
in_str = False
|
||||
escape = False
|
||||
stack: list[str] = []
|
||||
last_ok = 0
|
||||
for i, ch in enumerate(text):
|
||||
if in_str:
|
||||
if escape:
|
||||
escape = False
|
||||
elif ch == "\\":
|
||||
escape = True
|
||||
elif ch == '"':
|
||||
in_str = False
|
||||
continue
|
||||
if ch == '"':
|
||||
in_str = True
|
||||
continue
|
||||
if ch in "{[":
|
||||
stack.append("}" if ch == "{" else "]")
|
||||
last_ok = i
|
||||
elif ch in "}]":
|
||||
if stack and stack[-1] == ch:
|
||||
stack.pop()
|
||||
last_ok = i
|
||||
if in_str:
|
||||
# close string then truncate junk after last complete structure
|
||||
text = text + '"'
|
||||
if stack:
|
||||
# trim to last comma / incomplete key if needed
|
||||
text = text.rstrip()
|
||||
if text.endswith(","):
|
||||
text = text[:-1]
|
||||
text = text + "".join(reversed(stack))
|
||||
return text
|
||||
|
||||
|
||||
def _extract_json(text: str) -> dict[str, Any]:
|
||||
candidates = [_strip_fences(text), _repair_json_text(text)]
|
||||
m = re.search(r"\{[\s\S]*\}", text)
|
||||
if m:
|
||||
candidates.append(m.group(0))
|
||||
candidates.append(_repair_json_text(m.group(0)))
|
||||
|
||||
errors: list[str] = []
|
||||
for cand in candidates:
|
||||
try:
|
||||
data = json.loads(cand)
|
||||
if isinstance(data, dict):
|
||||
return data
|
||||
except json.JSONDecodeError as e:
|
||||
errors.append(str(e))
|
||||
raise json.JSONDecodeError(
|
||||
errors[-1] if errors else "Unable to parse LLM JSON",
|
||||
text[:200],
|
||||
0,
|
||||
)
|
||||
|
||||
|
||||
def _compact_solids(solids_geom: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
|
||||
"""Shrink AABB payload so the model reply is less likely to truncate."""
|
||||
out: list[dict[str, Any]] = []
|
||||
for s in solids_geom or []:
|
||||
center = s.get("center")
|
||||
if not center:
|
||||
mn, mx = s.get("min"), s.get("max")
|
||||
if mn and mx and len(mn) == 3 and len(mx) == 3:
|
||||
center = [(mn[i] + mx[i]) / 2 for i in range(3)]
|
||||
if not center:
|
||||
continue
|
||||
out.append(
|
||||
{
|
||||
"name": s.get("name") or s.get("id"),
|
||||
"center": [round(float(x), 1) for x in center[:3]],
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def _draft_from_llm(data: dict[str, Any], fallback: RobotDraft) -> RobotDraft:
|
||||
links = [
|
||||
LinkDef(name=l["name"], part_names=list(l.get("part_names") or []))
|
||||
for l in data.get("links") or []
|
||||
]
|
||||
joints = []
|
||||
for j in data.get("joints") or []:
|
||||
jtype = j.get("joint_type") or "revolute"
|
||||
if jtype not in {"revolute", "prismatic", "fixed", "continuous"}:
|
||||
jtype = "revolute"
|
||||
joints.append(
|
||||
JointDef(
|
||||
name=j.get("name") or "joint",
|
||||
joint_type=jtype, # type: ignore[arg-type]
|
||||
parent=j["parent"],
|
||||
child=j["child"],
|
||||
origin_xyz=[float(x) for x in (j.get("origin_xyz") or [0, 0, 0])],
|
||||
origin_rpy=[float(x) for x in (j.get("origin_rpy") or [0, 0, 0])],
|
||||
axis=[float(x) for x in (j.get("axis") or [0, 0, 1])],
|
||||
lower=float(j.get("lower", -3.14)),
|
||||
upper=float(j.get("upper", 3.14)),
|
||||
effort=float(j.get("effort", 100)),
|
||||
velocity=float(j.get("velocity", 1)),
|
||||
rationale=str(j.get("rationale") or "")[:80],
|
||||
)
|
||||
)
|
||||
if not links:
|
||||
return fallback.model_copy(
|
||||
update={
|
||||
"llm_raw": data,
|
||||
"notes": fallback.notes + ["LLM returned no links; kept seed."],
|
||||
}
|
||||
)
|
||||
return RobotDraft(
|
||||
name=str(data.get("robot_name") or fallback.name),
|
||||
profile=fallback.profile,
|
||||
links=links,
|
||||
joints=joints,
|
||||
notes=list(data.get("notes") or []),
|
||||
raw_parts=fallback.raw_parts,
|
||||
axis_candidates=fallback.axis_candidates,
|
||||
llm_raw=data,
|
||||
)
|
||||
|
||||
|
||||
async def _chat(messages: list[dict[str, str]], *, temperature: float = 0.2) -> str:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {settings.zhipu_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
body = {
|
||||
"model": settings.zhipu_model,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"max_tokens": 8192,
|
||||
}
|
||||
url = f"{settings.zhipu_base_url.rstrip('/')}/chat/completions"
|
||||
async with httpx.AsyncClient(timeout=180.0) as client:
|
||||
resp = await client.post(url, headers=headers, json=body)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
return data["choices"][0]["message"]["content"]
|
||||
|
||||
|
||||
async def propose_joints_with_zhipu(
|
||||
seed: RobotDraft,
|
||||
*,
|
||||
extra_hint: str = "",
|
||||
stats: dict[str, Any] | None = None,
|
||||
solids_geom: list[dict[str, Any]] | None = None,
|
||||
) -> RobotDraft:
|
||||
if not settings.zhipu_api_key:
|
||||
seed.notes = list(seed.notes) + [
|
||||
"ZHIPU_API_KEY 未配置：无法生成关节，请配置后重试。"
|
||||
]
|
||||
return seed
|
||||
|
||||
compact = _compact_solids(solids_geom)
|
||||
payload_user = {
|
||||
"profile": seed.profile,
|
||||
"robot_name": seed.name,
|
||||
"part_names": seed.raw_parts,
|
||||
"solids_center_mm": compact,
|
||||
"extra_hint": extra_hint,
|
||||
"instruction": (
|
||||
"直接生成完整 links+joints JSON。"
|
||||
"origin_xyz 用毫米；第一个关节靠近零件 center。"
|
||||
),
|
||||
}
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": json.dumps(payload_user, ensure_ascii=False)},
|
||||
]
|
||||
|
||||
content = await _chat(messages, temperature=0.2)
|
||||
try:
|
||||
parsed = _extract_json(content)
|
||||
except json.JSONDecodeError:
|
||||
# One repair pass: ask model to fix truncated / invalid JSON
|
||||
repair_messages = [
|
||||
{"role": "system", "content": "只输出修复后的合法 JSON 对象，不要解释。"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"下面内容不是合法 JSON（可能被截断或有尾逗号）。"
|
||||
"请修复并只输出完整 JSON：\n\n" + content[:12000]
|
||||
),
|
||||
},
|
||||
]
|
||||
repaired = await _chat(repair_messages, temperature=0.0)
|
||||
try:
|
||||
parsed = _extract_json(repaired)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(
|
||||
f"智谱返回的 JSON 无法解析（常见于输出被截断）。请重试或减少零件数量。原始错误: {e.msg}"
|
||||
) from e
|
||||
|
||||
return _draft_from_llm(parsed, seed)
|
||||
Reference in New Issue
Block a user