diff --git a/foam/utility/__init__.py b/foam/utility/__init__.py index ba5f9ea..b20b451 100644 --- a/foam/utility/__init__.py +++ b/foam/utility/__init__.py @@ -103,6 +103,13 @@ def load_urdf(urdf_path: Path) -> URDFDict: return xml +def load_srdf(srdf_path: Path) -> URDFDict: + with open(srdf_path, 'r') as f: + xml = xmltodict.parse(f.read()) + xml['robot']['@path'] = srdf_path + return xml + + def get_urdf_primitives(urdf: URDFDict, shrinkage: float = 1.) -> list[URDFPrimitive]: primitives = [] for link in urdf['robot']['link']: @@ -256,6 +263,45 @@ def set_urdf_spheres(urdf: URDFDict, spheres): print(f"spheres: {total_spheres}") +def set_srdf_link_sphere_approximations(srdf: URDFDict, spheres): + spheres_by_link = {} + total_spheres = 0 + for key, spherization in spheres.items(): + link_name = key.split('::')[0] + if link_name not in spheres_by_link: + spheres_by_link[link_name] = [] + + spheres_by_link[link_name].extend(spherization.spheres) + total_spheres += len(spherization.spheres) + + approximations = [] + for link_name in sorted(spheres_by_link.keys()): + link_spheres = [ + { + '@center': ' '.join(map(str, sphere.origin.tolist())), + '@radius': sphere.radius, + } + for sphere in spheres_by_link[link_name] + ] + + approximations.append( + { + '@link': link_name, + 'sphere': link_spheres, + } + ) + + if approximations: + srdf['robot']['link_sphere_approximation'] = approximations + elif 'link_sphere_approximation' in srdf['robot']: + del srdf['robot']['link_sphere_approximation'] + + def save_urdf(urdf: URDFDict, filename: Path): with open(filename, 'w') as f: f.write(xmltodict.unparse(urdf, pretty = True)) + + +def save_srdf(srdf: URDFDict, filename: Path): + with open(filename, 'w') as f: + f.write(xmltodict.unparse(srdf, pretty = True)) diff --git a/scripts/generate_sphere_urdf.py b/scripts/generate_sphere_urdf.py index 867c12d..7490454 100644 --- a/scripts/generate_sphere_urdf.py +++ b/scripts/generate_sphere_urdf.py @@ -135,8 +135,16 @@ def main( for primitive in primitives } - set_urdf_spheres(urdf, mesh_spheres | primitive_spheres) - save_urdf(urdf, Path(output)) + all_spheres = mesh_spheres | primitive_spheres + + output_path = Path(output) + if output_path.suffix.lower() == ".srdf": + srdf = load_srdf(output_path) + set_srdf_link_sphere_approximations(srdf, all_spheres) + save_srdf(srdf, output_path) + else: + set_urdf_spheres(urdf, all_spheres) + save_urdf(urdf, output_path) if __name__ == "__main__":