Merge pull request #1155 from StanfordVL/og-dev-improved

Improved version of import_custom_object.py and additional function for generate_urdf_for_mesh()

Author: williampai0704

omnigibson/examples/objects/import_custom_object.py

@@ -1,20 +1,20 @@
 """
 Helper script to download OmniGibson dataset and assets.
+Improved version that can import obj file and articulated file (glb, gltf).
 """
 
-import math
 from typing import Literal
-
 import click
-import trimesh
-
+import sys
+import shutil
+import select
+import tempfile
 import omnigibson as og
+
 from omnigibson.utils.asset_conversion_utils import (
-    generate_collision_meshes,
-    generate_urdf_for_obj,
     import_og_asset_from_urdf,
+    generate_urdf_for_mesh,
 )
-from omnigibson.utils.python_utils import assert_valid_key
 
 
 @click.command()
@@ -43,108 +43,106 @@
     default=32,
     help="Maximum number of convex hulls to decompose individual visual meshes into. Only relevant if --collision-method=coacd",
 )
-@click.option("--scale", type=float, default=1.0, help="Scale factor to apply to the mesh.")
 @click.option("--up-axis", type=click.Choice(["z", "y"]), default="z", help="Up axis for the mesh.")
 @click.option("--headless", is_flag=True, help="Run the script in headless mode.")
-@click.option("--confirm-bbox", default=True, help="Whether to confirm the scale factor.")
+@click.option("--scale", type=int, default=1, help="User choice scale, will be overwritten if check_scale and rescale")
+@click.option("--check_scale", is_flag=True, help="Check meshes scale based on heuristic")
+@click.option("--rescale", is_flag=True, help="Rescale meshes based on heuristic if check_scale ")
 @click.option("--overwrite", is_flag=True, help="Overwrite any pre-existing files")
+@click.option("--n_submesh", type=int, help="Maximum of submesh numnber")
 def import_custom_object(
     asset_path: str,
     category: str,
     model: str,
     collision_method: Literal["coacd", "convex", "none"],
     hull_count: int,
-    scale: float,
     up_axis: Literal["z", "y"],
     headless: bool,
-    confirm_bbox: bool,
+    scale: int,
+    check_scale: bool,
+    rescale: bool,
     overwrite: bool,
+    n_submesh: int,
 ):
     """
     Imports a custom-defined object asset into an OmniGibson-compatible USD format and saves the imported asset
     files to the custom dataset directory (gm.CUSTOM_DATASET_PATH)
     """
+
     assert len(model) == 6 and model.isalpha(), "Model name must be 6 characters long and contain only letters."
     collision_method = None if collision_method == "none" else collision_method
 
     # Sanity check mesh type
-    valid_formats = trimesh.available_formats()
     mesh_format = asset_path.split(".")[-1]
 
     # If we're not a URDF, import the mesh directly first
+    urdf_dep_paths = None
+    temp_dirs = []
     if mesh_format != "urdf":
-        assert_valid_key(key=mesh_format, valid_keys=valid_formats, name="mesh format")
-
-        # Load the mesh
-        visual_mesh: trimesh.Trimesh = trimesh.load(asset_path, force="mesh", process=False)
-
-        # Generate collision meshes if requested
-        collision_meshes = (
-            generate_collision_meshes(visual_mesh, method=collision_method, hull_count=hull_count)
-            if collision_method is not None
-            else []
+        temp_urdf_dir = tempfile.mkdtemp()
+        temp_dirs.append(temp_urdf_dir)
+
+        # Try to generate URDF, may raise ValueError if too many submeshes
+        urdf_path = generate_urdf_for_mesh(
+            asset_path,
+            temp_urdf_dir,
+            category,
+            model,
+            collision_method,
+            hull_count,
+            up_axis,
+            scale=scale,
+            check_scale=check_scale,
+            rescale=rescale,
+            overwrite=overwrite,
+            n_submesh=n_submesh,
         )
-
-        # If the up axis is y, we need to rotate the meshes
-        if up_axis == "y":
-            rotation_matrix = trimesh.transformations.rotation_matrix(math.pi / 2, [1, 0, 0])
-            visual_mesh.apply_transform(rotation_matrix)
-            for collision_mesh in collision_meshes:
-                collision_mesh.apply_transform(rotation_matrix)
-
-        # If the scale is nonzero, we apply it to the meshes
-        if scale != 1.0:
-            scale_transform = trimesh.transformations.scale_matrix(scale)
-            visual_mesh.apply_transform(scale_transform)
-            for collision_mesh in collision_meshes:
-                collision_mesh.apply_transform(scale_transform)
-
-        # Check the bounding box size and complain if it's larger than 3 meters
-        bbox_size = visual_mesh.bounding_box.extents
-        click.echo(f"Visual mesh bounding box size: {bbox_size}")
-
-        if confirm_bbox:
-            if any(size > 3.0 for size in bbox_size):
-                click.echo(
-                    "Warning: The bounding box sounds a bit large. Are you sure you don't need to scale? "
-                    "We just wanted to confirm this is intentional. You can skip this check by passing --no-confirm-bbox."
-                )
-                click.confirm("Do you want to continue?", abort=True)
-
-            elif any(size < 0.01 for size in bbox_size):
-                click.echo(
-                    "Warning: The bounding box sounds a bit small. Are you sure you don't need to scale? "
-                    "We just wanted to confirm this is intentional. You can skip this check by passing --no-confirm-bbox."
-                )
-                click.confirm("Do you want to continue?", abort=True)
-
-        # Generate the URDF
-        click.echo(f"Generating URDF for {category}/{model}...")
-        generate_urdf_for_obj(visual_mesh, collision_meshes, category, model, overwrite=overwrite)
-        click.echo("URDF generation complete!")
-
-        urdf_path = None
-        collision_method = None
+        if urdf_path is not None:
+            click.echo("URDF generation complete!")
+            urdf_dep_paths = ["material"]
+            collision_method = None
+        else:
+            # Clean up temp directories before exiting
+            for tmp_dir in temp_dirs:
+                shutil.rmtree(tmp_dir)
+            click.echo("Error during URDF generation")
+            raise click.Abort()
     else:
         urdf_path = asset_path
         collision_method = collision_method
 
-    # Convert to USD
-    import_og_asset_from_urdf(
-        category=category,
-        model=model,
-        urdf_path=urdf_path,
-        collision_method=collision_method,
-        hull_count=hull_count,
-        overwrite=overwrite,
-        use_usda=False,
-    )
+    try:
+        # Convert to USD
+        import_og_asset_from_urdf(
+            category=category,
+            model=model,
+            urdf_path=urdf_path,
+            urdf_dep_paths=urdf_dep_paths,
+            collision_method=collision_method,
+            hull_count=hull_count,
+            overwrite=overwrite,
+            use_usda=False,
+        )
+
+    except Exception as e:
+        click.echo(f"Error during USD conversion: {str(e)}")
+        # Clean up temp directories before exiting
+        for tmp_dir in temp_dirs:
+            shutil.rmtree(tmp_dir)
+        raise click.Abort()
+
+    # Clean up temp directories
+    for tmp_dir in temp_dirs:
+        shutil.rmtree(tmp_dir)
 
     # Visualize if not headless
     if not headless:
         click.echo("The asset has been successfully imported. You can view it and make changes and save if you'd like.")
         while True:
             og.sim.render()
+            if select.select([sys.stdin], [], [], 0)[0]:
+                sys.stdin.readline()  # Clear the input buffer
+                break
 
 
 if __name__ == "__main__":