Introduce (2D) Draw order component (#2056)

* create depth offset from draw order and patch it through

* transparent layers now know about depth offset - only images on the same layer will be made transparent

* add some points to layering demo

* comments and renames on draw order constants

* doc & sample details

* images with a shared root & an explicit draw order all go to the same space view now

* formatting

crates/re_log_types/src/component_types/draw_order.rs

@@ -0,0 +1,87 @@
+use arrow2_convert::{ArrowDeserialize, ArrowField, ArrowSerialize};
+
+use crate::Component;
+
+/// Draw order used for the display order of 2D elements.
+///
+/// Higher values are drawn on top of lower values.
+/// An entity can have only a single draw order component.
+/// Within an entity draw order is governed by the order of the components.
+///
+/// Draw order for entities with the same draw order is generally undefined.
+///
+/// ```
+/// use re_log_types::component_types::DrawOrder;
+/// use arrow2_convert::field::ArrowField;
+/// use arrow2::datatypes::{DataType, Field};
+///
+/// assert_eq!(DrawOrder::data_type(), DataType::Float32);
+/// ```
+#[derive(Debug, Clone, ArrowField, ArrowSerialize, ArrowDeserialize)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+#[arrow_field(transparent)]
+pub struct DrawOrder(pub f32);
+
+impl DrawOrder {
+    /// Draw order used for images if no draw order was specified.
+    pub const DEFAULT_IMAGE: DrawOrder = DrawOrder(-10.0);
+
+    /// Draw order used for 2D boxes if no draw order was specified.
+    pub const DEFAULT_BOX2D: DrawOrder = DrawOrder(10.0);
+
+    /// Draw order used for 2D lines if no draw order was specified.
+    pub const DEFAULT_LINES2D: DrawOrder = DrawOrder(20.0);
+
+    /// Draw order used for 2D points if no draw order was specified.
+    pub const DEFAULT_POINTS2D: DrawOrder = DrawOrder(30.0);
+}
+
+impl Component for DrawOrder {
+    #[inline]
+    fn name() -> crate::ComponentName {
+        "rerun.draw_order".into()
+    }
+}
+
+impl std::cmp::PartialEq for DrawOrder {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        self.0.is_nan() && other.0.is_nan() || self.0 == other.0
+    }
+}
+
+impl std::cmp::Eq for DrawOrder {}
+
+impl std::cmp::PartialOrd for DrawOrder {
+    #[inline]
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        if other == self {
+            Some(std::cmp::Ordering::Equal)
+        } else if other.0.is_nan() || self.0 < other.0 {
+            Some(std::cmp::Ordering::Less)
+        } else {
+            Some(std::cmp::Ordering::Greater)
+        }
+    }
+}
+
+impl std::cmp::Ord for DrawOrder {
+    #[inline]
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.partial_cmp(other).unwrap()
+    }
+}
+
+impl From<f32> for DrawOrder {
+    #[inline]
+    fn from(value: f32) -> Self {
+        Self(value)
+    }
+}
+
+impl From<DrawOrder> for f32 {
+    #[inline]
+    fn from(value: DrawOrder) -> Self {
+        value.0
+    }
+}

crates/re_log_types/src/component_types/mod.rs

@@ -23,6 +23,7 @@ mod class_id;
 mod color;
 pub mod context;
 pub mod coordinates;
+mod draw_order;
 mod instance_key;
 mod keypoint_id;
 mod label;
@@ -46,6 +47,7 @@ pub use class_id::ClassId;
 pub use color::ColorRGBA;
 pub use context::{AnnotationContext, AnnotationInfo, ClassDescription};
 pub use coordinates::ViewCoordinates;
+pub use draw_order::DrawOrder;
 pub use instance_key::InstanceKey;
 pub use keypoint_id::KeypointId;
 pub use label::Label;

crates/re_log_types/src/lib.rs

@@ -36,6 +36,7 @@ pub use self::component_types::coordinates;
 pub use self::component_types::AnnotationContext;
 pub use self::component_types::Arrow3D;
 pub use self::component_types::DecodedTensor;
+pub use self::component_types::DrawOrder;
 pub use self::component_types::{EncodedMesh3D, Mesh3D, MeshFormat, MeshId, RawMesh3D};
 pub use self::component_types::{Tensor, ViewCoordinates};
 pub use self::data::*;

crates/re_sdk/src/lib.rs

@@ -69,9 +69,9 @@ pub mod time {
 pub mod components {
     pub use re_log_types::component_types::{
         AnnotationContext, AnnotationInfo, Arrow3D, Box3D, ClassDescription, ClassId, ColorRGBA,
-        EncodedMesh3D, InstanceKey, KeypointId, Label, LineStrip2D, LineStrip3D, Mat3x3, Mesh3D,
-        MeshFormat, MeshId, Pinhole, Point2D, Point3D, Quaternion, Radius, RawMesh3D, Rect2D,
-        Rigid3, Scalar, ScalarPlotProps, Size3D, Tensor, TensorData, TensorDataMeaning,
+        DrawOrder, EncodedMesh3D, InstanceKey, KeypointId, Label, LineStrip2D, LineStrip3D, Mat3x3,
+        Mesh3D, MeshFormat, MeshId, Pinhole, Point2D, Point3D, Quaternion, Radius, RawMesh3D,
+        Rect2D, Rigid3, Scalar, ScalarPlotProps, Size3D, Tensor, TensorData, TensorDataMeaning,
         TensorDimension, TensorId, TextEntry, Transform, Vec2D, Vec3D, Vec4D, ViewCoordinates,
     };
 }

crates/re_viewer/src/ui/space_view_heuristics.rs

@@ -171,7 +171,13 @@ fn default_created_space_views_from_candidates(
 
         // Spatial views with images get extra treatment as well.
         if candidate.category == ViewCategory::Spatial {
-            let mut images_by_size: HashMap<(u64, u64), Vec<EntityPath>> = HashMap::default();
+            #[derive(Hash, PartialEq, Eq)]
+            enum ImageBucketing {
+                BySize((u64, u64)),
+                ExplicitDrawOrder,
+            }
+
+            let mut images_by_bucket: HashMap<ImageBucketing, Vec<EntityPath>> = HashMap::default();
 
             // For this we're only interested in the direct children.
             for entity_path in &candidate.data_blueprint.root_group().entities {
@@ -184,24 +190,41 @@ fn default_created_space_views_from_candidates(
                     for tensor in entity_view.iter_primary_flattened() {
                         if tensor.is_shaped_like_an_image() {
                             debug_assert!(matches!(tensor.shape.len(), 2 | 3));
-                            let dim = (tensor.shape[0].size, tensor.shape[1].size);
-                            images_by_size
-                                .entry(dim)
-                                .or_default()
-                                .push(entity_path.clone());
+
+                            if query_latest_single::<re_log_types::DrawOrder>(
+                                entity_db,
+                                entity_path,
+                                &query,
+                            )
+                            .is_some()
+                            {
+                                // Put everything in the same bucket if it has a draw order.
+                                images_by_bucket
+                                    .entry(ImageBucketing::ExplicitDrawOrder)
+                                    .or_default()
+                                    .push(entity_path.clone());
+                            } else {
+                                // Otherwise, distinguish buckets by image size.
+                                let dim = (tensor.shape[0].size, tensor.shape[1].size);
+                                images_by_bucket
+                                    .entry(ImageBucketing::BySize(dim))
+                                    .or_default()
+                                    .push(entity_path.clone());
+                            }
                         }
                     }
                 }
             }
 
-            // If all images are the same size, proceed with the candidate as is. Otherwise...
-            if images_by_size.len() > 1 {
-                // ...stack images of the same size, but no others.
-                for dim in images_by_size.keys() {
-                    // Ignore every image that has a different size.
-                    let images_of_different_size = images_by_size
+            if images_by_bucket.len() > 1 {
+                // If all images end up in the same bucket, proceed as normal. Otherwise stack images as instructed.
+                for bucket in images_by_bucket.keys() {
+                    // Ignore every image from antoher bucket. Keep all other entities.
+                    let images_of_different_size = images_by_bucket
                         .iter()
-                        .filter_map(|(other_dim, images)| (dim != other_dim).then_some(images))
+                        .filter_map(|(other_bucket, images)| {
+                            (bucket != other_bucket).then_some(images)
+                        })
                         .flatten()
                         .cloned()
                         .collect::<IntSet<_>>();

crates/re_viewer/src/ui/view_spatial/scene/mod.rs

@@ -1,14 +1,16 @@
-use std::sync::Arc;
+use std::{collections::BTreeMap, sync::Arc};
 
 use ahash::HashMap;
 
-use re_data_store::{EntityPath, InstancePathHash};
+use nohash_hasher::IntMap;
+use re_data_store::{query_latest_single, EntityPath, InstancePathHash};
 use re_log_types::{
     component_types::{ClassId, InstanceKey, KeypointId},
-    DecodedTensor,
+    DecodedTensor, DrawOrder, EntityPathHash,
 };
 use re_renderer::{renderer::TexturedRect, Color32, OutlineMaskPreference, Size};
 use re_viewer_context::{auto_color, AnnotationMap, Annotations, SceneQuery, ViewerContext};
+use smallvec::SmallVec;
 
 use crate::misc::{mesh_loader::LoadedMesh, SpaceViewHighlights, TransformCache};
 
@@ -91,6 +93,21 @@ pub struct SceneSpatial {
 
 pub type Keypoints = HashMap<(ClassId, i64), HashMap<KeypointId, glam::Vec3>>;
 
+#[derive(Default)]
+pub struct EntityDepthOffsets {
+    pub per_entity: IntMap<EntityPathHash, re_renderer::DepthOffset>,
+    pub box2d: re_renderer::DepthOffset,
+    pub lines2d: re_renderer::DepthOffset,
+    pub image: re_renderer::DepthOffset,
+    pub points: re_renderer::DepthOffset,
+}
+
+impl EntityDepthOffsets {
+    pub fn get(&self, ent_path: &EntityPath) -> Option<re_renderer::DepthOffset> {
+        self.per_entity.get(&ent_path.hash()).cloned()
+    }
+}
+
 impl SceneSpatial {
     pub fn new(re_ctx: &mut re_renderer::RenderContext) -> Self {
         Self {
@@ -103,6 +120,89 @@ impl SceneSpatial {
         }
     }
 
+    fn determine_depth_offsets(
+        ctx: &mut ViewerContext<'_>,
+        query: &SceneQuery<'_>,
+    ) -> EntityDepthOffsets {
+        crate::profile_function!();
+
+        enum DrawOrderTarget {
+            Entity(EntityPathHash),
+            DefaultBox2D,
+            DefaultLines2D,
+            DefaultImage,
+            DefaultPoints,
+        }
+
+        let mut entities_per_draw_order = BTreeMap::<DrawOrder, SmallVec<[_; 4]>>::new();
+        for (ent_path, _) in query.iter_entities() {
+            if let Some(draw_order) = query_latest_single::<DrawOrder>(
+                &ctx.log_db.entity_db,
+                ent_path,
+                &ctx.rec_cfg.time_ctrl.current_query(),
+            ) {
+                entities_per_draw_order
+                    .entry(draw_order)
+                    .or_default()
+                    .push(DrawOrderTarget::Entity(ent_path.hash()));
+            }
+        }
+
+        // Push in default draw orders. All of them using the none hash.
+        entities_per_draw_order
+            .entry(DrawOrder::DEFAULT_BOX2D)
+            .or_default()
+            .push(DrawOrderTarget::DefaultBox2D);
+        entities_per_draw_order
+            .entry(DrawOrder::DEFAULT_IMAGE)
+            .or_default()
+            .push(DrawOrderTarget::DefaultImage);
+        entities_per_draw_order
+            .entry(DrawOrder::DEFAULT_LINES2D)
+            .or_default()
+            .push(DrawOrderTarget::DefaultLines2D);
+        entities_per_draw_order
+            .entry(DrawOrder::DEFAULT_POINTS2D)
+            .or_default()
+            .push(DrawOrderTarget::DefaultPoints);
+
+        // Determine re_renderer draw order from this.
+        // We want to be as tightly around 0 as possible.
+        let mut offsets = EntityDepthOffsets::default();
+        let mut draw_order = -((entities_per_draw_order.len() / 2) as re_renderer::DepthOffset);
+        offsets.per_entity = entities_per_draw_order
+            .into_values()
+            .flat_map(move |targets| {
+                let pairs = targets
+                    .into_iter()
+                    .filter_map(|target| match target {
+                        DrawOrderTarget::Entity(entity) => Some((entity, draw_order)),
+                        DrawOrderTarget::DefaultBox2D => {
+                            offsets.box2d = draw_order;
+                            None
+                        }
+                        DrawOrderTarget::DefaultLines2D => {
+                            offsets.lines2d = draw_order;
+                            None
+                        }
+                        DrawOrderTarget::DefaultImage => {
+                            offsets.image = draw_order;
+                            None
+                        }
+                        DrawOrderTarget::DefaultPoints => {
+                            offsets.points = draw_order;
+                            None
+                        }
+                    })
+                    .collect::<SmallVec<[_; 4]>>();
+                draw_order += 1;
+                pairs
+            })
+            .collect();
+
+        offsets
+    }
+
     /// Loads all 3D objects into the scene according to the given query.
     pub(crate) fn load(
         &mut self,
@@ -133,8 +233,10 @@ impl SceneSpatial {
             &scene_part::CamerasPart,
         ];
 
+        let depth_offsets = Self::determine_depth_offsets(ctx, query);
+
         for part in parts {
-            part.load(self, ctx, query, transforms, highlights);
+            part.load(self, ctx, query, transforms, highlights, &depth_offsets);
         }
 
         self.primitives.any_outlines = highlights.any_outlines();

crates/re_viewer/src/ui/view_spatial/scene/scene_part/arrows3d.rs

@@ -9,7 +9,7 @@ use re_viewer_context::{DefaultColor, SceneQuery, ViewerContext};
 
 use crate::{
     misc::{SpaceViewHighlights, TransformCache},
-    ui::view_spatial::SceneSpatial,
+    ui::view_spatial::{scene::EntityDepthOffsets, SceneSpatial},
 };
 
 use super::{instance_key_to_picking_id, ScenePart};
@@ -94,6 +94,7 @@ impl ScenePart for Arrows3DPart {
         query: &SceneQuery<'_>,
         transforms: &TransformCache,
         highlights: &SpaceViewHighlights,
+        _depth_offsets: &EntityDepthOffsets,
     ) {
         crate::profile_scope!("Points2DPart");
 

crates/re_viewer/src/ui/view_spatial/scene/scene_part/boxes2d.rs

@@ -10,7 +10,8 @@ use re_viewer_context::{DefaultColor, SceneQuery, ViewerContext};
 use crate::{
     misc::{SpaceViewHighlights, TransformCache},
     ui::view_spatial::{
-        scene::scene_part::instance_path_hash_for_picking, SceneSpatial, UiLabel, UiLabelTarget,
+        scene::{scene_part::instance_path_hash_for_picking, EntityDepthOffsets},
+        SceneSpatial, UiLabel, UiLabelTarget,
     },
 };
 
@@ -25,6 +26,7 @@ impl Boxes2DPart {
         ent_path: &EntityPath,
         world_from_obj: glam::Affine3A,
         highlights: &SpaceViewHighlights,
+        depth_offset: re_renderer::DepthOffset,
     ) -> Result<(), QueryError> {
         scene.num_logged_2d_objects += 1;
 
@@ -37,6 +39,7 @@ impl Boxes2DPart {
             .primitives
             .line_strips
             .batch("2d boxes")
+            .depth_offset(depth_offset)
             .world_from_obj(world_from_obj)
             .outline_mask_ids(entity_highlight.overall)
             .picking_object_id(re_renderer::PickingLayerObjectId(ent_path.hash64()));
@@ -105,6 +108,7 @@ impl ScenePart for Boxes2DPart {
         query: &SceneQuery<'_>,
         transforms: &TransformCache,
         highlights: &SpaceViewHighlights,
+        depth_offsets: &EntityDepthOffsets,
     ) {
         crate::profile_scope!("Boxes2DPart");
 
@@ -136,6 +140,7 @@ impl ScenePart for Boxes2DPart {
                         ent_path,
                         world_from_obj,
                         highlights,
+                        depth_offsets.get(ent_path).unwrap_or(depth_offsets.box2d),
                     )?;
                 }
                 Ok(())