Source code for concepts.simulator.pybullet.urdf_utils.obj2urdf

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : obj2urdf.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
import tempfile
import os
import copy
import trimesh
import xml.etree.ElementTree as ET

from scipy.spatial.transform import Rotation

"""
Convert obj file to urdf file.

Copied from https://github.com/harvard-microrobotics/object2urdf
"""


[docs]class ObjectUrdfBuilder:
[docs] def __init__(self, object_folder, log_file=None, urdf_prototype=None): self.object_folder = os.path.abspath(object_folder) if log_file is not None: self.log_file = os.path.abspath(log_file) else: f = tempfile.NamedTemporaryFile('r', suffix='log', prefix='vhacd', delete=False) f.close() self.log_file = f.name self.suffix = "vhacd" if urdf_prototype is None: self.urdf_base = self._read_xml(os.path.join(os.path.dirname(__file__), "urdf_prototype.urdf")) else: self.urdf_base = self._read_xml(os.path.join(object_folder, urdf_prototype))
# Recursively get all files with a specific extension, excluding a certain suffix def _get_files_recursively(self, start_directory, filter_extension, exclude_suffix): for root, dirs, files in os.walk(start_directory): for file in files: if file.lower().endswith(filter_extension): if not file.lower().endswith(exclude_suffix + filter_extension): yield (root, file, os.path.abspath(os.path.join(root, file))) # Read and parse a URDF from a file def _read_xml(self, filename): root = ET.parse(filename).getroot() return root # Convert a list to a space-separated string def _list2str(self, in_list): out = "" for el in in_list: out += str(el) + " " return out[:-1] # Convert a space-separated string to a list def _str2list(self, in_str): out = in_str.split(' ') out = [float(el) for el in out] return out # Find the center of mass of the object
[docs] def get_center_of_mass(self, filename): mesh = trimesh.load(filename) if isinstance(mesh, trimesh.Scene): print("Imported combined mesh: using centroid rather than center of mass") return mesh.centroid else: return mesh.center_mass
# Find the geometric center of the object
[docs] def get_geometric_center(self, filename): mesh = trimesh.load(filename) return copy.deepcopy(mesh.centroid)
# Get the middle of a face of the bounding box
[docs] def get_face(self, filename, edge): mesh = trimesh.load(filename) bounds = mesh.bounds face = copy.deepcopy(mesh.centroid) if edge in ['top', 'xy_pos']: face[2] = bounds[1][2] elif edge in ['bottom', 'xy_neg']: face[2] = bounds[0][2] elif edge in ['xz_pos']: face[1] = bounds[1][1] elif edge in ['xz_neg']: face[1] = bounds[0][1] elif edge in ['yz_pos']: face[0] = bounds[1][0] elif edge in ['yz_neg']: face[0] = bounds[0][0] return face
# Do a convex decomposition
[docs] def do_vhacd(self, filename, outfile, debug=False, **kwargs): try: mesh = trimesh.load(filename) convex_list = trimesh.interfaces.vhacd.convex_decomposition(mesh, debug=debug, **kwargs) convex = trimesh.util.concatenate(convex_list) convex.export(outfile) except ValueError: print("No direct VHACD backend available, trying pybullet") pass try: import pybullet as p p.vhacd(filename, outfile, self.log_file, **kwargs) except ModuleNotFoundError: print( '\n' + "ERROR - pybullet module not found: If you want to do convex decomposisiton, make sure you install pybullet (https://pypi.org/project/pybullet) or install VHACD directly (https://github.com/mikedh/trimesh/issues/404)" + '\n' ) raise
# Find the center of mass of the object
[docs] def save_to_obj(self, filename): name, ext = os.path.splitext(filename) obj_filename = name + '.obj' mesh = trimesh.load(filename) mesh.export(obj_filename) return obj_filename
# Replace an attribute in a feild of a URDF
[docs] def replace_urdf_attribute(self, urdf, feild, attribute, value): urdf = self.replace_urdf_attributes(urdf, feild, {attribute: value}) return urdf
# Replace several attributes in a feild of a URDF
[docs] def replace_urdf_attributes(self, urdf, feild, attribute_dict, sub_feild=None): if sub_feild is None: sub_feild = [] field_obj = urdf.find(feild) if field_obj is not None: if len(sub_feild) > 0: for child in reversed(sub_feild): field_obj = ET.SubElement(field_obj, child) field_obj.attrib.update(attribute_dict) # field_obj.attrib = attribute_dict else: feilds = feild.split("/") new_feild = "/".join(feilds[0:-1]) sub_feild.append(feilds[-1]) self.replace_urdf_attributes(urdf, new_feild, attribute_dict, sub_feild)
# Make an updated copy of the URDF for the current object
[docs] def update_urdf(self, object_file, object_name, collision_file=None, override=None, mass_center=None): # If no separate collision geometry is provided, use the object file if collision_file is None: collision_file = object_file # Update the filenames and object name new_urdf = copy.deepcopy(self.urdf_base) self.replace_urdf_attribute(new_urdf, './/visual/geometry/mesh', 'filename', object_file) self.replace_urdf_attribute(new_urdf, './/collision/geometry/mesh', 'filename', collision_file) new_urdf.attrib['name'] = object_name # Update the overrides if override is not None: for orverride_el in override: # Update attributes out_el_all = new_urdf.findall('.//' + orverride_el.tag) for out_el in out_el_all: for key in orverride_el.attrib: out_el.set(key, orverride_el.attrib[key]) # Remove fields that will be overwritten for child in orverride_el: el = out_el.find(child.tag) if el is not None: out_el.remove(el) # Add updated feilds out_el.extend(orverride_el) # Output the center of mass if provided if mass_center is not None: # Check if there's a geometry offset offset_ob = new_urdf.find('.//collision/origin') if offset_ob is not None: offset_str = offset_ob.attrib.get('xyz', '0 0 0') rot_str = offset_ob.attrib.get('rpy', '0 0 0') offset = self._str2list(offset_str) rpy = self._str2list(rot_str) else: offset = [0, 0, 0] rpy = [0, 0, 0] # Check if there's a scale factor and apply it scale_ob = new_urdf.find('.//collision/geometry/mesh') if scale_ob is not None: scale_str = scale_ob.attrib.get('scale', '1 1 1') scale = self._str2list(scale_str) else: scale = [1, 1, 1] for idx, axis in enumerate(mass_center): mass_center[idx] = -mass_center[idx] * scale[idx] + offset[idx] rot = Rotation.from_euler('xyz', rpy) rot_matrix = rot.as_matrix() mass_center = np.matmul(rot_matrix, np.vstack(np.asarray(mass_center))).squeeze() self.replace_urdf_attributes( new_urdf, './/visual/origin', {'xyz': self._list2str(mass_center), 'rpy': self._list2str(rpy)} ) self.replace_urdf_attributes( new_urdf, './/collision/origin', {'xyz': self._list2str(mass_center), 'rpy': self._list2str(rpy)} ) return new_urdf
# Save a URDF to a file
[docs] def save_urdf(self, new_urdf, filename, overwrite=False): out_file = os.path.join(self.object_folder, filename) # Do not overwrite the file unless the option is True if os.path.exists(out_file) and not overwrite: return # Save the file mydata = ET.tostring(new_urdf) with open(out_file, "wb") as f: f.write(mydata)
# Build a URDF from an object file
[docs] def build_urdf( self, filename, output_folder=None, force_overwrite=False, decompose_concave=False, force_decompose=False, center='mass', **kwargs ): # If no output folder is specified, use the base object folder if output_folder is None: output_folder = self.object_folder # Generate a relative path from the output folder to the geometry files filename = os.path.abspath(filename) common = os.path.commonprefix([output_folder, filename]) rel = os.path.join(filename.replace(common, '')) if rel[0] == os.path.sep: rel = rel[1:] name = rel.split(os.path.sep)[0] rel = rel.replace(os.path.sep, '/') file_name_raw, file_extension = os.path.splitext(filename) # If an override file exists, include its data in the URDF override_file = filename.replace(file_extension, '.ovr') if os.path.exists(override_file): overrides = self._read_xml(override_file) else: overrides = None # Calculate the center of mass if center == 'mass': mass_center = self.get_center_of_mass(filename) elif center == 'geometric': mass_center = self.get_geometric_center(filename) elif center in ['top', 'bottom', 'xy_pos', 'xy_neg', 'xz_pos', 'xz_neg', 'yz_pos', 'yz_neg']: mass_center = self.get_face(filename, center) else: mass_center = None mesh = trimesh.load(filename) print(mesh.bounds) # If the user wants to run convex decomposition on concave objects, do it. if decompose_concave: if file_extension == '.stl': obj_filename = self.save_to_obj(filename) visual_file = rel.replace(file_extension, '.obj') elif file_extension == '.obj': obj_filename = filename visual_file = rel else: raise ValueError("Your filetype needs to be an STL or OBJ to perform concave decomposition") outfile = obj_filename.replace('.obj', '_' + self.suffix + '.obj') collision_file = visual_file.replace('.obj', '_' + self.suffix + '.obj') # Only run a decomposition if one does not exist, or if the user forces an overwrite if not os.path.exists(outfile) or force_decompose: self.do_vhacd(obj_filename, outfile, **kwargs) urdf_out = self.update_urdf(visual_file, name, collision_file=collision_file, override=overrides, mass_center=mass_center) else: urdf_out = self.update_urdf(rel, name, override=overrides, mass_center=mass_center) self.save_urdf(urdf_out, name + '.urdf', force_overwrite)
# Build the URDFs for all objects in your library.
[docs] def build_library(self, **kwargs): print("\nFOLDER: %s" % (self.object_folder)) # Get all OBJ files obj_files = self._get_files_recursively(self.object_folder, filter_extension='.obj', exclude_suffix=self.suffix) stl_files = self._get_files_recursively(self.object_folder, filter_extension='.stl', exclude_suffix=self.suffix) obj_folders = [] for root, _, full_file in obj_files: obj_folders.append(root) self.build_urdf(full_file, **kwargs) common = os.path.commonprefix([self.object_folder, full_file]) rel = os.path.join(full_file.replace(common, '')) print('\tBuilding: %s' % (rel)) for root, _, full_file in stl_files: if root not in obj_folders: self.build_urdf(full_file, **kwargs) common = os.path.commonprefix([self.object_folder, full_file]) rel = os.path.join(full_file.replace(common, '')) print('Building: %s' % (rel))