Adds bounding boxes to blender_kitti

README.md

@@ -6,6 +6,7 @@
 |:-------------------------:|:-------------------------:|
 | ![KITTI Point Cloud](img/blender_kitti_render_voxels_main.png?raw=true "Main view voxels") |![KITTI Point Cloud](img/blender_kitti_render_voxels_top.png?raw=true "Top view voxels") |
 | ![KITTI Scene Flow](img/blender_kitti_render_scene_flow_main.png?raw=true "Main view scene flow") |![KITTI Scene Flow](img/blender_kitti_render_scene_flow_top.png?raw=true "Top view scene flow") |
+| ![KITTI Bounding Boxes](img/blender_kitti_render_boxes_main.png?raw=true "Main view boxes") |![KITTI Bounding Boxes](img/blender_kitti_render_boxes_top.png?raw=true "Top view bounding boxes") |
 
 ## About
 
@@ -17,9 +18,9 @@ and all colors have the exact RGB-value as specified. All particles can be color
 * Performance of the scrips is acceptable. It should not take much longer than a second to create a 100k point cloud.
 
 Together, these qualities enable `blender-kitti` to render large scale data from the KITTI dataset
-(hence the name) or related datasets. 
+(hence the name) or related datasets.
 
-When it comes to visualization, everyone has a different usecase. So this is not a 
+When it comes to visualization, everyone has a different usecase. So this is not a
 one-fits-all solution but rather a collection of techniques that can be adapted
 to individual usecases.
 
@@ -70,11 +71,20 @@ $ blender --background --python-console
 >>> blender_kitti_examples.render_kitti_scene_flow()
 ```
 
+Render the bundled KITTI point cloud with some random bounding boxes (with random colors). This writes two image files to the `/tmp` folder.
+
+```
+$ blender --background --python-console
+
+>>> import blender_kitti_examples
+>>> blender_kitti_examples.render_kitti_bounding_boxes()
+
+
 ## Work on a scene in Blender
 
 You can import and use `blender-kitti` in the python console window in the Blender-GUI
 itself to work on a given scene.
- 
+
 ```
 # Create a random [Nx3] numpy array and add as point cloud to a scene in blender.
 import bpy

blender_kitti/__init__.py

@@ -3,7 +3,7 @@
 __author__ = """Christoph Rist"""
 __email__ = "c.rist@posteo.de"
 
-from .particles import add_voxels, add_point_cloud, add_flow_mesh
+from .particles import add_voxels, add_point_cloud, add_flow_mesh, add_boxes
 from .scene_setup import setup_scene, add_cameras_default
 from .system_setup import setup_system
 from .object_spotlight import add_spotlight_ground
@@ -12,6 +12,7 @@
 
 
 __all__ = [
+    "add_boxes",
     "add_voxels",
     "add_point_cloud",
     "add_cameras_default",

blender_kitti/particles.py

@@ -184,7 +184,6 @@ def _add_material_to_particle(name_prefix, colors, obj_particle, material=None):
 
 
 def create_cube(name_prefix: str, *, edge_length: float = 0.16):
-
     bm = bmesh.new()
     bmesh.ops.create_cube(
         bm,
@@ -207,7 +206,6 @@ def create_icosphere(
     radius: float = 0.02,
     use_smooth: bool = True,
 ):
-
     bm = bmesh.new()
     bmesh.ops.create_icosphere(
         bm,
@@ -418,7 +416,7 @@ def bmesh_join(list_of_bmeshes, list_of_matrices, *, normal_update=False, bmesh)
 
 
 def simple_scale_matrix(factor: np.array, direction: np.array):
-    dir_len = np.sqrt(np.sum(direction ** 2))
+    dir_len = np.sqrt(np.sum(direction**2))
     assert dir_len > 0.0, "direction vector of scale matrix may not have length zero"
     normalized_dir = direction / dir_len
     factor = 1.0 - factor
@@ -544,7 +542,7 @@ def add_flow_mesh(
 
         flow_vec = flow[flow_vec_idx]
 
-        flow_vec_unit = flow_vec / np.sqrt(np.sum(flow_vec ** 2))
+        flow_vec_unit = flow_vec / np.sqrt(np.sum(flow_vec**2))
 
         # rotation matrix R that rotates unit vector a onto unit vector b.
         v = np.cross(arrow_head_unit, flow_vec_unit)
@@ -622,3 +620,101 @@ def add_flow_mesh(
         scene.collection.objects.link(obj)
 
     return obj
+
+
+def create_cube_with_wireframe(
+    position: np.ndarray,
+    scale: np.ndarray,
+    rotation: np.ndarray,
+    wireframe_scale: float,
+    color: np.ndarray,
+):
+    # create a mesh cube
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.ops.mesh.primitive_cube_add(
+        size=1, enter_editmode=False, align="WORLD", location=position
+    )
+    cube = bpy.context.object
+
+    # set scale and rotation
+    cube.scale = scale
+    cube.rotation_euler = rotation
+    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
+
+    # add wireframe modifier
+    bpy.ops.object.modifier_add(type="WIREFRAME")
+    wireframe_modifier = cube.modifiers[-1]
+    wireframe_modifier.thickness = wireframe_scale
+
+    # create a new material with the given color
+    material = bpy.data.materials.new(name="CubeMaterial")
+    material.use_nodes = False
+    material.diffuse_color = color
+
+    # assign the material to the cube
+    if len(cube.data.materials) > 0:
+        cube.data.materials[0] = material
+    else:
+        cube.data.materials.append(material)
+
+    return cube
+
+
+def add_boxes(
+    *,
+    scene,
+    boxes: typing.Dict[str, np.ndarray],
+    box_colors_rgba_f64: np.ndarray,
+    confidence_threshold: float = 0.0,
+    bounding_box_wire_frame_scale: float = 0.2,
+    verbose: bool = False,
+):
+    """
+    supports only boxes with yaw rotation
+
+    scene: blender py scene
+    boxes: dictionairy with
+        * 'pos': np.ndarray with shape [num_boxes, 3] (i.e. box positions in 3d)
+        * 'rot': np.ndarray with shape [num_boxes, 1] (i.e. box yaw angles)
+        * 'dims': np.ndarray with shape [num_boxes, 3] (i.e. box size length, width, height)
+        * 'probs': np.ndarray with shape [num_boxes, 1] (i.e. box confidence)
+    box_colors_rgba_f64: np.ndarray with shape [num_boxes, 4], i.e. a color for each box
+    confidence_threshold: boxes below this threshold are discarded
+    bounding_box_wire_frame_scale: this is the thickness of the box wireframe (in meters I think)
+    """
+
+    assert "pos" in boxes, "need box positions with key 'pos' to work!"
+    assert "dims" in boxes, "need box dimensions with key 'dims' to work!"
+    assert "rot" in boxes, "need box rotations (yaw angle) with key 'rot' to work!"
+
+    assert (
+        box_colors_rgba_f64 <= 1.0
+    ).all(), "this code is only tested with f64 colors <= 1.0!"
+
+    num_boxes = boxes["pos"].shape[0]
+    for box_idx in range(num_boxes):
+        box_confidence = np.squeeze(boxes["probs"][box_idx])
+        if box_confidence < confidence_threshold:
+            if verbose:
+                print(f"Discarding box #{box_idx} with confidence {box_confidence}")
+            continue
+        if verbose:
+            print(
+                f"Add box #{box_idx} at position: ",
+                boxes["pos"][box_idx],
+                ", rotation: ",
+                boxes["rot"][box_idx],
+                f", confidence: {box_confidence}",
+            )
+        cube = create_cube_with_wireframe(
+            position=boxes["pos"][box_idx],
+            scale=boxes["dims"][box_idx],
+            rotation=(
+                0.0,
+                0.0,
+                boxes["rot"][box_idx],
+            ),
+            wireframe_scale=bounding_box_wire_frame_scale,
+            color=box_colors_rgba_f64[box_idx],
+        )
+        scene.collection.objects.link(cube)

blender_kitti_examples/__init__.py

@@ -4,12 +4,14 @@
 __email__ = "c.rist@posteo.de"
 
 from .example_render_kitti import (
+    render_kitti_bounding_boxes,
     render_kitti_point_cloud,
     render_kitti_voxels,
     render_kitti_scene_flow,
 )
 
 __all__ = [
+    "render_kitti_bounding_boxes",
     "render_kitti_point_cloud",
     "render_kitti_voxels",
     "render_kitti_scene_flow",

blender_kitti_examples/example_render_kitti.py

@@ -5,6 +5,7 @@
 import numpy as np
 from blender_kitti.bpy_helper import needs_bpy_bmesh
 from blender_kitti import (
+    add_boxes,
     add_point_cloud,
     add_voxels,
     setup_scene,
@@ -20,6 +21,7 @@
     get_semantic_kitti_voxels,
     get_pseudo_flow,
 )
+import bpy
 
 
 def dry_render(_scene, cameras, output_path):
@@ -57,11 +59,10 @@ def render(scene, cameras, output_path, *, bpy):
 
 
 def render_kitti_point_cloud(gpu_compute=False):
-
     scene = setup_scene()
     cameras = add_cameras_default(scene)
 
-    scene.view_layers["View Layer"].cycles.use_denoising = True
+    scene.view_layers["ViewLayer"].cycles.use_denoising = True
     scene.render.resolution_percentage = 100
     scene.render.resolution_x = 640
     scene.render.resolution_y = 480
@@ -75,16 +76,14 @@ def render_kitti_point_cloud(gpu_compute=False):
 
     point_cloud, colors = get_semantic_kitti_point_cloud()
     _ = add_point_cloud(points=point_cloud, colors=colors, scene=scene)
-    render(scene, cameras, "/tmp/blender_kitti_render_point_cloud_{}.png")
+    render(scene, cameras, "/tmp/blender_kitti_render_point_cloud_{}.png", bpy=bpy)
 
 
 def render_kitti_scene_flow(gpu_compute=False):
-
     scene = setup_scene()
     cameras = add_cameras_default(scene)
 
-    scene.view_layers["View Layer"].cycles.use_denoising = False
-    scene.view_layers["View Layer"].cycles.use_denoising = True
+    scene.view_layers["ViewLayer"].cycles.use_denoising = True
     scene.render.resolution_percentage = 100
     scene.render.resolution_x = 640
     scene.render.resolution_y = 480
@@ -113,8 +112,119 @@ def render_kitti_scene_flow(gpu_compute=False):
     render(scene, cameras, "/tmp/blender_kitti_render_scene_flow_{}.png")
 
 
-def render_kitti_voxels(gpu_compute=False):
+def render_kitti_bounding_boxes(gpu_compute=True):
+    scene = setup_scene()
+    cameras = add_cameras_default(scene)
+    scene.view_layers["ViewLayer"].cycles.use_denoising = True
+    scene.render.resolution_percentage = 100
+    scene.render.resolution_x = 1280
+    scene.render.resolution_y = 1024
+    # alpha background
+    scene.render.film_transparent = True
+    #
+    if gpu_compute:
+        scene.cycles.device = "GPU"
+    else:
+        scene.cycles.device = "CPU"
 
+    point_cloud, colors = get_semantic_kitti_point_cloud()
+    _ = add_point_cloud(points=point_cloud, colors=colors, scene=scene)
+
+    box_range_max = point_cloud.max(axis=0) / 2
+    box_range_min = point_cloud.min(axis=0) / 2
+    print(f"Create random boxes in range {box_range_min} - {box_range_max}")
+    num_pred_boxes = 20
+    num_gt_boxes = 10
+
+    max_box_dims = np.array([7.0, 3.0, 2.0])
+
+    # random boxes
+    boxes_pred = {
+        "pos": box_range_min[None, ...]
+        + np.random.rand(num_pred_boxes, 3)
+        * (box_range_max - box_range_min)[None, ...],
+        "dims": 1 + np.random.rand(num_pred_boxes, 3) * max_box_dims[None, ...],
+        "rot": 2 * np.pi * np.random.rand(num_pred_boxes, 1),
+        "probs": np.random.rand(num_pred_boxes, 1),
+    }
+
+    pred_box_colors = np.random.rand(boxes_pred["pos"].shape[0], 4)
+    pred_box_colors[:, -1] = 1.0
+
+    _ = add_boxes(
+        scene=scene,
+        boxes=boxes_pred,
+        box_colors_rgba_f64=pred_box_colors,
+        confidence_threshold=0.3,
+        verbose=True,
+    )
+
+    boxes_gt = {
+        "pos": box_range_min[None, ...]
+        + np.random.rand(num_gt_boxes, 3) * (box_range_max - box_range_min)[None, ...],
+        "dims": np.random.rand(num_gt_boxes, 3) * max_box_dims[None, ...],
+        "rot": 2 * np.pi * np.random.rand(num_gt_boxes, 1),
+        "probs": np.ones((num_gt_boxes, 1)),
+    }
+
+    boxes_gt = {
+        "pos": np.array(
+            [
+                [
+                    5.0,
+                    0.0,
+                    0.0,
+                ],
+                [
+                    5.0,
+                    10.0,
+                    0.0,
+                ],
+            ]
+        ),
+        "dims": np.array(
+            [
+                [
+                    3.0,
+                    1.0,
+                    2.0,
+                ],
+                [
+                    5.0,
+                    2.0,
+                    2.0,
+                ],
+            ]
+        ),
+        "rot": np.array(
+            [
+                [
+                    np.pi / 4,
+                ],
+                [3 * np.pi / 4],
+            ]
+        ),
+        "probs": np.ones((2, 1)),
+    }
+
+    gt_box_colors = np.ones((boxes_gt["pos"].shape[0], 4)) * np.array(
+        [
+            [1.0, 0.0, 0.0, 1.0],
+        ]
+    )
+
+    _ = add_boxes(
+        scene=scene,
+        boxes=boxes_gt,
+        box_colors_rgba_f64=gt_box_colors,
+        confidence_threshold=0.3,
+        verbose=True,
+    )
+
+    render(scene, cameras, "/tmp/blender_kitti_render_boxes_{}.png", bpy=bpy)
+
+
+def render_kitti_voxels(gpu_compute=False):
     scene = setup_scene()
     cam_main = create_camera_perspective(
         location=(2.86, 17.52, 3.74),
@@ -128,7 +238,7 @@ def render_kitti_voxels(gpu_compute=False):
     scene.collection.objects.link(cam_top)
     scene.camera = cam_main
 
-    scene.view_layers["View Layer"].cycles.use_denoising = True
+    scene.view_layers["ViewLayer"].cycles.use_denoising = True
     scene.render.resolution_percentage = 100
     scene.render.resolution_x = 640
     scene.render.resolution_y = 480