Draw order for python api (#2138)

* draw order for images and points

* python api for draw order on lines & rects. improved api demo

* py formatting & linting

crates/re_log_types/src/component_types/mod.rs

@@ -73,12 +73,13 @@ pub use vec::{Vec2D, Vec3D, Vec4D};
 
 lazy_static! {
     //TODO(john): use a run-time type registry
-    static ref FIELDS: [Field; 26] = [
+    static ref FIELDS: [Field; 27] = [
         <AnnotationContext as Component>::field(),
         <Arrow3D as Component>::field(),
         <Box3D as Component>::field(),
         <ClassId as Component>::field(),
         <ColorRGBA as Component>::field(),
+        <DrawOrder as Component>::field(),
         <InstanceKey as Component>::field(),
         <KeypointId as Component>::field(),
         <Label as Component>::field(),
@@ -94,8 +95,8 @@ lazy_static! {
         <ScalarPlotProps as Component>::field(),
         <Size3D as Component>::field(),
         <Tensor as Component>::field(),
-        <TextEntry as Component>::field(),
         <TextBox as Component>::field(),
+        <TextEntry as Component>::field(),
         <Transform as Component>::field(),
         <Vec2D as Component>::field(),
         <Vec3D as Component>::field(),

examples/python/api_demo/main.py

@@ -77,6 +77,44 @@ def run_segmentation() -> None:
     rr.log_text_entry("logs/seg_demo_log", "label1 disappears and everything with label3 is now default colored again")
 
 
+def run_2d_layering() -> None:
+    rr.set_time_seconds("sim_time", 1)
+
+    # Large gray background.
+    img = np.full((512, 512), 64, dtype="uint8")
+    rr.log_image("2d_layering/background", img, draw_order=0.0)
+
+    # Smaller gradient in the middle.
+    img = np.zeros((256, 256, 3), dtype="uint8")
+    img[:, :, 0] = np.linspace(0, 255, 256, dtype="uint8")
+    img[:, :, 1] = np.linspace(0, 255, 256, dtype="uint8")
+    img[:, :, 1] = img[:, :, 1].transpose()
+    rr.log_image("2d_layering/middle_gradient", img, draw_order=1.0)
+
+    # Slightly smaller blue in the middle, on the same layer as the previous.
+    img = np.full((192, 192, 3), (0, 0, 255), dtype="uint8")
+    rr.log_image("2d_layering/middle_blue", img, draw_order=1.0)
+
+    # Small white on top.
+    img = np.full((128, 128), 255, dtype="uint8")
+    rr.log_image("2d_layering/top", img, draw_order=2.0)
+
+    # Rectangle in between the top and the middle.
+    rr.log_rect("2d_layering/rect_between_top_and_middle", (64, 64, 256, 256), draw_order=1.5)
+
+    # Lines behind the rectangle.
+    rr.log_line_strip(
+        "2d_layering/lines_behind_rect", [(i * 20, i % 2 * 100 + 100) for i in range(20)], draw_order=1.25
+    )
+
+    # And some points in front of the rectangle.
+    rr.log_points(
+        "2d_layering/points_between_top_and_middle",
+        [(32.0 + int(i / 16) * 16.0, 64.0 + (i % 16) * 16.0) for i in range(16 * 16)],
+        draw_order=1.51,
+    )
+
+
 def run_2d_lines() -> None:
     import numpy as np
 
@@ -333,6 +371,7 @@ def main() -> None:
         "segmentation": run_segmentation,
         "text": run_text_logs,
         "transforms_3d": transforms_rigid_3d,
+        "2d_layering": run_2d_layering,
     }
 
     parser = argparse.ArgumentParser(description="Logs rich data using the Rerun SDK.")

examples/rust/api_demo/src/main.rs

@@ -21,8 +21,8 @@ use itertools::Itertools;
 use rerun::{
     components::{
         AnnotationContext, AnnotationInfo, Box3D, ClassDescription, ClassId, ColorRGBA, DrawOrder,
-        Label, LineStrip3D, Point2D, Point3D, Quaternion, Radius, Rect2D, Rigid3, Tensor,
-        TensorDataMeaning, TextEntry, Transform, Vec3D, ViewCoordinates,
+        Label, LineStrip2D, LineStrip3D, Point2D, Point3D, Quaternion, Radius, Rect2D, Rigid3,
+        Tensor, TensorDataMeaning, TextEntry, Transform, Vec2D, Vec3D, ViewCoordinates,
     },
     coordinates::SignedAxis3,
     external::{
@@ -303,46 +303,59 @@ fn colored_tensor<F: Fn(usize, usize) -> [u8; 3]>(
 fn demo_2d_layering(rec_stream: &RecordingStream) -> anyhow::Result<()> {
     use ndarray::prelude::*;
 
+    let time = sim_time(1.0);
+
     // Add several overlapping images.
     // Large dark gray in the background
     let img = Array::<u8, _>::from_elem((512, 512, 1).f(), 64);
     MsgSender::new("2d_layering/background")
-        .with_timepoint(sim_time(1.0))
+        .with_timepoint(time.clone())
         .with_component(&[Tensor::try_from(img.as_standard_layout().view())?])?
         .with_component(&[DrawOrder(0.0)])?
         .send(rec_stream)?;
     // Smaller gradient in the middle
     let img = colored_tensor(256, 256, |x, y| [x as u8, y as u8, 0]);
-    MsgSender::new("2d_layering/middle_red")
-        .with_timepoint(sim_time(1.0))
+    MsgSender::new("2d_layering/middle_gradient")
+        .with_timepoint(time.clone())
         .with_component(&[Tensor::try_from(img.as_standard_layout().view())?])?
         .with_component(&[DrawOrder(1.0)])?
         .send(rec_stream)?;
-    // Slightly smaller red in the middle, on the same layer as the previous.
+    // Slightly smaller blue in the middle, on the same layer as the previous.
     let img = colored_tensor(192, 192, |_, _| [0, 0, 255]);
     MsgSender::new("2d_layering/middle_blue")
-        .with_timepoint(sim_time(1.0))
+        .with_timepoint(time.clone())
         .with_component(&[Tensor::try_from(img.as_standard_layout().view())?])?
         .with_component(&[DrawOrder(1.0)])?
         .send(rec_stream)?;
     // Small white on top.
     let img = Array::<u8, _>::from_elem((128, 128, 1).f(), 255);
     MsgSender::new("2d_layering/top")
-        .with_timepoint(sim_time(1.0))
+        .with_timepoint(time.clone())
         .with_component(&[Tensor::try_from(img.as_standard_layout().view())?])?
         .with_component(&[DrawOrder(2.0)])?
         .send(rec_stream)?;
 
-    // Put a rectangle in between
+    // Rectangle in between the top and the middle.
     MsgSender::new("2d_layering/rect_between_top_and_middle")
-        .with_timepoint(sim_time(2.0))
+        .with_timepoint(time.clone())
         .with_component(&[Rect2D::from_xywh(64.0, 64.0, 256.0, 256.0)])?
         .with_component(&[DrawOrder(1.5)])?
         .send(rec_stream)?;
 
+    // Lines behind the rectangle.
+    MsgSender::new("2d_layering/lines_behind_rect")
+        .with_timepoint(time.clone())
+        .with_component(&[LineStrip2D(
+            (0..20)
+                .map(|i| Vec2D([(i * 20) as f32, (i % 2 * 100 + 100) as f32]))
+                .collect(),
+        )])?
+        .with_component(&[DrawOrder(1.25)])?
+        .send(rec_stream)?;
+
     // And some points in front of the rectangle.
     MsgSender::new("2d_layering/points_between_top_and_middle")
-        .with_timepoint(sim_time(1 as _))
+        .with_timepoint(time)
         .with_component(
             &(0..256)
                 .map(|i| Point2D::new(32.0 + (i / 16) as f32 * 16.0, 64.0 + (i % 16) as f32 * 16.0))

rerun_py/rerun_sdk/rerun/components/__init__.py

@@ -13,14 +13,15 @@
     "arrow",
     "box",
     "color",
+    "draw_order",
     "experimental",
     "label",
     "point",
     "quaternion",
     "radius",
     "rect2d",
-    "scalar",
     "scalar_plot_props",
+    "scalar",
     "tensor",
     "text_entry",
     "vec",

rerun_py/rerun_sdk/rerun/components/draw_order.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import pyarrow as pa
+
+from rerun.components import REGISTERED_COMPONENT_NAMES, ComponentTypeFactory
+
+__all__ = [
+    "DrawOrderArray",
+    "DrawOrder",
+]
+
+
+class DrawOrderArray(pa.ExtensionArray):  # type: ignore[misc]
+    def splat(draw_order: float) -> DrawOrderArray:  # type: ignore[misc]
+        storage = pa.array([draw_order], type=DrawOrder.storage_type)
+        return storage  # type: ignore[no-any-return]
+
+
+DrawOrder = ComponentTypeFactory("DrawOrder", DrawOrderArray, REGISTERED_COMPONENT_NAMES["rerun.draw_order"])
+
+pa.register_extension_type(DrawOrder())

rerun_py/rerun_sdk/rerun/log/image.py

@@ -21,6 +21,7 @@ def log_image(
     entity_path: str,
     image: Tensor,
     *,
+    draw_order: Optional[float] = None,
     ext: Optional[Dict[str, Any]] = None,
     timeless: bool = False,
     recording: Optional[RecordingStream] = None,
@@ -44,6 +45,10 @@ def log_image(
         Path to the image in the space hierarchy.
     image:
         A [Tensor][rerun.log.tensor.Tensor] representing the image to log.
+    draw_order:
+        An optional floating point value that specifies the 2D drawing order.
+        Objects with higher values are drawn on top of those with lower values.
+        The default for images is -10.0.
     ext:
         Optional dictionary of extension components. See [rerun.log_extension_components][]
     timeless:
@@ -83,14 +88,15 @@ def log_image(
     if interpretable_as_image and num_non_empty_dims != len(shape):
         image = np.squeeze(image)
 
-    _log_tensor(entity_path, image, ext=ext, timeless=timeless, recording=recording)
+    _log_tensor(entity_path, image, draw_order=draw_order, ext=ext, timeless=timeless, recording=recording)
 
 
 @log_decorator
 def log_depth_image(
     entity_path: str,
     image: Tensor,
     *,
+    draw_order: Optional[float] = None,
     meter: Optional[float] = None,
     ext: Optional[Dict[str, Any]] = None,
     timeless: bool = False,
@@ -111,6 +117,10 @@ def log_depth_image(
         Path to the image in the space hierarchy.
     image:
         A [Tensor][rerun.log.tensor.Tensor] representing the depth image to log.
+    draw_order:
+        An optional floating point value that specifies the 2D drawing order.
+        Objects with higher values are drawn on top of those with lower values.
+        The default for images is -10.0.
     meter:
         How long is a meter in the given dtype?
         For instance: with uint16, perhaps meter=1000 which would mean
@@ -151,6 +161,7 @@ def log_depth_image(
         _log_tensor(
             entity_path,
             image,
+            draw_order=draw_order,
             meter=meter,
             ext=ext,
             timeless=timeless,
@@ -164,6 +175,7 @@ def log_segmentation_image(
     entity_path: str,
     image: npt.ArrayLike,
     *,
+    draw_order: Optional[float] = None,
     ext: Optional[Dict[str, Any]] = None,
     timeless: bool = False,
     recording: Optional[RecordingStream] = None,
@@ -186,6 +198,10 @@ def log_segmentation_image(
         Path to the image in the space hierarchy.
     image:
         A [Tensor][rerun.log.tensor.Tensor] representing the segmentation image to log.
+    draw_order:
+        An optional floating point value that specifies the 2D drawing order.
+        Objects with higher values are drawn on top of those with lower values.
+        The default for images is -10.0.
     ext:
         Optional dictionary of extension components. See [rerun.log_extension_components][]
     timeless:
@@ -214,6 +230,7 @@ def log_segmentation_image(
         _log_tensor(
             entity_path,
             tensor=image,
+            draw_order=draw_order,
             ext=ext,
             timeless=timeless,
             recording=recording,
@@ -225,6 +242,7 @@ def log_segmentation_image(
         _log_tensor(
             entity_path,
             tensor=image,
+            draw_order=draw_order,
             meaning=bindings.TensorDataMeaning.ClassId,
             ext=ext,
             timeless=timeless,

rerun_py/rerun_sdk/rerun/log/lines.py

@@ -6,6 +6,7 @@
 
 from rerun import bindings
 from rerun.components.color import ColorRGBAArray
+from rerun.components.draw_order import DrawOrderArray
 from rerun.components.instance import InstanceArray
 from rerun.components.linestrip import LineStrip2DArray, LineStrip3DArray
 from rerun.components.radius import RadiusArray
@@ -44,6 +45,7 @@ def log_line_strip(
     *,
     stroke_width: Optional[float] = None,
     color: Optional[Color] = None,
+    draw_order: Optional[float] = None,
     ext: Optional[Dict[str, Any]] = None,
     timeless: bool = False,
     recording: Optional[RecordingStream] = None,
@@ -71,6 +73,10 @@ def log_line_strip(
         Optional width of the line.
     color:
         Optional RGB or RGBA in sRGB gamma-space as either 0-1 floats or 0-255 integers, with separate alpha.
+    draw_order:
+        An optional floating point value that specifies the 2D drawing order.
+        Objects with higher values are drawn on top of those with lower values.
+        The default for lines is 20.0.
     ext:
         Optional dictionary of extension components. See [rerun.log_extension_components][]
     timeless:
@@ -106,6 +112,9 @@ def log_line_strip(
         radii = _normalize_radii([stroke_width / 2])
         instanced["rerun.radius"] = RadiusArray.from_numpy(radii)
 
+    if draw_order is not None:
+        instanced["rerun.draw_order"] = DrawOrderArray.splat(draw_order)
+
     if ext:
         _add_extension_components(instanced, splats, ext, None)
 
@@ -125,6 +134,7 @@ def log_line_segments(
     *,
     stroke_width: Optional[float] = None,
     color: Optional[Color] = None,
+    draw_order: Optional[float] = None,
     ext: Optional[Dict[str, Any]] = None,
     timeless: bool = False,
     recording: Optional[RecordingStream] = None,
@@ -151,6 +161,10 @@ def log_line_segments(
         Optional width of the line.
     color:
         Optional RGB or RGBA in sRGB gamma-space as either 0-1 floats or 0-255 integers, with separate alpha.
+    draw_order:
+        An optional floating point value that specifies the 2D drawing order.
+        Objects with higher values are drawn on top of those with lower values.
+        The default for lines is 20.0.
     ext:
         Optional dictionary of extension components. See [rerun.log_extension_components][]
     timeless:
@@ -198,6 +212,9 @@ def log_line_segments(
         radii = _normalize_radii([stroke_width / 2])
         splats["rerun.radius"] = RadiusArray.from_numpy(radii)
 
+    if draw_order is not None:
+        instanced["rerun.draw_order"] = DrawOrderArray.splat(draw_order)
+
     if ext:
         _add_extension_components(instanced, splats, ext, None)