import xml.etree.ElementTree as ET from os import path ASSETS_PATH = path.join(path.dirname(path.dirname(__file__)), 'assets') MODELS_PATH = path.join(ASSETS_PATH, 'models') CHASSISES_DIR_PATH = path.join(MODELS_PATH, 'chassises') GRIPPERS_DIR_PATH = path.join(MODELS_PATH, 'grippers') MANIPULATORS_DIR_PATH = path.join(MODELS_PATH, 'manipulators') SCENES_DIR_PATH = path.join(ASSETS_PATH, 'scenes') class XMLSplicer: def __init__(self, name='robot', scene='default', chassis=None, manipulator=None, gripper=None, **kwargs, ): self.xml_name = name self.tree = None self.root = None self.splice_robot( scene=scene, chassis=chassis, manipulator=manipulator, gripper=gripper, **kwargs, ) def _init_scene(self, scene): """ Initialize the scene xml for the given scene file. We build the tree here as the global tree. Each other node will append to it. :param scene: xml file path of scene """ self.tree = ET.parse(scene) self.root = self.tree.getroot() ELEMENTS = ['asset', 'worldbody', 'actuator', 'default', 'contact', 'sensor', 'equality'] for element in ELEMENTS: if self.root.find(element) is None: self.root.append(ET.Element(element)) else: for node in self.root.findall(element): self._file_repath(scene, node) def add_component_from_xml(self, xml: str, goal_body: tuple): """ For each input xml file, we extract the component we need, and append it into global tree. :param xml: xml file path :param goal_body: goal body, a tuple consist of id and name """ if not isinstance(xml, str): raise ValueError("Please checkout your xml path.") tree = ET.parse(xml) root_node = tree.getroot() # preprocess self.tag_rename(goal_body[0], root_node) # add assets if root_node.find('asset') is not None: self._file_repath(xml, root_node.find('asset')) for asset in root_node.find('asset'): self.root.find('asset').append(asset) # add world-body if root_node.find('worldbody') is not None: for body in root_node.find('worldbody'): node = self.root.find(f'.//body[@name=\'{goal_body[1]}\']') if node is None: node = self.root.find('worldbody') node.append(body) # add actuator if root_node.find('actuator') is not None: for actuator in root_node.find('actuator'): self.root.find('actuator').append(actuator) # add default if root_node.find('default') is not None: for default in root_node.find('default'): self.root.find('default').append(default) # add contact if root_node.find('contact') is not None: for contact in root_node.find('contact'): self.root.find('contact').append(contact) # add sensor if root_node.find('sensor') is not None: for sensor in root_node.find('sensor'): self.root.find('sensor').append(sensor) # add equality if root_node.find('equality') is not None: for equality in root_node.find('equality'): self.root.find('equality').append(equality) def _file_repath(self, xml_path, asset_node: ET.Element): """ Reset file path of the assets to abstract path. Note all the texture files should put into the texture dict. :param xml_path: path of specified xml file :param asset_node: node """ for child in asset_node: if child.tag == 'mesh' and child.attrib['file'] is not None: child.attrib['file'] = path.join(path.dirname(xml_path), child.attrib['file']) elif child.tag == 'texture' and 'file' in child.attrib: child.attrib['file'] = path.join(path.dirname(path.dirname(xml_path)), 'textures', path.basename(child.attrib['file'])) def tag_rename(self, id, node: ET.Element): for mesh in node.findall('.//mesh[@name]'): for element in node.findall('.//geom[@mesh=\'{}\']'.format(mesh.attrib['name'])): element.attrib['mesh'] = '{}_{}'.format(id, element.attrib['mesh']) mesh.attrib['name'] = '{}_{}'.format(id, mesh.attrib['name']) for body in node.findall('.//body[@name]'): for element in node.findall('.//'): for key, value in element.attrib.items(): if value == body.attrib['name'] and key != 'mesh': element.attrib[key] = '{}_{}'.format(id, element.attrib[key]) for default in node.findall('.//default'): if 'class' in default.attrib: for geom in node.findall('.//geom'): if 'class' in geom.attrib and geom.get('class') == default.attrib['class']: geom.set('class', '{}_{}'.format(id, geom.attrib['class'])) for body in node.findall('.//body'): if 'childclass' in body.attrib and body.get('childclass') == default.attrib['class']: body.set('childclass', '{}_{}'.format(id, body.attrib['childclass'])) default.set('class', '{}_{}'.format(id, default.attrib['class'])) for joint in node.findall('.//joint[@name]'): target = joint.attrib['name'] for element in node.findall('.//'): for key, value in element.attrib.items(): if value == target: element.attrib[key] = '{}_{}'.format(id, element.attrib[key]) for connect in node.findall('.//connect[@name]'): for element in node.findall('.//'): for key, value in element.attrib.items(): if value == connect.attrib['name']: element.attrib[key] = '{}_{}'.format(id, element.attrib[key]) for site in node.findall('.//site[@name]'): for element in node.findall('.//'): for key, value in element.attrib.items(): if value == site.attrib['name'] and element.tag != 'site': element.attrib[key] = '{}_{}'.format(id, element.attrib[key]) site.set('name', '{}_{}'.format(id, site.attrib['name'])) for camera in node.findall('.//camera[@name]'): camera.set('name', '{}_{}'.format(id, camera.attrib['name'])) for sensor in node.findall('.//sensor/*[@name]'): sensor.set('name', '{}_{}'.format(id, sensor.attrib['name'])) for actuator in node.findall('.//actuator/*[@name]'): actuator.set('name', '{}_{}'.format(id, actuator.attrib['name'])) def add_node_from_xml(self, xml: str): parent_element = self.root.find('worldbody').find('body') new_tree = ET.parse('/assets/models/manipulators/Bimanual/Bimanual.xml') new_node = new_tree.getroot().find('worldbody').find('body') parent_element.append(new_node) def add_node_from_str(self, father_node: str, xml_text: str): parent_element = self.root.find(father_node) new_element = ET.fromstring(xml_text) parent_element.append(new_element) def add_texture(self, name: str, type: str, file: str): parent_element = self.root.find("asset") texture_element = ET.Element('texture') texture_element.set('name', name) texture_element.set('type', type) texture_element.set('file', file) parent_element.append(texture_element) def add_material(self, name: str, texture: str, texrepeat: str, texuniform: str): parent_element = self.root.find("asset") material_element = ET.Element('material') material_element.set('name', name) material_element.set('texture', texture) material_element.set('texrepeat', texrepeat) material_element.set('texuniform', texuniform) parent_element.append(material_element) def add_mesh(self, name: str, file: str): parent_element = self.root.find("asset") mesh_element = ET.Element('mesh') mesh_element.set('name', name) mesh_element.set('file', file) parent_element.append(mesh_element) def add_geom(self, node: str, **kwargs): if node == 'worldbody': geom_element = self.root.find('worldbody') else: geom_element = self.root.find(f'.//body[@name=\'{node}\']') geom = ET.Element('geom') for key in kwargs: geom.set(key, kwargs[key]) geom_element.append(geom) def add_body(self, node: str, **kwargs): if node == 'worldbody': body_element = self.root.find('worldbody') else: body_element = self.root.find('worldbody').find(node) body = ET.Element('body') for key in kwargs: body.set(key, kwargs[key]) body_element.append(body) def add_joint(self, node: str, **kwargs): if node == 'worldbody': joint_element = self.root.find('worldbody') else: joint_element = self.root.find(f'.//body[@name=\'{node}\']') joint = ET.Element('joint') for key in kwargs: joint.set(key, kwargs[key]) joint_element.append(joint) def save_xml(self, output_path='../assets'): """ Save xml file with identity path""" self.tree.write(path.join(output_path, f"{self.xml_name}.xml")) @property def xml(self, output_path=path.join(path.dirname(__file__), '../assets')): """ Save xml file and get its path""" self.save_xml(output_path) return path.abspath(path.join(output_path, f"{self.xml_name}.xml")) def splice_robot(self, scene=None, chassis=None, manipulator=None, gripper=None, **kwargs): """ :param scene: scene name :param chassis: chassis name :param manipulator: manipulator name :param gripper: gripper name :param kwargs: """ if isinstance(scene, str): if scene.endswith('.xml'): scene_path = scene else: scene_path = path.join(SCENES_DIR_PATH, '{}.xml'.format(scene)) self._init_scene(scene_path) else: raise ValueError("Must has scene.xml to generate the world.") if isinstance(chassis, str): chassis_path = path.join(CHASSISES_DIR_PATH, chassis, '{}.xml'.format(chassis)) self.add_component_from_xml(chassis_path, goal_body=(0, 'worldbody')) if isinstance(manipulator, str): if manipulator.endswith('.xml'): manipulator_path = manipulator else: manipulator_path = path.join(MANIPULATORS_DIR_PATH, manipulator, '{}.xml'.format(manipulator)) self.add_component_from_xml(manipulator_path, goal_body=(0, '0_base_car_link') if chassis is not None else (0, 'worldbody')) if isinstance(gripper, str): for goal_body in enumerate(kwargs['g2m_body']): gripper_path = path.join(GRIPPERS_DIR_PATH, gripper, '{}.xml'.format(gripper)) self.add_component_from_xml(gripper_path, goal_body=goal_body) if __name__ == "__main__": BiDianaMed = XMLSplicer( name='Bidiana', scene='default', chassis='Omnidirect', manipulator='Bimanual', gripper='robotiq_gripper', g2m_body=['0_right_link7', '0_left_link7'] ) DianaMed = XMLSplicer( name='diana_omni', scene='default', chassis='Omnidirect', manipulator='DianaMed', gripper='robotiq_gripper', g2m_body=['0_link7'] ) DianaGrasping = XMLSplicer( name='diana_omni_grasping', scene='grasping', chassis='Omnidirect', manipulator='DianaMed', gripper='robotiq_gripper', g2m_body=['0_link7'] ) print(BiDianaMed.xml) print(DianaMed.xml)