311 lines
11 KiB
Python
311 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Render Gen2 collision boxes in the configured default standing pose."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import ast
|
|
import math
|
|
import re
|
|
import xml.etree.ElementTree as ET
|
|
from pathlib import Path
|
|
|
|
import matplotlib
|
|
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
DEFAULT_URDF = (
|
|
REPO_ROOT
|
|
/ "source"
|
|
/ "engineai_lab"
|
|
/ "assets"
|
|
/ "gen2"
|
|
/ "urdf"
|
|
/ "robot_simplified_collision.urdf"
|
|
)
|
|
# 生成图属于检查产物，不写回运行时资产包。
|
|
DEFAULT_OUTPUT = REPO_ROOT / "outputs" / "gen2_asset_audit" / "gen2_collision_visualization.png"
|
|
GEN2_CFG = REPO_ROOT / "source" / "engineai_lab" / "robots" / "gen2.py"
|
|
|
|
|
|
def _parse_floats(text: str | None, default: tuple[float, ...]) -> tuple[float, ...]:
|
|
if not text:
|
|
return default
|
|
return tuple(float(value) for value in text.split())
|
|
|
|
|
|
def _rpy_matrix(rpy: tuple[float, float, float]) -> np.ndarray:
|
|
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 np.array(
|
|
[
|
|
[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],
|
|
],
|
|
dtype=float,
|
|
)
|
|
|
|
|
|
def _axis_angle(axis: tuple[float, float, float], angle: float) -> np.ndarray:
|
|
vec = np.asarray(axis, dtype=float)
|
|
norm = np.linalg.norm(vec)
|
|
if norm < 1.0e-12:
|
|
return np.eye(3)
|
|
x, y, z = vec / norm
|
|
c, s = math.cos(angle), math.sin(angle)
|
|
one_c = 1.0 - c
|
|
return np.array(
|
|
[
|
|
[c + x * x * one_c, x * y * one_c - z * s, x * z * one_c + y * s],
|
|
[y * x * one_c + z * s, c + y * y * one_c, y * z * one_c - x * s],
|
|
[z * x * one_c - y * s, z * y * one_c + x * s, c + z * z * one_c],
|
|
],
|
|
dtype=float,
|
|
)
|
|
|
|
|
|
def _transform(xyz=(0.0, 0.0, 0.0), rpy=(0.0, 0.0, 0.0)) -> np.ndarray:
|
|
matrix = np.eye(4)
|
|
matrix[:3, :3] = _rpy_matrix(rpy)
|
|
matrix[:3, 3] = np.asarray(xyz, dtype=float)
|
|
return matrix
|
|
|
|
|
|
def _load_default_joint_positions() -> dict[str, float]:
|
|
if not GEN2_CFG.exists():
|
|
return {}
|
|
text = GEN2_CFG.read_text(encoding="utf-8")
|
|
match = re.search(r"joint_pos=\{(.*?)\},\n\s*joint_vel=", text, flags=re.S)
|
|
if match is None:
|
|
return {}
|
|
return ast.literal_eval("{" + match.group(1) + "}")
|
|
|
|
|
|
def _compute_link_world_transforms(root: ET.Element, joint_positions: dict[str, float]) -> dict[str, np.ndarray]:
|
|
link_names = [link.attrib["name"] for link in root.findall("link")]
|
|
children: dict[str, list[str]] = {}
|
|
parent_of: dict[str, str] = {}
|
|
edge = {}
|
|
|
|
for joint in root.findall("joint"):
|
|
parent = joint.find("parent").attrib["link"]
|
|
child = joint.find("child").attrib["link"]
|
|
origin = joint.find("origin")
|
|
axis_elem = joint.find("axis")
|
|
xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
|
|
rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
|
|
axis = (
|
|
_parse_floats(axis_elem.attrib.get("xyz"), (1.0, 0.0, 0.0)) if axis_elem is not None else (1.0, 0.0, 0.0)
|
|
)
|
|
children.setdefault(parent, []).append(child)
|
|
parent_of[child] = parent
|
|
edge[(parent, child)] = (joint.attrib["name"], joint.attrib.get("type", "fixed"), xyz, rpy, axis)
|
|
|
|
root_link = next(name for name in link_names if name not in parent_of)
|
|
world = {root_link: np.eye(4)}
|
|
stack = [root_link]
|
|
while stack:
|
|
parent = stack.pop()
|
|
for child in children.get(parent, []):
|
|
joint_name, joint_type, xyz, rpy, axis = edge[(parent, child)]
|
|
joint_transform = _transform(xyz, rpy)
|
|
if joint_type in {"revolute", "continuous"}:
|
|
rotation = np.eye(4)
|
|
rotation[:3, :3] = _axis_angle(axis, joint_positions.get(joint_name, 0.0))
|
|
joint_transform = joint_transform @ rotation
|
|
elif joint_type == "prismatic":
|
|
translation = np.eye(4)
|
|
translation[:3, 3] = np.asarray(axis, dtype=float) * joint_positions.get(joint_name, 0.0)
|
|
joint_transform = joint_transform @ translation
|
|
world[child] = world[parent] @ joint_transform
|
|
stack.append(child)
|
|
return world
|
|
|
|
|
|
def _box_vertices(transform: np.ndarray, size: tuple[float, float, float]) -> np.ndarray:
|
|
sx, sy, sz = (np.asarray(size, dtype=float) * 0.5).tolist()
|
|
local = np.array(
|
|
[
|
|
[-sx, -sy, -sz, 1.0],
|
|
[sx, -sy, -sz, 1.0],
|
|
[sx, sy, -sz, 1.0],
|
|
[-sx, sy, -sz, 1.0],
|
|
[-sx, -sy, sz, 1.0],
|
|
[sx, -sy, sz, 1.0],
|
|
[sx, sy, sz, 1.0],
|
|
[-sx, sy, sz, 1.0],
|
|
]
|
|
)
|
|
return (transform @ local.T).T[:, :3]
|
|
|
|
|
|
def _box_faces(vertices: np.ndarray) -> list[np.ndarray]:
|
|
indices = [(0, 1, 2, 3), (4, 5, 6, 7), (0, 1, 5, 4), (1, 2, 6, 5), (2, 3, 7, 6), (3, 0, 4, 7)]
|
|
return [vertices[list(face)] for face in indices]
|
|
|
|
|
|
def _module_color(link_name: str) -> str:
|
|
if link_name.startswith("left_leg"):
|
|
return "#2a9d8f"
|
|
if link_name.startswith("right_leg"):
|
|
return "#e76f51"
|
|
if link_name.startswith("left_arm"):
|
|
return "#457b9d"
|
|
if link_name.startswith("right_arm"):
|
|
return "#f4a261"
|
|
if link_name.startswith("waist"):
|
|
return "#8d5a97"
|
|
return "#5c677d"
|
|
|
|
|
|
def _collect_boxes(root: ET.Element) -> list[dict]:
|
|
joint_positions = _load_default_joint_positions()
|
|
link_world = _compute_link_world_transforms(root, joint_positions)
|
|
boxes = []
|
|
for link in root.findall("link"):
|
|
link_name = link.attrib["name"]
|
|
collision = link.find("collision")
|
|
box = link.find("./collision/geometry/box")
|
|
if collision is None or box is None:
|
|
continue
|
|
origin = collision.find("origin")
|
|
xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
|
|
rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
|
|
size = _parse_floats(box.attrib["size"], (0.0, 0.0, 0.0))
|
|
transform = link_world[link_name] @ _transform(xyz, rpy)
|
|
boxes.append(
|
|
{
|
|
"name": link_name,
|
|
"vertices": _box_vertices(transform, size),
|
|
"center": transform[:3, 3],
|
|
"size": size,
|
|
"color": _module_color(link_name),
|
|
}
|
|
)
|
|
return boxes
|
|
|
|
|
|
def _set_equal_axes_3d(ax, points: np.ndarray) -> None:
|
|
mins = points.min(axis=0)
|
|
maxs = points.max(axis=0)
|
|
center = (mins + maxs) * 0.5
|
|
radius = float((maxs - mins).max() * 0.56)
|
|
ax.set_xlim(center[0] - radius, center[0] + radius)
|
|
ax.set_ylim(center[1] - radius, center[1] + radius)
|
|
ax.set_zlim(center[2] - radius, center[2] + radius)
|
|
|
|
|
|
def _draw_projection(ax, boxes: list[dict], dims: tuple[int, int], labels: tuple[str, str], title: str) -> None:
|
|
edges = [(0, 1), (1, 2), (2, 3), (3, 0), (4, 5), (5, 6), (6, 7), (7, 4), (0, 4), (1, 5), (2, 6), (3, 7)]
|
|
for box in boxes:
|
|
vertices = box["vertices"]
|
|
for start, end in edges:
|
|
ax.plot(
|
|
[vertices[start, dims[0]], vertices[end, dims[0]]],
|
|
[vertices[start, dims[1]], vertices[end, dims[1]]],
|
|
color=box["color"],
|
|
linewidth=1.2,
|
|
)
|
|
ax.set_title(title)
|
|
ax.set_xlabel(labels[0])
|
|
ax.set_ylabel(labels[1])
|
|
ax.set_aspect("equal", adjustable="box")
|
|
ax.grid(True, alpha=0.25)
|
|
|
|
|
|
def _save_projection(output_path: Path, boxes: list[dict], dims: tuple[int, int], labels: tuple[str, str], title: str) -> None:
|
|
fig, ax = plt.subplots(figsize=(8, 8), dpi=180)
|
|
_draw_projection(ax, boxes, dims, labels, title)
|
|
fig.tight_layout()
|
|
fig.savefig(output_path)
|
|
plt.close(fig)
|
|
print(f"wrote {output_path}")
|
|
|
|
|
|
def _save_3d(output_path: Path, boxes: list[dict], all_points: np.ndarray) -> None:
|
|
fig = plt.figure(figsize=(8, 8), dpi=180)
|
|
ax3d = fig.add_subplot(1, 1, 1, projection="3d")
|
|
for box in boxes:
|
|
collection = Poly3DCollection(_box_faces(box["vertices"]), facecolor=box["color"], edgecolor="#202020", alpha=0.32)
|
|
collection.set_linewidth(0.45)
|
|
ax3d.add_collection3d(collection)
|
|
_set_equal_axes_3d(ax3d, all_points)
|
|
ax3d.set_title("Gen2 collision boxes - 3D")
|
|
ax3d.set_xlabel("X forward (m)")
|
|
ax3d.set_ylabel("Y left (m)")
|
|
ax3d.set_zlabel("Z up (m)")
|
|
ax3d.view_init(elev=18, azim=-58)
|
|
fig.tight_layout()
|
|
fig.savefig(output_path)
|
|
plt.close(fig)
|
|
print(f"wrote {output_path}")
|
|
|
|
|
|
def render(urdf_path: Path, output_path: Path) -> None:
|
|
root = ET.parse(urdf_path).getroot()
|
|
boxes = _collect_boxes(root)
|
|
if not boxes:
|
|
raise RuntimeError(f"No box collisions found in {urdf_path}")
|
|
|
|
all_points = np.concatenate([box["vertices"] for box in boxes], axis=0)
|
|
fig = plt.figure(figsize=(15, 11), dpi=170)
|
|
|
|
ax3d = fig.add_subplot(2, 2, 1, projection="3d")
|
|
for box in boxes:
|
|
collection = Poly3DCollection(_box_faces(box["vertices"]), facecolor=box["color"], edgecolor="#202020", alpha=0.32)
|
|
collection.set_linewidth(0.45)
|
|
ax3d.add_collection3d(collection)
|
|
_set_equal_axes_3d(ax3d, all_points)
|
|
ax3d.set_title("Gen2 collision boxes - 3D")
|
|
ax3d.set_xlabel("X forward (m)")
|
|
ax3d.set_ylabel("Y left (m)")
|
|
ax3d.set_zlabel("Z up (m)")
|
|
ax3d.view_init(elev=18, azim=-58)
|
|
|
|
_draw_projection(fig.add_subplot(2, 2, 2), boxes, (1, 2), ("Y left (m)", "Z up (m)"), "Front view")
|
|
_draw_projection(fig.add_subplot(2, 2, 3), boxes, (0, 2), ("X forward (m)", "Z up (m)"), "Side view")
|
|
_draw_projection(fig.add_subplot(2, 2, 4), boxes, (0, 1), ("X forward (m)", "Y left (m)"), "Top view")
|
|
|
|
legend_items = [
|
|
("body", "#5c677d"),
|
|
("waist", "#8d5a97"),
|
|
("left leg", "#2a9d8f"),
|
|
("right leg", "#e76f51"),
|
|
("left arm", "#457b9d"),
|
|
("right arm", "#f4a261"),
|
|
]
|
|
handles = [plt.Line2D([0], [0], color=color, lw=4, label=name) for name, color in legend_items]
|
|
fig.legend(handles=handles, loc="lower center", ncol=6, frameon=False)
|
|
fig.suptitle(f"{urdf_path.name}: {len(boxes)} collision boxes", fontsize=14)
|
|
fig.tight_layout(rect=(0, 0.04, 1, 0.96))
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
fig.savefig(output_path)
|
|
plt.close(fig)
|
|
print(f"wrote {output_path}")
|
|
|
|
stem = output_path.with_suffix("")
|
|
_save_3d(stem.with_name(f"{stem.name}_3d.png"), boxes, all_points)
|
|
_save_projection(stem.with_name(f"{stem.name}_front.png"), boxes, (1, 2), ("Y left (m)", "Z up (m)"), "Front view")
|
|
_save_projection(stem.with_name(f"{stem.name}_side.png"), boxes, (0, 2), ("X forward (m)", "Z up (m)"), "Side view")
|
|
_save_projection(stem.with_name(f"{stem.name}_top.png"), boxes, (0, 1), ("X forward (m)", "Y left (m)"), "Top view")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--urdf", type=Path, default=DEFAULT_URDF)
|
|
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
|
|
args = parser.parse_args()
|
|
render(args.urdf.resolve(), args.output.resolve())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|