Files
step2urdf/app/services/pipeline.py
T
sunxianghui 93773f3887 Initial commit: step2urdf tool with handtuned JSON import.
Includes FastAPI backend, vendored step2urdf frontend, and A7 handtuned arm JSON for URDF generation.
2026-08-26 15:32:47 +08:00

120 lines
3.9 KiB
Python

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