improve point cloud visualization

include/sapien/sapien_renderer/point_cloud_component.h

@@ -15,6 +15,8 @@ class PointCloudComponent : public Component {
 
   std::shared_ptr<PointCloudComponent>
   setVertices(Eigen::Matrix<float, Eigen::Dynamic, 3, Eigen::RowMajor> const &vertices);
+  Eigen::Matrix<float, Eigen::Dynamic, 3, Eigen::RowMajor> getVertices();
+
   std::shared_ptr<PointCloudComponent> setAttribute(
       std::string const &name,
       Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> const &attribute);

include/sapien/sapien_renderer/sapien_renderer_system.h

@@ -81,6 +81,10 @@ class SapienRendererSystem : public System {
     return std::vector<std::shared_ptr<SapienRenderBodyComponent>>{mRenderBodyComponents.begin(),
                                                                    mRenderBodyComponents.end()};
   }
+  std::vector<std::shared_ptr<PointCloudComponent>> getPointCloudComponents() const {
+    return std::vector<std::shared_ptr<PointCloudComponent>>{mPointCloudComponents.begin(),
+                                                             mPointCloudComponents.end()};
+  }
   std::vector<std::shared_ptr<SapienRenderCameraComponent>> getCameraComponents() const {
     return std::vector<std::shared_ptr<SapienRenderCameraComponent>>{
         mRenderCameraComponents.begin(), mRenderCameraComponents.end()};

python/py_package/__init__.pyi

@@ -34,5 +34,5 @@ from . import utils
 from . import version
 from . import wrapper
 __all__ = ['ActorBuilder', 'ArticulationBuilder', 'Component', 'CudaArray', 'Device', 'Engine', 'Entity', 'Path', 'PinocchioModel', 'Pose', 'SapienRenderer', 'Scene', 'SceneConfig', 'System', 'Widget', 'asset', 'internal_renderer', 'math', 'os', 'physx', 'pkg_resources', 'platform', 'profile', 'pysapien', 'pysapien_pinocchio', 'render', 'set_log_level', 'simsense', 'utils', 'version', 'warn', 'wrapper']
-__version__: str = '3.0.0.dev20240520+55267edc'
+__version__: str = '3.0.0.dev20240521+6b6d61d2'
 __warningregistry__: dict = {'version': 0}

python/py_package/pysapien/render.pyi

@@ -397,6 +397,10 @@ class RenderPointCloudComponent(sapien.pysapien.Component):
         """
     def get_cuda_vertices(self) -> sapien.pysapien.CudaArray:
         ...
+    def get_vertices(self) -> numpy.ndarray[tuple[M, typing.Literal[3]], numpy.dtype[numpy.float32]]:
+        """
+        Get previously set vertices. This function does not reflect any changes directly made to the GPU.
+        """
     def set_attribute(self, name: str, attribute: numpy.ndarray[tuple[M, N], numpy.dtype[numpy.float32]] | list | tuple) -> RenderPointCloudComponent:
         ...
     def set_vertices(self, vertices: numpy.ndarray[tuple[M, typing.Literal[3]], numpy.dtype[numpy.float32]] | list | tuple) -> RenderPointCloudComponent:
@@ -617,6 +621,8 @@ class RenderSystem(sapien.pysapien.System):
         ...
     def get_lights(self) -> list[RenderLightComponent]:
         ...
+    def get_point_clouds(self) -> list[RenderPointCloudComponent]:
+        ...
     def get_render_bodies(self) -> list[RenderBodyComponent]:
         ...
     def set_ambient_light(self, color: numpy.ndarray[typing.Literal[3], numpy.dtype[numpy.float32]] | list[float] | tuple) -> None:
@@ -639,6 +645,9 @@ class RenderSystem(sapien.pysapien.System):
     def lights(self) -> list[RenderLightComponent]:
         ...
     @property
+    def point_clouds(self) -> list[RenderPointCloudComponent]:
+        ...
+    @property
     def render_bodies(self) -> list[RenderBodyComponent]:
         ...
 class RenderSystemGroup:

python/py_package/show_anything.py

@@ -32,11 +32,28 @@ def add_urdf_file(self, filename):
         loader = self.scene.create_urdf_loader()
         loader.load(filename)
 
+    def add_pcd_file(self, filename):
+        import trimesh
+
+        tp = trimesh.load(filename)
+        pcd = sapien.render.RenderPointCloudComponent(tp.vertices.shape[0])
+        pcd.set_vertices(tp.vertices.astype(np.float32))
+        pcd.set_attribute("color", (tp.colors / 255).astype(np.float32))
+        pcd.set_attribute(
+            "scale", np.ones(tp.vertices.shape[0], dtype=np.float32) * 0.1
+        )
+        self.scene.add_entity(sapien.Entity().add_component(pcd))
+
     def get_scene_aabb(self):
         aabb = np.array([[1e10] * 3, [-1e10] * 3])
         for b in self.scene.render_system.render_bodies:
             aabb_ = b.get_global_aabb_fast()
             aabb = np.minimum(aabb[0], aabb_[0]), np.maximum(aabb[1], aabb_[1])
+
+        for p in self.scene.render_system.point_clouds:
+            vs = p.get_vertices()
+            aabb = np.minimum(aabb[0], vs.min(0)), np.maximum(aabb[1], vs.max(0))
+
         return aabb
 
     def focus_scene(self):
@@ -156,6 +173,29 @@ def is_mesh_file(anything):
         return False
     if not any(anything.lower().endswith(x) for x in suffix):
         return False
+
+    if anything.lower().endswith(".ply"):
+        import trimesh
+
+        result = trimesh.load(anything)
+        if isinstance(result, trimesh.PointCloud):
+            return False
+    return True
+
+
+def is_pcd_file(anything):
+    suffix = [".ply", ".pcd"]
+    if not isinstance(anything, str):
+        return False
+    if not any(anything.lower().endswith(x) for x in suffix):
+        return False
+
+    if anything.lower().endswith(".ply"):
+        import trimesh
+
+        result = trimesh.load(anything)
+        if not isinstance(result, trimesh.PointCloud):
+            return False
     return True
 
 
@@ -175,6 +215,10 @@ def show_anything(*args, update_camera=True, loop=True):
             viewer.add_mesh_file(thing)
             continue
 
+        if is_pcd_file(thing):
+            viewer.add_pcd_file(thing)
+            continue
+
         if is_urdf_file(thing):
             viewer.add_urdf_file(thing)
             continue

python/pybind/sapien_renderer.cpp

@@ -597,6 +597,9 @@ This function waits for any pending CUDA operations on cuda stream provided by :
       .def_property_readonly("render_bodies", &SapienRendererSystem::getRenderBodyComponents)
       .def("get_render_bodies", &SapienRendererSystem::getRenderBodyComponents)
 
+      .def_property_readonly("point_clouds", &SapienRendererSystem::getPointCloudComponents)
+      .def("get_point_clouds", &SapienRendererSystem::getPointCloudComponents)
+
       .def_property_readonly("lights", &SapienRendererSystem::getLightComponents)
       .def("get_lights", &SapienRendererSystem::getLightComponents)
 
@@ -906,6 +909,9 @@ This function waits for any pending CUDA operations on cuda stream provided by :
 
   PyRenderPointCloudComponent.def(py::init<uint32_t>(), py::arg("capacity") = 0)
       .def("set_vertices", &PointCloudComponent::setVertices, py::arg("vertices"))
+      .def("get_vertices", &PointCloudComponent::getVertices,
+           "Get previously set vertices. This function does not reflect any changes directly made "
+           "to the GPU.")
       .def("set_attribute", &PointCloudComponent::setAttribute, py::arg("name"),
            py::arg("attribute"))
       .def("get_cuda_vertices", &PointCloudComponent::getCudaArray)

src/sapien_renderer/point_cloud_component.cpp

@@ -16,6 +16,12 @@ std::shared_ptr<PointCloudComponent> PointCloudComponent::setVertices(
   return std::static_pointer_cast<PointCloudComponent>(shared_from_this());
 }
 
+Eigen::Matrix<float, Eigen::Dynamic, 3, Eigen::RowMajor> PointCloudComponent::getVertices() {
+  auto pos = mPointSet->getVertexAttribute("position");
+  return Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, 3, Eigen::RowMajor>>(pos.data(),
+                                                                              pos.size() / 3, 3);
+}
+
 std::shared_ptr<PointCloudComponent> PointCloudComponent::setAttribute(
     std::string const &name,
     Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> const &attribute) {

vulkan_shader/point/composite0.frag

@@ -1,62 +1,114 @@
 #version 450
 
-layout (constant_id = 0) const int NUM_DIRECTIONAL_LIGHTS = 3;
-layout (constant_id = 1) const int NUM_POINT_LIGHTS = 10;
-layout (constant_id = 2) const int NUM_DIRECTIONAL_LIGHT_SHADOWS = 1;
-layout (constant_id = 3) const int NUM_POINT_LIGHT_SHADOWS = 3;
-layout (constant_id = 4) const int NUM_TEXTURED_LIGHT_SHADOWS = 1;
-layout (constant_id = 5) const int NUM_SPOT_LIGHT_SHADOWS = 10;
-layout (constant_id = 6) const int NUM_SPOT_LIGHTS = 10;
-
-layout(set = 0, binding = 0) uniform sampler2D samplerPointColor;
-layout(set = 0, binding = 1) uniform sampler2D samplerPointNormal;
-layout(set = 0, binding = 2) uniform sampler2D samplerPointPosition;
-
-#define SET_NUM 2
-#include "./scene_set.glsl"
-#undef SET_NUM
+layout (constant_id = 0) const float exposure = 1.0;
+
+layout(set = 0, binding = 0) uniform sampler2D samplerLighting;
+layout(set = 0, binding = 1) uniform usampler2D samplerSegmentation;
+layout(set = 0, binding = 2) uniform sampler2D samplerPosition;
+
+layout(set = 0, binding = 3) uniform sampler2D samplerLineDepth;
+layout(set = 0, binding = 4) uniform sampler2D samplerLine;
+layout(set = 0, binding = 5) uniform sampler2D samplerPointDepth;
+layout(set = 0, binding = 6) uniform sampler2D samplerPoint;
+layout(set = 0, binding = 7) uniform sampler2D samplerGbufferDepth;
 
-#define SET_NUM 1
-#include "./camera_set.glsl"
-#undef SET_NUM
 
 layout(location = 0) in vec2 inUV;
-layout(location = 0) out vec4 outPoint;
 
-vec3 diffuseIBL(vec3 albedo, vec3 N) {
-  N = vec3(-N.y, N.z, -N.x);
-  vec3 color = textureLod(samplerEnvironment, N, 5).rgb;
-  return color * albedo;
-}
+layout(location = 0) out vec4 outColor;
+layout(location = 1) out vec4 outDepthLinear;
+layout(location = 2) out vec4 outSegmentationView0;
+layout(location = 3) out vec4 outSegmentationView1;
 
-vec3 specularIBL(vec3 fresnel, float roughness, vec3 N, vec3 V) {
-  float dotNV = max(dot(N, V), 0);
-  vec3 R = 2 * dot(N, V) * N - V;
-  R = vec3(-R.y, R.z, -R.x);
-  vec3 color = textureLod(samplerEnvironment, R, roughness * 5).rgb;
-  vec2 envBRDF = texture(samplerBRDFLUT, vec2(roughness, dotNV)).xy;
-  return color * (fresnel * envBRDF.x + envBRDF.y);
-}
+#define SET_NUM 1
+#include "./camera_set.glsl"
+#undef SET_NUM
 
-void main() {
-  vec3 fresnel = vec3(0.04);
-  float roughness = 0.01;
+vec4 colors[60] = {
+  vec4(0.8,  0.4,  0.4 , 1 ),
+  vec4(0.8,  0.41, 0.24, 1 ),
+  vec4(0.8,  0.75, 0.32, 1 ),
+  vec4(0.6,  0.8,  0.4 , 1 ),
+  vec4(0.35, 0.8,  0.24, 1 ),
+  vec4(0.32, 0.8,  0.51, 1 ),
+  vec4(0.4,  0.8,  0.8 , 1 ),
+  vec4(0.24, 0.63, 0.8 , 1 ),
+  vec4(0.32, 0.37, 0.8 , 1 ),
+  vec4(0.6,  0.4,  0.8 , 1 ),
+  vec4(0.69, 0.24, 0.8 , 1 ),
+  vec4(0.8,  0.32, 0.61, 1 ),
+  vec4(0.8,  0.32, 0.32, 1 ),
+  vec4(0.8,  0.64, 0.4 , 1 ),
+  vec4(0.8,  0.74, 0.24, 1 ),
+  vec4(0.56, 0.8,  0.32, 1 ),
+  vec4(0.4,  0.8,  0.44, 1 ),
+  vec4(0.24, 0.8,  0.46, 1 ),
+  vec4(0.32, 0.8,  0.8 , 1 ),
+  vec4(0.4,  0.56, 0.8 , 1 ),
+  vec4(0.24, 0.3,  0.8 , 1 ),
+  vec4(0.56, 0.32, 0.8 , 1 ),
+  vec4(0.8,  0.4,  0.76, 1 ),
+  vec4(0.8,  0.24, 0.58, 1 ),
+  vec4(0.8,  0.24, 0.24, 1 ),
+  vec4(0.8,  0.61, 0.32, 1 ),
+  vec4(0.72, 0.8,  0.4 , 1 ),
+  vec4(0.52, 0.8,  0.24, 1 ),
+  vec4(0.32, 0.8,  0.37, 1 ),
+  vec4(0.4,  0.8,  0.68, 1 ),
+  vec4(0.24, 0.8,  0.8 , 1 ),
+  vec4(0.32, 0.51, 0.8 , 1 ),
+  vec4(0.48, 0.4,  0.8 , 1 ),
+  vec4(0.52, 0.24, 0.8 , 1 ),
+  vec4(0.8,  0.32, 0.75, 1 ),
+  vec4(0.8,  0.4,  0.52, 1 ),
+  vec4(0.8,  0.52, 0.4 , 1 ),
+  vec4(0.8,  0.58, 0.24, 1 ),
+  vec4(0.7,  0.8,  0.32, 1 ),
+  vec4(0.48, 0.8,  0.4 , 1 ),
+  vec4(0.24, 0.8,  0.3 , 1 ),
+  vec4(0.32, 0.8,  0.66, 1 ),
+  vec4(0.4,  0.68, 0.8 , 1 ),
+  vec4(0.24, 0.46, 0.8 , 1 ),
+  vec4(0.42, 0.32, 0.8 , 1 ),
+  vec4(0.72, 0.4,  0.8 , 1 ),
+  vec4(0.8,  0.24, 0.74, 1 ),
+  vec4(0.8,  0.32, 0.46, 1 ),
+  vec4(0.8,  0.46, 0.32, 1 ),
+  vec4(0.8,  0.76, 0.4 , 1 ),
+  vec4(0.69, 0.8,  0.24, 1 ),
+  vec4(0.42, 0.8,  0.32, 1 ),
+  vec4(0.4,  0.8,  0.56, 1 ),
+  vec4(0.24, 0.8,  0.63, 1 ),
+  vec4(0.32, 0.66, 0.8 , 1 ),
+  vec4(0.4,  0.44, 0.8 , 1 ),
+  vec4(0.35, 0.24, 0.8 , 1 ),
+  vec4(0.7,  0.32, 0.8 , 1 ),
+  vec4(0.8,  0.4,  0.64, 1 ),
+  vec4(0.8,  0.24, 0.41, 1 )
+};
 
-  vec3 color = vec3(0.0);
 
-  vec3 diffuseAlbedo = texture(samplerPointColor, inUV).xyz;
-  vec3 normal = texture(samplerPointNormal, inUV).xyz;
-  vec3 position = texture(samplerPointPosition, inUV).xyz;
+void main() {
+  outColor = texture(samplerPoint, inUV);
+  // outColor = texture(samplerLighting, inUV);
+  // outColor.rgb = pow(outColor.rgb * exposure, vec3(1/2.2));
+  // outColor = clamp(outColor, vec4(0), vec4(1));
 
-  vec3 camDir = -normalize(position);
+  // outDepthLinear.x = -texture(samplerPosition, inUV).z;
 
-  vec3 wnormal = mat3(cameraBuffer.viewMatrixInverse) * normal;
-  color += diffuseIBL(diffuseAlbedo, wnormal);
-  color += specularIBL(fresnel, roughness,
-                       wnormal,
-                       mat3(cameraBuffer.viewMatrixInverse) * camDir);
+  // uvec2 seg = texture(samplerSegmentation, inUV).xy;
+  // outSegmentationView0 = mix(vec4(0,0,0,1), colors[seg.x % 60], sign(seg.x));
+  // outSegmentationView1 = mix(vec4(0,0,0,1), colors[seg.y % 60], sign(seg.y));
 
-  color += sceneBuffer.ambientLight.rgb * diffuseAlbedo;
+  // vec4 lineColor = texture(samplerLine, inUV);
+  // if (texture(samplerLineDepth, inUV).x < 1) {
+  //   outColor = vec4(lineColor.xyz, 1);
+  // }
 
-  outPoint = vec4(color, 1.0);
+  // vec4 pointColor = texture(samplerPoint, inUV);
+  // pointColor.rgb = pow(pointColor.rgb * exposure, vec3(1/2.2));
+  // float pointDepth = texture(samplerPointDepth, inUV).x;
+  // if (pointDepth < 1 && pointDepth < texture(samplerGbufferDepth, inUV).x) {
+  //   outColor = vec4(pointColor.xyz, 1);
+  // }
 }