44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import xml.etree.ElementTree as ET
|
|
import math
|
|
|
|
def round_axis(x, y, z):
|
|
# Normalize vector
|
|
mag = math.sqrt(x**2 + y**2 + z**2)
|
|
if mag == 0:
|
|
return (0, 0, 0)
|
|
x, y, z = x/mag, y/mag, z/mag
|
|
|
|
# Compare to cardinal directions
|
|
directions = {
|
|
(1, 0, 0): 'X+',
|
|
(-1, 0, 0): 'X-',
|
|
(0, 1, 0): 'Y+',
|
|
(0, -1, 0): 'Y-',
|
|
(0, 0, 1): 'Z+',
|
|
(0, 0, -1): 'Z-'
|
|
}
|
|
|
|
best = max(directions.keys(), key=lambda d: x*d[0] + y*d[1] + z*d[2])
|
|
return best
|
|
|
|
def process_urdf(file_path, output_path):
|
|
tree = ET.parse(file_path)
|
|
root = tree.getroot()
|
|
|
|
for joint in root.findall('joint'):
|
|
axis_elem = joint.find('axis')
|
|
if axis_elem is not None:
|
|
axis_str = axis_elem.attrib.get('xyz', '')
|
|
try:
|
|
x, y, z = map(float, axis_str.strip().split())
|
|
rx, ry, rz = round_axis(x, y, z)
|
|
axis_elem.set('xyz', f"{rx} {ry} {rz}")
|
|
except:
|
|
continue
|
|
|
|
tree.write(output_path)
|
|
print(f"Rounded axes written to {output_path}")
|
|
|
|
# Example usage
|
|
process_urdf("./urdf/sample.urdf", "./urdf/output.urdf")
|