Merge pull request #1075 from DLR-RM/depth_at_points

feat(camera): Adds method to compute depth at points

blenderproc/api/camera/__init__.py

@@ -5,4 +5,4 @@
 from blenderproc.python.camera.CameraValidation import perform_obstacle_in_view_check, visible_objects, \
     scene_coverage_score, decrease_interest_score, check_novel_pose
 from blenderproc.python.camera.LensDistortionUtility import set_lens_distortion, set_camera_parameters_from_config_file
-from blenderproc.python.camera.CameraProjection import depth_via_raytracing, pointcloud_from_depth, project_points, unproject_points
+from blenderproc.python.camera.CameraProjection import depth_via_raytracing, depth_at_points_via_raytracing, pointcloud_from_depth, project_points, unproject_points

blenderproc/python/camera/CameraProjection.py

@@ -22,6 +22,34 @@ def depth_via_raytracing(bvh_tree: BVHTree, frame: Optional[int] = None, return_
     :param return_dist: If True, a distance image instead of a depth image is returned.
     :return: The depth image with shape [H, W].
     """
+    resolution_x = bpy.context.scene.render.resolution_x
+    resolution_y = bpy.context.scene.render.resolution_y
+
+    # Generate 2D coordinates of all pixels
+    y = np.arange(resolution_y)   
+    x = np.arange(resolution_x)
+    points = np.stack(np.meshgrid(x, y), -1).astype(np.float32)
+
+    # Calc depth at points
+    depth = depth_at_points_via_raytracing(bvh_tree, points.reshape(-1, 2), frame, return_dist)
+
+    # Reshape back into depth image
+    depth = np.reshape(depth, [resolution_y, resolution_x])
+    return depth
+
+
+def depth_at_points_via_raytracing(bvh_tree: BVHTree, points_2d: np.ndarray, frame: Optional[int] = None, return_dist: bool = False) -> np.ndarray:
+    """ Computes the depth values at the given 2D points.
+
+    All points that correspond to rays which do not hit any object are set to inf.
+
+    :param bvh_tree: The BVH tree to use for raytracing.
+    :param points_2d: An array of N 2D points with shape [N, 2].
+    :param frame: The frame number whose assigned camera pose should be used. If None is given, the current frame
+                  is used.
+    :param return_dist: If True, distance values instead of depth are returned.
+    :return: The depth values with shape [N].
+    """
     with KeyFrame(frame):
         cam_ob = bpy.context.scene.camera
         cam = cam_ob.data
@@ -42,25 +70,23 @@ def depth_via_raytracing(bvh_tree: BVHTree, frame: Optional[int] = None, return_
         dists = []
         # Go in discrete grid-like steps over plane
         position = cam2world_matrix.to_translation()
-        for y in range(0, resolution_y):
-            for x in reversed(range(0, resolution_x)):
-                # Compute current point on plane
-                end = frame[0] + vec_x * (x + 0.5) / float(resolution_x) \
-                        + vec_y * (y + 0.5) / float(resolution_y)
-                # Send ray from the camera position through the current point on the plane
-                _, _, _, dist = bvh_tree.ray_cast(position, end - position)
-                if dist is None:
-                    dist = np.inf
-
-                dists.append(dist)
+        for p in points_2d:
+            # Compute current point on plane
+            end = frame[0] + vec_x * (resolution_x - (p[0] + 0.5)) / float(resolution_x) \
+                    + vec_y * (p[1] + 0.5) / float(resolution_y)
+            # Send ray from the camera position through the current point on the plane
+            _, _, _, dist = bvh_tree.ray_cast(position, end - position)
+            if dist is None:
+                dist = np.inf
+
+            dists.append(dist)
         dists = np.array(dists)
-        dists = np.reshape(dists, [resolution_y, resolution_x])
 
         if not return_dist:
-            return dist2depth(dists)
+            return dist2depth(dists, points_2d)
         else:
             return dists
-
+    
 def unproject_points(points_2d: np.ndarray, depth: np.ndarray, frame: Optional[int] = None, depth_cut_off: float = 1e6) -> np.ndarray:
     """ Unproject 2D points into 3D
 

blenderproc/python/postprocessing/PostProcessingUtility.py

@@ -12,11 +12,13 @@
 from blenderproc.python.utility.BlenderUtility import get_all_blender_mesh_objects
 
 
-def dist2depth(dist: Union[List[np.ndarray], np.ndarray]) -> Union[List[np.ndarray], np.ndarray]:
+def dist2depth(dist: Union[List[np.ndarray], np.ndarray], points_2d: Optional[np.ndarray] = None) -> Union[List[np.ndarray], np.ndarray]:
     """
-    Maps a distance image to depth image, also works with a list of images.
+    Maps a distance image to depth image, also works with a list of images or a 1-dim array of dist values.
 
     :param dist: The distance data.
+    :param points_2d: Can be used to specify the 2D points corresponding to the given distance values:
+                      Is necessary, if the given distance data is not a full distance image.
     :return: The depth data
     """
 
@@ -28,7 +30,10 @@ def dist2depth(dist: Union[List[np.ndarray], np.ndarray]) -> Union[List[np.ndarr
     K = CameraUtility.get_intrinsics_as_K_matrix()
     f, cx, cy = K[0, 0], K[0, 2], K[1, 2]
 
-    xs, ys = np.meshgrid(np.arange(dist.shape[1]), np.arange(dist.shape[0]))
+    if points_2d is None:
+        xs, ys = np.meshgrid(np.arange(dist.shape[1]), np.arange(dist.shape[0]))
+    else:
+        xs, ys = points_2d[:, 0], points_2d[:, 1]
 
     # coordinate distances to principal point
     x_opt = np.abs(xs - cx)

tests/testCameraProjection.py

@@ -33,6 +33,32 @@ def test_unproject_project(self):
 
         np.testing.assert_almost_equal(pixels, pixels_gt, decimal=3)
 
+
+    def test_depth_at_points_via_raytracing(self):
+        """ Test if depth_at_points_via_raytracing leads to the same results as depth_via_raytracing + unproject + project
+        """
+        bproc.clean_up(True)
+        resource_folder = os.path.join("examples", "resources")
+        objs = bproc.loader.load_obj(os.path.join(resource_folder, "scene.obj"))
+
+        cam2world_matrix = np.array([[1.0, 0.0, 0.0, 0.0], [0.0, 0.2674988806247711, -0.9635581970214844, -13.741], [-0.0, 0.9635581970214844, 0.2674988806247711, 4.1242], [0.0, 0.0, 0.0, 1.0]])
+        bproc.camera.add_camera_pose(cam2world_matrix)
+        bproc.camera.set_resolution(640, 480)
+
+        bvh_tree = bproc.object.create_bvh_tree_multi_objects(objs)
+
+        depth = bproc.camera.depth_via_raytracing(bvh_tree)
+        pc = bproc.camera.pointcloud_from_depth(depth)
+        pixels = bproc.camera.project_points(pc.reshape(-1, 3))
+
+        depth2 = bproc.camera.depth_at_points_via_raytracing(bvh_tree, pixels)
+        depth2[np.isnan(depth2)] = np.inf
+        pc2 = bproc.camera.unproject_points(pixels, depth2)
+
+        np.testing.assert_almost_equal(depth.flatten(), depth2, decimal=3)
+        np.testing.assert_almost_equal(pc.reshape(-1, 3), pc2, decimal=3)
+
+
     def test_depth_via_raytracing(self):
         """ Tests if depth image via raytracing and rendered depth image are identical.
         """
@@ -59,7 +85,6 @@ def test_depth_via_raytracing(self):
         data["depth"][0][data["depth"][0] >= 65504] = np.inf
 
         diff = np.abs(depth[~np.isinf(depth)] - data["depth"][0][~np.isinf(depth)])
-
         self.assertTrue(np.median(diff) < 1e-4)
         self.assertTrue((diff < 1e-4).mean() > 0.99)