Use file's md5sum for compressed urdf

ros/kxr_models/node_scripts/urdf_model_server.py

@@ -7,6 +7,7 @@
 from filelock import FileLock
 from kxr_models.md5sum_utils import checksum_md5
 from kxr_models.urdf import aggregate_urdf_mesh_files
+from kxr_models.urdf import replace_urdf_path
 import rospkg
 import rospy
 
@@ -19,8 +20,9 @@ def __init__(self):
             full_namespace[:last_slash_pos] if last_slash_pos != 0 else ""
         )
 
-    def set_urdf_hash(self, md5sum, urdf_data):
-        rospy.set_param(self.clean_namespace + "/urdf_hash", md5sum)
+    def set_urdf_hash(self, urdf_md5sum, compressed_md5sum, urdf_data):
+        rospy.set_param(self.clean_namespace + "/urdf_hash", urdf_md5sum)
+        rospy.set_param(self.clean_namespace + "/compressed_urdf_hash", compressed_md5sum)
         rospy.set_param(self.clean_namespace + "/robot_description_viz", urdf_data)
 
     def run(self):
@@ -43,13 +45,14 @@ def run(self):
                 temp_file.write(urdf)
                 urdf_path = temp_file.name
 
-            md5sum = checksum_md5(urdf_path)
+            urdf_md5sum = checksum_md5(urdf_path)
             compressed_urdf_path = os.path.join(
-                kxr_models_path, "models", "urdf", f"{md5sum}.tar.gz"
+                kxr_models_path, "models", "urdf", f"{urdf_md5sum}.tar.gz"
             )
 
             if os.path.exists(compressed_urdf_path):
-                self._update_robot_description(urdf_path, md5sum, kxr_models_path)
+                compressed_md5sum = checksum_md5(compressed_urdf_path)
+                self._update_robot_description(urdf_path, urdf_md5sum, compressed_md5sum)
                 continue
 
             lock_path = compressed_urdf_path + ".lock"
@@ -60,18 +63,22 @@ def run(self):
                     compress=True,
                 )
                 rospy.loginfo(f"Compressed URDF model saved to {compressed_urdf_path}")
-                self._update_robot_description(urdf_path, md5sum, kxr_models_path)
+                compressed_md5sum = checksum_md5(compressed_urdf_path)
+                self._update_robot_description(urdf_path, urdf_md5sum, compressed_md5sum)
 
             os.remove(urdf_path)
 
-    def _update_robot_description(self, urdf_path, md5sum, kxr_models_path):
+    def _update_robot_description(self, urdf_path, md5sum, compressed_md5sum):
         parser = ET.XMLParser(target=ET.TreeBuilder(insert_comments=True))
         tree = ET.parse(urdf_path, parser)
         root = tree.getroot()
         robot_name = root.get("name", "default_robot")
-        with open(urdf_path) as urdf_file:
-            urdf_data = urdf_file.read()
-        self.set_urdf_hash(md5sum, urdf_data)
+        with tempfile.NamedTemporaryFile(
+                "w", suffix=".urdf", delete=True) as temp_file:
+            replace_urdf_path(urdf_path, temp_file.name)
+            with open(temp_file.name) as urdf_file:
+                urdf_data = urdf_file.read()
+        self.set_urdf_hash(md5sum, compressed_md5sum, urdf_data)
         rospy.loginfo(f"Updated robot description for {robot_name}")
 
 

ros/kxr_models/python/kxr_models/download_urdf.py

@@ -1,11 +1,15 @@
+from distutils.version import StrictVersion
 import os.path as osp
 
 import gdown
+import pkg_resources
 import rospkg
 import rospy
 
 from kxr_models.ros import get_namespace
 
+gdown_version = pkg_resources.get_distribution("gdown").version
+
 
 def download_urdf_mesh_files(namespace=None):
     if namespace is None:
@@ -14,11 +18,15 @@ def download_urdf_mesh_files(namespace=None):
     server_ip = rospy.get_param(namespace + "/model_server_ip", "localhost")
 
     urdf_hash = rospy.get_param(namespace + "/urdf_hash", None)
+    compressed_urdf_hash = rospy.get_param(
+        namespace + "/compressed_urdf_hash", None)
     rate = rospy.Rate(1)
-    while not rospy.is_shutdown() and urdf_hash is None:
+    while not rospy.is_shutdown() and urdf_hash is None and compressed_urdf_hash is None:
         rate.sleep()
-        rospy.loginfo("Waiting rosparam {} set".format(namespace + "/urdf_hash"))
+        rospy.loginfo("Waiting rosparam {} and {} set".format(
+            namespace + "/urdf_hash", namespace + "/compressed_urdf_hash"))
         urdf_hash = rospy.get_param(namespace + "/urdf_hash", None)
+        compressed_urdf_hash = rospy.get_param(namespace + "/compressed_urdf_hash", None)
 
     server_url = f"http://{server_ip}:{port}/urdf/{urdf_hash}.tar.gz"
 
@@ -31,7 +39,12 @@ def download_urdf_mesh_files(namespace=None):
     while not osp.exists(compressed_urdf_path):
         rospy.loginfo(f"Waiting {compressed_urdf_path} from server")
         rospy.sleep(1.0)
-        gdown.cached_download(url=server_url,
-                              hash=f'md5:{urdf_hash}',
-                              path=compressed_urdf_path)
+        if StrictVersion(gdown_version) < StrictVersion("5.1.0"):
+            gdown.cached_download(url=server_url,
+                                  md5=f'{compressed_urdf_hash}',
+                                  path=compressed_urdf_path)
+        else:
+            gdown.cached_download(url=server_url,
+                                  hash=f'md5:{compressed_urdf_hash}',
+                                  path=compressed_urdf_path)
     gdown.extractall(osp.join(kxr_models_path, "models", "urdf", f"{urdf_hash}.tar.gz"))

ros/kxr_models/python/kxr_models/urdf.py

@@ -57,3 +57,39 @@ def aggregate_urdf_mesh_files(input_urdf_path, output_directory, compress=False)
     if compress:
         return make_tarfile(output_dir)
     return output_urdf_path
+
+
+def replace_urdf_path(input_urdf_path, output_urdf_path):
+    urdf_path = Path(input_urdf_path)
+    if not urdf_path.exists():
+        raise OSError(f"No such urdf {urdf_path}")
+    robot_md5sum = checksum_md5(urdf_path)
+    parser = ET.XMLParser(target=ET.TreeBuilder(insert_comments=True))
+    tree = ET.parse(urdf_path, parser)
+    root = tree.getroot()
+    for mesh in root.findall(".//mesh"):
+        filename = mesh.get("filename")
+        if filename is None:
+            continue
+        abs_path = resolve_filepath(filename, urdf_path)
+        if abs_path is not None:
+            mesh_robot_md5sum = checksum_md5(abs_path)
+            suffix = Path(abs_path).suffix
+            mesh.set(
+                "filename",
+                f"package://kxr_models/models/urdf/{robot_md5sum}/{mesh_robot_md5sum}{suffix}",
+            )
+    for mesh in root.findall(".//texture"):
+        filename = mesh.get("filename")
+        if filename is None:
+            continue
+        abs_path = resolve_filepath(filename, urdf_path)
+        if abs_path is not None:
+            mesh_robot_md5sum = checksum_md5(abs_path)
+            suffix = Path(abs_path).suffix
+            mesh.set(
+                "filename",
+                f"package://kxr_models/models/urdf/{robot_md5sum}/{mesh_robot_md5sum}{suffix}",
+            )
+    tree.write(output_urdf_path)
+    return output_urdf_path