Merge pull request #14 from yuecideng/yuecideng/RayCastRenderer

Add RayRastRenderer

README.md

@@ -23,6 +23,7 @@ Core modules:
     2. 3D rigid transformation solver including SVD, RANSAC and [TEASERPP](https://github.com/MIT-SPARK/TEASER-plusplus).
 - `pose_estimation`: 
     1. Point Pair Features (PPF) based 6D pose estimator. (This implementation is evaluated on Linemod, Linemod-Occluded and YCB-Video dataset, the performance can be found in [BOP Leaderboards/PPF-3D-ICP](https://bop.felk.cvut.cz/leaderboards/))
+    2. A RayCastRenderer, whcih is useful for partial view point clouds, depth map and instance map generation.
 
 - `segmentation`: 
     1. Proximity extraction in scalable implementation with different vriants, including distance, and normal angle.
@@ -73,59 +74,59 @@ import misc3d as m3d
 ```
 
 ```python
-# estimate normals inplace
+# Estimate normals inplace.
 m3d.common.estimate_normals(pcd, (640, 480), 3)
 ```
 
 ```python
-# ransac for primitives fitting
+# Ransac for primitives fitting.
 w, indices = m3d.common.fit_plane(pcd, 0.01, 100)
 w, indices = m3d.common.fit_sphere(pcd, 0.01, 100)
 w, indices = m3d.common.fit_cylinder(pcd, 0.01, 100)
 ```
 
 ```python
-# farthest point sampling
+# Farthest point sampling.
 indices = m3d.preprocessing.farthest_point_sampling(pcd, 1000)
 ```
 
 ```python
-# crop ROI point clouds
+# Crop ROI point clouds.
 pcd_roi = m3d.preprocessing.crop_roi_pointcloud(pcd, (500, 300, 600, 400), (640, 480))
 ```
 
 ```python
-# project point clouds into a plane
+# Project point clouds into a plane.
 pcd_plane = m3d.preprocessing.project_into_plane(pcd)
 ```
 
 ```python
-# boundary points detection
+# Boundary points detection.
 index = m3d.features.detect_boundary_points(
     pcd, o3d.geometry.KDTreeSearchParamHybrid(0.02, 30))
 boundary = pcd.select_by_index(index)
 ```
 
 ```python
-# features matching using FLANN or ANNOY
+# Features matching using FLANN or ANNOY
 # `fpfh_src` is open3d.pipeline.registration.Feature instance which is computed using FPFH 3d descriptor.
 index1, index2 = m3d.registration.match_correspondence(fpfh_src, fpfh_dst, m3d.registration.MatchMethod.FLANN)
 index1, index2 = m3d.registration.match_correspondence(fpfh_src, fpfh_dst, m3d.registration.MatchMethod.ANNOY)
 ```
 
 ```python
-# solve 3d rigid transformation
-# ransac solver
+# Solve 3d rigid transformation.
+# Ransac solver
 pose = m3d.registration.compute_transformation_ransac(pc_src, pc_dst, (index1, index2), 0.03, 100000)
-# svd solver
+# SVD solver
 pose = m3d.registration.compute_transformation_svd(pc_src, pc_dst)
-# teaser solver
+# Teaser solver
 pose = m3d.registration.compute_transformation_teaser(pc_src, pc_dst, 0.01)
 ```
 
 ```python
-# ppf pose estimator
-# init config for ppf pose estimator
+# PPF pose estimator.
+# Init config for ppf pose estimator.
 config = m3d.pose_estimation.PPFEstimatorConfig()
 config.training_param.rel_sample_dist = 0.04
 config.score_thresh = 0.1
@@ -135,6 +136,21 @@ ret = ppf.train(model)
 ret, results = ppf.estimate(scene)
 ```
 
+```python
+# Create a ray cast renderer.
+intrinsic = o3d.camera.PinholeCameraIntrinsic(
+    640, 480, 572.4114, 573.5704, 325.2611, 242.0489)
+renderer = m3d.pose_estimation.RayCastRenderer(intrinsic)
+
+# Cast rays for single mesh with a associated pose.
+ret = renderer.cast_rays([mesh], [pose])
+depth = renderer.get_depth_map()
+# Convert depth map into numpy array. 
+depth = depth.numpy()
+# Get partial view point clouds of the mesh in the scene.
+pcd = renderer.get_point_cloud()
+```
+
 ```python
 # proximity extraction
 pe = m3d.segmentation.ProximityExtractor(100)

examples/cpp/CMakeLists.txt

@@ -7,6 +7,9 @@ target_link_libraries(normal_estimation PUBLIC misc3d::misc3d)
 add_executable(ppf_estimator ppf_estimator.cpp)
 target_link_libraries(ppf_estimator PUBLIC misc3d::misc3d)
 
+add_executable(ray_cast_rendering ray_cast_rendering.cpp)
+target_link_libraries(ray_cast_rendering PUBLIC misc3d::misc3d)
+
 add_executable(preprocessing preprocessing.cpp)
 target_link_libraries(preprocessing PUBLIC misc3d::misc3d)
 
@@ -22,7 +25,7 @@ target_link_libraries(segment_plane_iterative PUBLIC misc3d::misc3d)
 add_executable(transform_estimation transform_estimation.cpp)
 target_link_libraries(transform_estimation PUBLIC misc3d::misc3d)
 
-install(TARGETS farthest_point_sampling normal_estimation ppf_estimator 
+install(TARGETS farthest_point_sampling normal_estimation ppf_estimator ray_cast_rendering
                 preprocessing ransac_and_boundary segmentation segment_plane_iterative 
                 transform_estimation
                 RUNTIME DESTINATION "${CMAKE_INSTALL_PREFIX}/misc3d/bin")

examples/cpp/ray_cast_rendering.cpp

@@ -0,0 +1,51 @@
+#include <iostream>
+#include <memory>
+
+#include <misc3d/pose_estimation/ray_cast_renderer.h>
+#include <misc3d/utils.h>
+#include <misc3d/vis/vis_utils.h>
+#include <open3d/camera/PinholeCameraIntrinsic.h>
+#include <open3d/geometry/PointCloud.h>
+#include <open3d/geometry/TriangleMesh.h>
+#include <open3d/io/TriangleMeshIO.h>
+#include <open3d/visualization/utility/DrawGeometry.h>
+
+int main(int argc, char *argv[]) {
+    const open3d::camera::PinholeCameraIntrinsic camera_intrinsic(
+        640, 480, 572.4114, 573.5704, 325.2611, 242.0489);
+
+    misc3d::pose_estimation::RayCastRenderer renderer(camera_intrinsic);
+
+    open3d::geometry::TriangleMesh mesh;
+    bool ret = open3d::io::ReadTriangleMesh(
+        "../examples/data/pose_estimation/model/obj.ply", mesh);
+    mesh.Scale(0.001, Eigen::Vector3d::Zero());
+    const open3d::geometry::TriangleMesh mesh2 = mesh;
+
+    Eigen::Matrix<double, 4, 4, Eigen::RowMajor> pose, pose2;
+    pose << 0.29493218, 0.95551309, 0.00312103, -0.14527225, 0.89692822,
+        -0.27572004, -0.34568516, 0.12533501, -0.32944616, 0.10475302,
+        -0.93834537, 0.99371838, 0., 0., 0., 1.;
+    pose2 << 0.29493218, 0.95551309, 0.00312103, -0.04527225, 0.89692822,
+        -0.27572004, -0.34568516, 0.02533501, -0.32944616, 0.10475302,
+        -0.93834537, 0.99371838, 0., 0., 0., 1.;
+
+    misc3d::Timer timer;
+    timer.Start();
+    ret = renderer.CastRays({mesh, mesh2}, {pose, pose2});
+    std::cout << "Ray cast time: " << timer.Stop() << std::endl;
+
+    auto instance_pcds = renderer.GetInstancePointCloud();
+    instance_pcds[0].PaintUniformColor({1, 0, 0});
+    instance_pcds[1].PaintUniformColor({0, 1, 0});
+    const auto axis = open3d::geometry::TriangleMesh::CreateCoordinateFrame(
+        0.1, Eigen::Vector3d::Zero());
+
+    open3d::visualization::DrawGeometries(
+        {std::make_shared<const open3d::geometry::PointCloud>(instance_pcds[0]),
+         std::make_shared<const open3d::geometry::PointCloud>(instance_pcds[1]),
+         axis},
+        "Ray Cast Rendering");
+
+    return 0;
+}

examples/python/ray_cast_rendering.py

@@ -0,0 +1,54 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+import numpy as np
+import open3d as o3d
+from torch import embedding
+import misc3d as m3d
+from matplotlib import pyplot as plt
+import time
+
+# Create a ray cast renderer.
+camera_intrinsic = o3d.camera.PinholeCameraIntrinsic(640, 480, 572.4114,
+                                                     573.5704, 325.2611,
+                                                     242.0489)
+renderer = m3d.pose_estimation.RayCastRenderer(camera_intrinsic)
+
+# Load mesh and create a copy.
+mesh = o3d.io.read_triangle_mesh('../data/pose_estimation/model/obj.ply')
+mesh = mesh.scale(0.001, np.array([0, 0, 0]))
+mesh2 = mesh
+
+# Create mesh pose relative to camera.
+pose = np.array([[0.29493218, 0.95551309, 0.00312103, -0.14527225],
+                 [0.89692822, -0.27572004, -0.34568516, 0.12533501],
+                 [-0.32944616, 0.10475302, -0.93834537, 0.99371838],
+                 [0., 0., 0., 1.]])
+pose2 = np.array([[0.29493218, 0.95551309, 0.00312103, -0.04527225],
+                  [0.89692822, -0.27572004, -0.34568516, 0.02533501],
+                  [-0.32944616, 0.10475302, -0.93834537, 0.99371838],
+                  [0., 0., 0., 1.]])
+
+# Cast rays for meshes in the scene.
+t0 = time.time()
+ret = renderer.cast_rays([mesh, mesh2], [pose, pose2])
+print('cast_rays:', time.time() - t0)
+
+# Visualize the depth and instance map. Both of them are open3d.core.Tensor type,
+# which can be converted into numpy array.
+depth = renderer.get_depth_map().numpy()
+instance_map = renderer.get_instance_map().numpy()
+instance_pcds = renderer.get_instance_point_cloud()
+instance_pcds[0].paint_uniform_color([1, 0, 0])
+instance_pcds[1].paint_uniform_color([0, 1, 0])
+
+plt.subplot(1, 2, 1)
+plt.imshow(depth)
+plt.subplot(1, 2, 2)
+plt.imshow(instance_map, vmax=2)
+plt.show()
+
+axis = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1,
+                                                         origin=[0, 0, 0])
+o3d.visualization.draw_geometries([instance_pcds[0], instance_pcds[1], axis],
+                                  "Ray Cast Rendering")

include/misc3d/pose_estimation/ray_cast_renderer.h

@@ -0,0 +1,90 @@
+#pragma once
+
+#include <unordered_map>
+#include <vector>
+
+#include <open3d/camera/PinholeCameraIntrinsic.h>
+#include <open3d/geometry/Image.h>
+#include <open3d/t/geometry/RaycastingScene.h>
+#include <open3d/t/geometry/TriangleMesh.h>
+#include <Eigen/Dense>
+
+namespace misc3d {
+
+namespace pose_estimation {
+
+/**
+ * @brief Ray cast renderer for depth map, instance map, point clouds and
+ * normals map generation.
+ * Currently only support CPU.
+ * TODO: Add CUDA suppot.
+ */
+class RayCastRenderer {
+public:
+    /**
+     * @brief Construct a new Ray Cast Renderer defined by a camera intrinsic.
+     * The rays start from the origin [0, 0, 0] will be created based on the
+     * intrinsic inside constructor.
+     *
+     * @param intrinsic
+     */
+    explicit RayCastRenderer(
+        const open3d::camera::PinholeCameraIntrinsic &intrinsic);
+
+    virtual ~RayCastRenderer() {}
+
+    /**
+     * @brief Compute the first intersection of the rays with the scene.
+     * The result is stored in a map with the following keys:
+     * \b t_hit: the distance of each hit
+     * \b geometry_ids: the instance id of the geometries
+     * \b primitive_normals: the normals of the geometries
+     *
+     * @param mesh_list
+     * @param pose_list
+     * @return true
+     * @return false
+     */
+    bool CastRays(const std::vector<open3d::geometry::TriangleMesh> &mesh_list,
+                  const std::vector<Eigen::Matrix4d> &pose_list);
+
+    /**
+     * @brief Get Depth Map
+     *
+     * @return open3d::core::Tensor
+     */
+    open3d::core::Tensor GetDepthMap() const;
+
+    /**
+     * @brief Get Instance Map
+     * The id is from 0 to num_instances - 1.
+     *
+     * @return open3d::core::Tensor
+     */
+    open3d::core::Tensor GetInstanceMap() const;
+
+    /**
+     * @brief Get PointCloud from depth map with valid value.
+     *
+     * @return open3d::geometry::PointCloud
+     */
+    open3d::geometry::PointCloud GetPointCloud() const;
+
+    /**
+     * @brief Get Instance PointCloud individully.
+     *
+     * @return std::vector<open3d::geometry::PointCloud>
+     */
+    std::vector<open3d::geometry::PointCloud> GetInstancePointCloud() const;
+
+private:
+    std::unordered_map<std::string, open3d::core::Tensor> ray_cast_results_;
+    int width_;
+    int height_;
+    open3d::core::Tensor rays_;
+    size_t num_instance_;
+};
+
+}  // namespace pose_estimation
+
+}  // namespace misc3d

python/py_misc3d.cpp

@@ -13,6 +13,9 @@ PYBIND11_MODULE(py_misc3d, m) {
     py::object o3d_camera =
         (py::object)py::module_::import("open3d").attr("camera");
 
+    py::object o3d_core =
+        (py::object)py::module_::import("open3d").attr("core");
+
     py::object o3d_vis =
         (py::object)py::module_::import("open3d").attr("visualization");
 

python/py_pose_estimation.cpp

@@ -1,6 +1,7 @@
 #include <py_misc3d.h>
 
 #include <misc3d/pose_estimation/ppf_estimation.h>
+#include <misc3d/pose_estimation/ray_cast_renderer.h>
 #include <misc3d/utils.h>
 
 namespace misc3d {
@@ -104,6 +105,15 @@ void pybind_pose_estimation(py::module &m) {
     config.def_readwrite("score_thresh", &PPFEstimatorConfig::score_thresh_);
     config.def_readwrite("num_result", &PPFEstimatorConfig::num_result_);
     config.def_readwrite("object_id", &PPFEstimatorConfig::object_id_);
+
+    py::class_<RayCastRenderer>(m, "RayCastRenderer")
+        .def(py::init<const open3d::camera::PinholeCameraIntrinsic &>())
+        .def("cast_rays", &RayCastRenderer::CastRays)
+        .def("get_depth_map", &RayCastRenderer::GetDepthMap)
+        .def("get_instance_map", &RayCastRenderer::GetInstanceMap)
+        .def("get_point_cloud", &RayCastRenderer::GetPointCloud)
+        .def("get_instance_point_cloud",
+             &RayCastRenderer::GetInstancePointCloud);
 }
 
 }  // namespace pose_estimation

src/CMakeLists.txt

@@ -10,6 +10,7 @@ set(SRCS utils.cpp
          transform_estimation.cpp
          correspondence_matching.cpp
          ppf_estimation.cpp
+         ray_cast_renderer.cpp
          vis_utils.cpp)
 
 add_library(misc3d ${SRCS})