isl-org / Open3D

Open3D: A Modern Library for 3D Data Processing
http://www.open3d.org
Other
11.15k stars 2.27k forks source link

Slow viewer initialization with large shared texture image #6049

Open duburcqa opened 1 year ago

duburcqa commented 1 year ago

Checklist

My Question

I considering using open3d for robotics applications. For now, I'm using panda3d and I'm fairly satisfied but there are some issues I hoping to fix by using open3d instead, such as built-in notebook support. I wrote a sample script to start playing with open3d and get an idea of its performance.

First, once the viewer is fully loaded everything is smooth, but it takes a long time to initialize it (more than 5s on Intel i9-11900H with Nvidia RTX 3080 Max-Q). The total loading time including creating meshes in more than 8s while I'm used to about 3s using panda3d. I'm wondering if I'm doing anything bad and whether it is possible to improve performance.

Second, I would like to set the absolute pose of each mesh rather than applying a relative transform on them. Is it possible to do so without having to keep track of the current absolute transform manually ?

Here is the standalone snippet I wrote:

python3 -m pip install open3d-cpu gym_jiminy_zoo pycollada
import os
import warnings
from pkg_resources import resource_filename
from typing import Optional

import numpy as np

import trimesh
import trimesh.viewer
import open3d as o3d
import open3d.visualization as vis

import gym
import jiminy_py
import hppfcl
import pinocchio as pin

# vis.webrtc_server.enable_webrtc()

def make_capsule(radius: float,
                 length: float,
                 num_segments: int = 16,
                 num_rings: int = 16
                 ) -> o3d.t.geometry.TriangleMesh:
    vertices, normals = [], []
    for u in np.linspace(0, np.pi, num_rings):
        for v in np.linspace(0, 2 * np.pi, num_segments):
            x, y, z = np.cos(v) * np.sin(u), np.sin(v) * np.sin(u), np.cos(u)
            offset = np.sign(z) * 0.5 * length
            vertices.append((x * radius, y * radius, z * radius + offset))
            normals.append((x, y, z))

    faces = []
    for i in range(num_rings - 1):
        for j in range(num_segments - 1):
            r0 = i * num_segments + j
            r1 = r0 + num_segments
            if i < num_rings - 2:
                faces.append((r0, r1, r1 + 1))
            if i > 0:
                faces.append((r0, r1 + 1, r0 + 1))

    mesh = o3d.t.geometry.TriangleMesh()
    mesh.vertex.positions = o3d.core.Tensor(vertices, o3d.core.float32)
    mesh.vertex.normals = o3d.core.Tensor(normals, o3d.core.float32)
    mesh.triangle.indices = o3d.core.Tensor(faces, o3d.core.int32)

    return mesh

def load_mesh_file(mesh_path: str,
                   scale: Optional[np.ndarray] = None
                   ) -> o3d.t.geometry.TriangleMesh:
    mesh = trimesh.load(mesh_path, force='mesh')

    mesh_o3d = o3d.t.geometry.TriangleMesh()
    mesh_o3d.vertex.positions = o3d.core.Tensor(mesh.vertices, o3d.core.float32)
    mesh_o3d.triangle.indices = o3d.core.Tensor(mesh.faces, o3d.core.int32)
    mesh_o3d.compute_vertex_normals()

    mesh_o3d.material.set_default_properties()
    mesh_o3d.material.material_name = 'defaultLit'
    if isinstance(mesh.visual, trimesh.visual.color.ColorVisuals):
        vertex_colors = mesh.visual.vertex_colors
        if vertex_colors is not None:
            mesh_o3d.vertex.colors = o3d.core.Tensor(
            np.divide(vertex_colors[:, :3], 255), o3d.core.float32)
        face_colors = mesh.visual.face_colors
        if face_colors is not None:
            mesh_o3d.triangle.colors = o3d.core.Tensor(
            np.divide(face_colors[:, :3], 255), o3d.core.float32)
    if isinstance(mesh.visual, trimesh.visual.texture.TextureVisuals):
        texture_uvs = mesh.visual.uv
        mesh_o3d.vertex.texture_uvs = o3d.core.Tensor(texture_uvs, o3d.core.float32)
        if isinstance(mesh.visual.material, trimesh.visual.material.PBRMaterial):
            texture = mesh.visual.material.baseColorTexture
            if texture is not None:
                mesh_o3d.material.texture_maps['albedo'] = o3d.t.geometry.Image(
                    o3d.core.Tensor(np.asarray(texture), o3d.core.uint8))
            roughness = mesh.visual.material.roughnessFactor
            if roughness is not None:
                mesh_o3d.material.scalar_properties['roughness'] = roughness
            base_color = mesh.visual.material.baseColorFactor
            if base_color is not None:
                mesh_o3d.material.vector_properties['base_color'][:] = np.divide(
                    base_color, 255)
            metallic = mesh.visual.material.metallicFactor
            if metallic is not None:
                mesh_o3d.material.scalar_properties['metallic'] = metallic

    if scale is not None:
        mesh_o3d.vertex.positions *= o3d.core.Tensor(scale, o3d.core.float32)

    return mesh_o3d

def load_geometry_object(geometry_object: pin.GeometryObject,
                         color: Optional[np.ndarray] = None
                         ) -> Optional[o3d.t.geometry.TriangleMesh]:
    """Load a single geometry object
    """
    # Extract geometry information
    geom = geometry_object.geometry
    mesh_path = geometry_object.meshPath
    texture_path = ""
    if geometry_object.overrideMaterial:
        # Get material from URDF. The color is only used if no texture or
        # if its value is not the default because meshColor is never unset.
        if os.path.exists(geometry_object.meshTexturePath):
            texture_path = geometry_object.meshTexturePath
        if color is None and (not texture_path or any(
                geometry_object.meshColor != [0.9, 0.9, 0.9, 1.0])):
            color = geometry_object.meshColor

    # Try to load mesh from path first
    is_success = True
    if os.path.exists(mesh_path):
        mesh = load_mesh_file(mesh_path, geometry_object.meshScale)
    else:
        # Each geometry must have at least a color or a texture
        if color is None and not texture_path:
            color = np.array([0.75, 0.75, 0.75, 1.0])

        # Append a primitive geometry
        if isinstance(geom, hppfcl.Capsule):
            mesh = make_capsule(geom.radius, 2.0 * geom.halfLength)
        elif isinstance(geom, hppfcl.Cylinder):
            mesh = o3d.t.geometry.TriangleMesh.create_cylinder(
                geom.radius, 2*geom.halfLength)
        elif isinstance(geom, hppfcl.Cone):
            mesh = o3d.t.geometry.TriangleMesh.create_cylinder(
                geom.radius, 2*geom.halfLength)
        elif isinstance(geom, hppfcl.Box):
            mesh = o3d.t.geometry.TriangleMesh.create_box(*(2.0*geom.halfSide))
        elif isinstance(geom, hppfcl.Sphere):
            mesh = o3d.t.geometry.TriangleMesh.create_sphere(geom.radius)
        elif isinstance(geom, (hppfcl.Convex, hppfcl.BVHModelBase)):
            # Extract vertices and faces from geometry
            if isinstance(geom, hppfcl.Convex):
                vertices = geom.points()
                num_faces, get_faces = geom.num_polygons, geom.polygons
            else:
                vertices = geom.vertices()
                num_faces, get_faces = geom.num_tris, geom.tri_indices
            faces = np.empty((num_faces, 3), dtype=np.int32)
            for i in range(num_faces):
                tri = get_faces(i)
                for j in range(3):
                    faces[i, j] = tri[j]

            # Return immediately if there is nothing to load
            if num_faces == 0:
                return None

            # Create primitive triangle geometry
            mesh = o3d.t.geometry.TriangleMesh()
            mesh.vertex.positions = o3d.core.Tensor(vertices, o3d.core.float32)
            mesh.triangle.indices = o3d.core.Tensor(faces, o3d.core.int32)
            mesh.compute_vertex_normals()
        else:
            is_success = False
        if is_success:
            mesh.material.set_default_properties()
            mesh.material.material_name = 'defaultLit'

    # Early return if impossible to load the geometry for some reason
    if not is_success:
        warnings.warn(
            f"Unsupported geometry type for {geometry_object.name} "
            f"({type(geom)})", category=UserWarning, stacklevel=2)
        return None

    # Set material
    if texture_path:
        mesh.material.texture_maps['albedo'] = o3d.t.io.read_image(texture_path)
    if color is not None:
        mesh.material.vector_properties['base_color'] = np.divide(
            color, 255, dtype=np.float32)

    return mesh

# data_root_dir = resource_filename("gym_jiminy.envs", "data/bipedal_robots/atlas")
# mesh = load_mesh_file(os.path.join(data_root_dir, "meshes/utorso.dae"))
# vis.draw([mesh])

env = gym.make("gym_jiminy.envs:atlas-v0")
env.reset()
# env.render()
# env.viewer.display_collisions(True)
for geom_model, geom_data in (
        # (env.robot.collision_model, env.robot.collision_data),
        (env.robot.visual_model, env.robot.visual_data),
    ):
    meshes = {
        geom.name: load_geometry_object(geom) if geom.name != "ground" else None
        for geom in geom_model.geometryObjects}
    pin.updateGeometryPlacements(
        env.robot.pinocchio_model, env.robot.pinocchio_data, geom_model, geom_data)
    for mesh, geom, oMg in zip(
            meshes.values(), geom_model.geometryObjects, geom_data.oMg):
        if mesh is not None:
            # mesh.translate(oMg.translation, relative=True)
            # mesh.rotate(oMg.rotation, center=mesh.get_center())
            mesh.transform(oMg.homogeneous)
    # window = vis.Visualizer()
    # for name, mesh in meshes.items():
    #     if mesh is not None:
    #         window.add_geometry(name, mesh)
    #     window.poll_events()
    #     window.update_renderer()
    vis.draw(
        [mesh for mesh in meshes.values() if mesh is not None],
        show_ui=False,
        # rpc_interface=True,
        )
duburcqa commented 1 year ago

I found out what is the issue. In fact, there is a single large texture image (9MB) that is shared by all meshes (30) and the uv map only points to part of it for each of them. Apparently this use case is not properly supported by open3d. Providing the same o3d.t.geometry.Image (same memory address) does not help.