Fix urdf conversion script

happypose/pose_estimators/cosypose/cosypose/libmesh/urdf_utils.py

@@ -30,15 +30,17 @@ def convert_rigid_body_dataset_to_urdfs(
         obj_to_urdf(obj_path, urdf_path)
 
 
-def ply_to_obj(ply_path, obj_path, texture_size=None):
+def ply_to_obj(ply_path: Path, obj_path: Path, texture_size=None):
+    assert obj_path.suffix == ".obj"
     mesh = trimesh.load(ply_path)
     obj_label = obj_path.with_suffix("").name
 
     # adapt materials according to previous example meshes
-    mesh.visual.material.ambient = np.array([51, 51, 51, 255], dtype=np.uint8)
-    mesh.visual.material.diffuse = np.array([255, 255, 255, 255], dtype=np.uint8)
-    mesh.visual.material.specular = np.array([255, 255, 255, 255], dtype=np.uint8)
-    mesh.visual.material.name = obj_label + "_texture"
+    if mesh.visual.defined:
+        mesh.visual.material.ambient = np.array([51, 51, 51, 255], dtype=np.uint8)
+        mesh.visual.material.diffuse = np.array([255, 255, 255, 255], dtype=np.uint8)
+        mesh.visual.material.specular = np.array([255, 255, 255, 255], dtype=np.uint8)
+        mesh.visual.material.name = obj_label + "_texture"
 
     # print(mesh.visual.uv)
     kwargs_export = {"mtl_name": f"{obj_label}.mtl"}

happypose/pose_estimators/cosypose/cosypose/scripts/convert_models_to_urdf.py

@@ -3,26 +3,37 @@
 
 from tqdm import tqdm
 
-from happypose.pose_estimators.cosypose.cosypose.datasets.datasets_cfg import (
-    make_object_dataset,
-)
 from happypose.pose_estimators.cosypose.cosypose.libmesh.urdf_utils import (
     obj_to_urdf,
     ply_to_obj,
 )
 
+# from happypose.pose_estimators.cosypose.cosypose.datasets.datasets_cfg import \
+#     make_object_dataset
+from happypose.toolbox.datasets.datasets_cfg import make_object_dataset
+
 
 def convert_bop_dataset_to_urdfs(
     obj_ds_name: str, urdf_dir: Path, texture_size=(1024, 1024)
 ):
+    """
+    For each object, generate these files:
+
+    {path_to_urdf_dir}/{ds_name}/{obj_label}/{obj_label}.obj
+    {path_to_urdf_dir}/{ds_name}/{obj_label}/{obj_label}.obj.mtl
+    {path_to_urdf_dir}/{ds_name}/{obj_label}/{obj_label}.png
+    {path_to_urdf_dir}/{ds_name}/{obj_label}/{obj_label}.urdf
+    """
     obj_dataset = make_object_dataset(obj_ds_name)
-    urdf_dir.mkdir(exist_ok=True, parents=True)
+    urdf_ds_dir = urdf_dir / obj_ds_name
+    urdf_ds_dir.mkdir(exist_ok=True, parents=True)
     for n in tqdm(range(len(obj_dataset))):
         obj = obj_dataset[n]
-        ply_path = Path(obj["mesh_path"])
-        out_dir = urdf_dir / obj["label"]
-        out_dir.mkdir(exist_ok=True)
-        obj_path = out_dir / ply_path.with_suffix(".obj").name
+        ply_path = obj.mesh_path
+        obj_name = ply_path.with_suffix("").name
+        obj_urdf_dir = urdf_ds_dir / obj_name
+        obj_urdf_dir.mkdir(exist_ok=True)
+        obj_path = (obj_urdf_dir / obj_name).with_suffix(".obj")
         ply_to_obj(ply_path, obj_path, texture_size=texture_size)
         obj_to_urdf(obj_path, obj_path.with_suffix(".urdf"))
 
@@ -34,8 +45,8 @@ def main():
 
     from happypose.pose_estimators.cosypose.cosypose.config import LOCAL_DATA_DIR
 
-    urdf_dir = LOCAL_DATA_DIR / "urdf"
-    convert_bop_dataset_to_urdfs(urdf_dir, args.models)
+    urdf_dir = LOCAL_DATA_DIR / "urdfs"
+    convert_bop_dataset_to_urdfs(args.models, urdf_dir)
 
 
 if __name__ == "__main__":