Generate clean URDFs (no xacro) for the Universal Robots models in ur_description

generate_ur_urdf.py

import argparse
import os
import subprocess
from xml.etree import ElementTree

def get_robot_names() -> list:
    """ Get the names of the robots in the config folder """

    # Check names of folders inside /config

    # Get the current path
    path = os.path.dirname(os.path.realpath(__file__))

    # Get the path to the config folder
    config_path = os.path.join(path, "config")

    # Get the names of the folders inside the config folder
    config_folders = os.listdir(config_path)

    # keep only folders
    robot_names = [folder for folder in config_folders if os.path.isdir(os.path.join(config_path, folder))]

    return robot_names

def select_robot() -> str:
    """Select the robot to generate the URDF for"""

    # Get the names of the robots
    config_folders = get_robot_names()

    # Print the names of the folders
    print("Available robots:")
    for i, folder in enumerate(config_folders):
        print(f"{i+1}: {folder}")

    # Ask the user to select a robot
    while True:
        try:
            selection = int(input("Select a robot: "))
            if selection < 1 or selection > len(config_folders):
                raise ValueError
            break
        except ValueError:
            print("Invalid selection, try again.")

    # Return the name of the selected robot
    return config_folders[selection-1]

def generate_urdf(robot_name: str) -> None:
    "generate the URDF for the selected robot"

    # get the config files
    path = os.path.dirname(os.path.realpath(__file__))
    config_path = os.path.join(path, "config")

    # get the initial positions file
    initial_positions_file = os.path.join(config_path, "initial_positions.yaml")

    # get the path to the config folder
    robot_configs = os.path.join(config_path, robot_name)

    default_kinematics_file = os.path.join(robot_configs, "default_kinematics.yaml")
    joint_limits_parameters_file = os.path.join(robot_configs, "joint_limits.yaml")
    physical_parameters_file = os.path.join(robot_configs, "physical_parameters.yaml")
    visual_parameters_file = os.path.join(robot_configs, "visual_parameters.yaml")

    # Create the xacro command
    command = f"xacro urdf/ur.urdf.xacro name:={robot_name} "
    command += f"initial_positions_file:={initial_positions_file} "
    command += f"kinematics_params:={default_kinematics_file} "
    command += f"joint_limit_params:={joint_limits_parameters_file} "
    command += f"physical_params:={physical_parameters_file} "
    command += f"visual_params:={visual_parameters_file} "

    # Run the command
    stdout = subprocess.run(command, shell=True, capture_output=True).stdout.decode("utf-8")

    return stdout

def clean_xacro(urdf_str) -> str:
    """ Clean URDF file from resifual xacro tags """

    # Parse the URDF into an xml tree
    urdf_xml = ElementTree.fromstring(urdf_str)

    # Remove ROS2 tags
    for child in urdf_xml:
        if 'ros2_control' in child.tag:
            urdf_xml.remove(child)

    # remove package tags
    for child in urdf_xml:
        if child.tag == 'link':
            for tag_name in ['visual', 'collision']:
                mesh = child.find(f'{tag_name}/geometry/mesh')
                if mesh is None:
                    continue

                filename = child.find(f'{tag_name}/geometry/mesh').attrib['filename']
                if 'package://' in filename:
                    filename = filename.replace('package://', '')
                    filename = '/'.join(filename.split('/')[1:])
                    child.find(f'{tag_name}/geometry/mesh').attrib['filename'] = filename

    # Return the cleaned URDF
    return ElementTree.tostring(urdf_xml, encoding="unicode")

def parse_args():
    """Parse arguments"""
    parser = argparse.ArgumentParser(description="Generate the URDF for a Universal Robots robot")
    parser.add_argument("-d", "--directory", type=str,
                        help="Directory to save the URDF in")
    parser.add_argument("-r", "--robot", type=str,
                        help=f"Name of the robot to generate the URDF for. available robots: {get_robot_names()}")
    parser.add_argument("-a", "--all", action="store_true",
                        help="Generate the URDF for all robots")
    return parser.parse_args()

def main() -> None:
    """Main function"""

    # Parse arguments
    args = parse_args()

    SAVE_DIR = args.directory
    ROBOT_NAME = args.robot
    GENERATE_ALL = args.all

    if ROBOT_NAME is not None:
        # Check if the robot name is valid
        if ROBOT_NAME not in get_robot_names():
            err_msg = f"Invalid robot name: {ROBOT_NAME}\n"
            err_msg += f"Available robots: {get_robot_names()}"
            raise ValueError(err_msg)
        robot_names = [ROBOT_NAME]
    elif GENERATE_ALL:
        # Get the names of the robots
        robot_names = get_robot_names()
        if not SAVE_DIR:
            # Ask for confirmation
            print("you have not selected a directory to save the URDF in.")
            print("This will print to terminal the URDF for all robots, are you sure you want to continue?")
            while True:
                try:
                    selection = input("y/n: ")
                    if selection not in ["y", "n"]:
                        print("Invalid selection, try again.")
                        continue
                    if selection == "n":
                        exit()
                    break
                except ValueError:
                    print("Invalid selection, try again.")
    else:
        robot_names = [select_robot()]

    for robot_name in robot_names:
        # Generate the URDF
        urdf = generate_urdf(robot_name)

        # Clean the URDF
        urdf = clean_xacro(urdf)

        if SAVE_DIR is not None:
            # Save the URDF
            with open(os.path.join(SAVE_DIR, f"{robot_name}.urdf"), "w") as f:
                f.write(urdf)
        else:
            # Print the URDF
            print(urdf)

if __name__ == "__main__":
    main()