298 lines
11 KiB
Python
298 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate simplified Gen2 collision proxies and URDF variants.
|
|
|
|
The script does not modify the original robot.urdf. It writes:
|
|
|
|
- source/engineai_lab/assets/gen2/meshes/collision_simplified/*.stl
|
|
- source/engineai_lab/assets/gen2/urdf/robot_simplified_collision.urdf
|
|
- source/engineai_lab/assets/gen2/urdf/robot_simplified_collision_mesh.urdf
|
|
- outputs/gen2_asset_audit/collision_simplified_audit.csv
|
|
|
|
The recommended URDF is robot_simplified_collision.urdf. It uses URDF box
|
|
primitives for collision. The generated STL files and mesh URDF are useful for
|
|
visual overlay checks in tools such as Blender.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
import math
|
|
import struct
|
|
import xml.etree.ElementTree as ET
|
|
from pathlib import Path
|
|
from xml.dom import minidom
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
ASSET_DIR = REPO_ROOT / "source" / "engineai_lab" / "assets" / "gen2"
|
|
URDF_DIR = ASSET_DIR / "urdf"
|
|
MESH_DIR = ASSET_DIR / "meshes"
|
|
INPUT_URDF = URDF_DIR / "robot.urdf"
|
|
OUTPUT_DIR = MESH_DIR / "collision_simplified"
|
|
OUTPUT_PRIMITIVE_URDF = URDF_DIR / "robot_simplified_collision.urdf"
|
|
OUTPUT_MESH_URDF = URDF_DIR / "robot_simplified_collision_mesh.urdf"
|
|
AUDIT_CSV = REPO_ROOT / "outputs" / "gen2_asset_audit" / "collision_simplified_audit.csv"
|
|
|
|
|
|
def _parse_floats(text: str | None, default: tuple[float, float, float]) -> tuple[float, float, float]:
|
|
if not text:
|
|
return default
|
|
values = [float(v) for v in text.split()]
|
|
if len(values) != 3:
|
|
raise ValueError(f"Expected 3 floats, got: {text}")
|
|
return values[0], values[1], values[2]
|
|
|
|
|
|
def _rpy_to_matrix(rpy: tuple[float, float, float]) -> list[list[float]]:
|
|
roll, pitch, yaw = rpy
|
|
cr, sr = math.cos(roll), math.sin(roll)
|
|
cp, sp = math.cos(pitch), math.sin(pitch)
|
|
cy, sy = math.cos(yaw), math.sin(yaw)
|
|
return [
|
|
[cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
|
|
[sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
|
|
[-sp, cp * sr, cp * cr],
|
|
]
|
|
|
|
|
|
def _mat_vec_mul(mat: list[list[float]], vec: tuple[float, float, float]) -> tuple[float, float, float]:
|
|
return tuple(sum(mat[i][j] * vec[j] for j in range(3)) for i in range(3))
|
|
|
|
|
|
def _fmt(values: tuple[float, ...] | list[float]) -> str:
|
|
return " ".join(f"{v:.9g}" for v in values)
|
|
|
|
|
|
def _binary_stl_bbox(path: Path) -> tuple[tuple[float, float, float], tuple[float, float, float], int]:
|
|
data = path.read_bytes()
|
|
if len(data) < 84:
|
|
raise ValueError(f"STL too small: {path}")
|
|
tri_count = struct.unpack("<I", data[80:84])[0]
|
|
if 84 + 50 * tri_count != len(data):
|
|
raise ValueError(f"Only binary STL is supported by this script: {path}")
|
|
|
|
mins = [float("inf"), float("inf"), float("inf")]
|
|
maxs = [float("-inf"), float("-inf"), float("-inf")]
|
|
offset = 84
|
|
for _ in range(tri_count):
|
|
values = struct.unpack("<12f", data[offset : offset + 48])
|
|
for start in (3, 6, 9):
|
|
for axis in range(3):
|
|
value = values[start + axis]
|
|
mins[axis] = min(mins[axis], value)
|
|
maxs[axis] = max(maxs[axis], value)
|
|
offset += 50
|
|
|
|
extent = tuple(maxs[i] - mins[i] for i in range(3))
|
|
center = tuple((maxs[i] + mins[i]) * 0.5 for i in range(3))
|
|
return extent, center, tri_count
|
|
|
|
|
|
def _link_scale(name: str) -> tuple[float, float, float]:
|
|
if name == "body_link":
|
|
return 0.74, 0.74, 0.74
|
|
if name.startswith("waist_link"):
|
|
return 0.22, 0.22, 0.22
|
|
if name.endswith("leg_link_6"):
|
|
return 0.90, 0.80, 0.45
|
|
if name.endswith("leg_link_4"):
|
|
return 0.40, 0.44, 0.58
|
|
if name.endswith("leg_link_5"):
|
|
return 0.42, 0.42, 0.50
|
|
if "leg_link_2" in name or "leg_link_3" in name:
|
|
return 0.30, 0.30, 0.30
|
|
if "leg_link_1" in name:
|
|
return 0.22, 0.22, 0.22
|
|
if "arm_link_7" in name:
|
|
return 0.20, 0.20, 0.20
|
|
if "arm_link_5" in name or "arm_link_6" in name:
|
|
return 0.22, 0.22, 0.22
|
|
if "arm_link_1" in name or "arm_link_2" in name:
|
|
return 0.30, 0.30, 0.30
|
|
if "arm_link" in name:
|
|
return 0.36, 0.36, 0.36
|
|
return 0.46, 0.46, 0.46
|
|
|
|
|
|
def _proxy_box(name: str, extent: tuple[float, float, float], center: tuple[float, float, float]):
|
|
scale = _link_scale(name)
|
|
size = tuple(max(extent[i] * scale[i], 0.01) for i in range(3))
|
|
proxy_center = center
|
|
|
|
# Keep the sole close to the original mesh bottom while reducing foot height.
|
|
if name.endswith("leg_link_6"):
|
|
min_z = center[2] - extent[2] * 0.5
|
|
size = (max(extent[0] * 0.90, 0.08), max(extent[1] * 0.80, 0.04), min(max(extent[2] * 0.45, 0.025), 0.045))
|
|
proxy_center = (center[0], center[1], min_z + size[2] * 0.5 + 0.002)
|
|
|
|
return proxy_center, size, scale
|
|
|
|
|
|
def _normal(a, b, c):
|
|
ux, uy, uz = b[0] - a[0], b[1] - a[1], b[2] - a[2]
|
|
vx, vy, vz = c[0] - a[0], c[1] - a[1], c[2] - a[2]
|
|
nx, ny, nz = uy * vz - uz * vy, uz * vx - ux * vz, ux * vy - uy * vx
|
|
length = math.sqrt(nx * nx + ny * ny + nz * nz)
|
|
if length <= 1e-12:
|
|
return 0.0, 0.0, 0.0
|
|
return nx / length, ny / length, nz / length
|
|
|
|
|
|
def _write_binary_stl(path: Path, triangles: list[tuple[tuple[float, float, float], ...]]) -> None:
|
|
header = b"generated simplified collision proxy".ljust(80, b" ")
|
|
with path.open("wb") as file:
|
|
file.write(header)
|
|
file.write(struct.pack("<I", len(triangles)))
|
|
for tri in triangles:
|
|
n = _normal(*tri)
|
|
file.write(struct.pack("<3f", *n))
|
|
for vertex in tri:
|
|
file.write(struct.pack("<3f", *vertex))
|
|
file.write(struct.pack("<H", 0))
|
|
|
|
|
|
def _box_triangles(center: tuple[float, float, float], size: tuple[float, float, float]):
|
|
cx, cy, cz = center
|
|
sx, sy, sz = (size[0] * 0.5, size[1] * 0.5, size[2] * 0.5)
|
|
v = [
|
|
(cx - sx, cy - sy, cz - sz),
|
|
(cx + sx, cy - sy, cz - sz),
|
|
(cx + sx, cy + sy, cz - sz),
|
|
(cx - sx, cy + sy, cz - sz),
|
|
(cx - sx, cy - sy, cz + sz),
|
|
(cx + sx, cy - sy, cz + sz),
|
|
(cx + sx, cy + sy, cz + sz),
|
|
(cx - sx, cy + sy, cz + sz),
|
|
]
|
|
faces = [
|
|
(0, 2, 1),
|
|
(0, 3, 2),
|
|
(4, 5, 6),
|
|
(4, 6, 7),
|
|
(0, 1, 5),
|
|
(0, 5, 4),
|
|
(1, 2, 6),
|
|
(1, 6, 5),
|
|
(2, 3, 7),
|
|
(2, 7, 6),
|
|
(3, 0, 4),
|
|
(3, 4, 7),
|
|
]
|
|
return [(v[a], v[b], v[c]) for a, b, c in faces]
|
|
|
|
|
|
def _replace_collision_with_box(collision: ET.Element, origin_xyz, origin_rpy, size) -> None:
|
|
for child in list(collision):
|
|
collision.remove(child)
|
|
|
|
origin = ET.SubElement(collision, "origin")
|
|
origin.set("xyz", _fmt(origin_xyz))
|
|
origin.set("rpy", _fmt(origin_rpy))
|
|
geometry = ET.SubElement(collision, "geometry")
|
|
box = ET.SubElement(geometry, "box")
|
|
box.set("size", _fmt(size))
|
|
|
|
|
|
def _replace_collision_with_mesh(collision: ET.Element, mesh_path: str) -> None:
|
|
geometry = collision.find("geometry")
|
|
if geometry is None:
|
|
geometry = ET.SubElement(collision, "geometry")
|
|
for child in list(geometry):
|
|
geometry.remove(child)
|
|
mesh = ET.SubElement(geometry, "mesh")
|
|
mesh.set("filename", mesh_path)
|
|
|
|
|
|
def _write_pretty_xml(tree: ET.ElementTree, path: Path) -> None:
|
|
rough = ET.tostring(tree.getroot(), encoding="utf-8")
|
|
pretty = minidom.parseString(rough).toprettyxml(indent=" ")
|
|
lines = [line for line in pretty.splitlines() if line.strip()]
|
|
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
def main() -> None:
|
|
if not INPUT_URDF.exists():
|
|
raise FileNotFoundError(INPUT_URDF)
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
for stale_proxy in OUTPUT_DIR.glob("*_box_collision.stl"):
|
|
stale_proxy.unlink()
|
|
|
|
primitive_tree = ET.parse(INPUT_URDF)
|
|
mesh_tree = ET.parse(INPUT_URDF)
|
|
audit_rows = []
|
|
|
|
mesh_links = {link.attrib["name"]: link for link in mesh_tree.getroot().findall("link")}
|
|
|
|
for primitive_link in primitive_tree.getroot().findall("link"):
|
|
name = primitive_link.attrib["name"]
|
|
collision = primitive_link.find("collision")
|
|
mesh_link = mesh_links[name]
|
|
mesh_collision = mesh_link.find("collision")
|
|
if collision is None or mesh_collision is None:
|
|
continue
|
|
|
|
mesh = collision.find("./geometry/mesh")
|
|
origin = collision.find("origin")
|
|
if mesh is None or origin is None:
|
|
continue
|
|
|
|
mesh_file = mesh.attrib["filename"]
|
|
# URDF 中只允许相对路径；从 URDF 所在目录解析，避免依赖当前工作目录。
|
|
source_mesh_path = (INPUT_URDF.parent / mesh_file).resolve()
|
|
extent, center, tri_count = _binary_stl_bbox(source_mesh_path)
|
|
proxy_center, size, scale = _proxy_box(name, extent, center)
|
|
|
|
original_xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0))
|
|
original_rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0))
|
|
rotated_center = _mat_vec_mul(_rpy_to_matrix(original_rpy), proxy_center)
|
|
proxy_origin_xyz = tuple(original_xyz[i] + rotated_center[i] for i in range(3))
|
|
|
|
stl_name = f"{name}_box_collision.stl"
|
|
stl_rel_path = f"../meshes/collision_simplified/{stl_name}"
|
|
_write_binary_stl(OUTPUT_DIR / stl_name, _box_triangles(proxy_center, size))
|
|
|
|
_replace_collision_with_box(collision, proxy_origin_xyz, original_rpy, size)
|
|
_replace_collision_with_mesh(mesh_collision, stl_rel_path)
|
|
|
|
audit_rows.append(
|
|
{
|
|
"link": name,
|
|
"source_mesh": mesh_file,
|
|
"proxy_mesh": stl_rel_path,
|
|
"source_triangles": tri_count,
|
|
"source_extent_x": extent[0],
|
|
"source_extent_y": extent[1],
|
|
"source_extent_z": extent[2],
|
|
"proxy_center_x": proxy_center[0],
|
|
"proxy_center_y": proxy_center[1],
|
|
"proxy_center_z": proxy_center[2],
|
|
"proxy_size_x": size[0],
|
|
"proxy_size_y": size[1],
|
|
"proxy_size_z": size[2],
|
|
"scale_x": scale[0],
|
|
"scale_y": scale[1],
|
|
"scale_z": scale[2],
|
|
}
|
|
)
|
|
|
|
if not audit_rows:
|
|
raise RuntimeError(f"No mesh collisions were converted from {INPUT_URDF}")
|
|
|
|
_write_pretty_xml(primitive_tree, OUTPUT_PRIMITIVE_URDF)
|
|
_write_pretty_xml(mesh_tree, OUTPUT_MESH_URDF)
|
|
|
|
AUDIT_CSV.parent.mkdir(parents=True, exist_ok=True)
|
|
with AUDIT_CSV.open("w", newline="", encoding="utf-8") as file:
|
|
writer = csv.DictWriter(file, fieldnames=list(audit_rows[0].keys()))
|
|
writer.writeheader()
|
|
writer.writerows(audit_rows)
|
|
|
|
print(f"wrote {OUTPUT_PRIMITIVE_URDF}")
|
|
print(f"wrote {OUTPUT_MESH_URDF}")
|
|
print(f"wrote {AUDIT_CSV}")
|
|
print(f"wrote {len(audit_rows)} collision STL files under {OUTPUT_DIR}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|