New 2d heuristic -- split any bucket with more than 1 image

crates/re_space_view_spatial/src/space_view_2d.rs

@@ -2,10 +2,14 @@ use ahash::HashMap;
 use itertools::Itertools;
 use nohash_hasher::IntSet;
 
-use re_entity_db::EntityProperties;
+use re_entity_db::{EntityProperties, EntityTree};
 use re_log_types::{EntityPath, EntityPathFilter};
 use re_tracing::profile_scope;
-use re_types::components::TensorData;
+use re_types::{
+    archetypes::{DepthImage, Image},
+    components::TensorData,
+    Archetype, ComponentName,
+};
 use re_viewer_context::{
     ApplicableEntities, IdentifiedViewSystem as _, PerSystemEntities, RecommendedSpaceView,
     SpaceViewClass, SpaceViewClassRegistryError, SpaceViewId, SpaceViewSpawnHeuristics,
@@ -176,8 +180,8 @@ impl SpaceViewClass for SpatialSpaceView2D {
         // Spawn a space view at each subspace that has any potential 2D content.
         // Note that visualizability filtering is all about being in the right subspace,
         // so we don't need to call the visualizers' filter functions here.
-        let mut heuristics = SpatialTopology::access(ctx.entity_db.store_id(), |topo| {
-            SpaceViewSpawnHeuristics {
+        let mut heuristics =
+            SpatialTopology::access(ctx.entity_db.store_id(), |topo| SpaceViewSpawnHeuristics {
                 recommended_space_views: topo
                     .iter_subspaces()
                     .flat_map(|subspace| {
@@ -188,41 +192,48 @@ impl SpaceViewClass for SpatialSpaceView2D {
                             return Vec::new();
                         }
 
-                        let images_by_bucket =
-                            bucket_images_in_subspace(ctx, subspace, image_entities);
+                        let mut recommended_space_views = Vec::<RecommendedSpaceView>::new();
+
+                        for bucket_entities in
+                            bucket_images_in_subspace(ctx, subspace, image_entities).values()
+                        {
+                            // Collect the bucket as a set
+                            // TODO(jleibs): might as well do this when bucketing in the first place
+                            let bucket_entities: IntSet<EntityPath> =
+                                bucket_entities.iter().cloned().collect();
+
+                            add_recommended_space_views_for_bucket(
+                                ctx,
+                                &bucket_entities,
+                                &mut recommended_space_views,
+                            );
+                        }
 
-                        if images_by_bucket.len() <= 1 {
-                            // If there's no or only a single image bucket, use the whole subspace to capture all the non-image entities!
-                            vec![RecommendedSpaceView {
+                        // If we only recommended 1 space view from the bucketing, we're better off using the
+                        // root of the subspace, below. If there were multiple subspaces, keep them, even if
+                        // they may be redundant with the root space.
+                        if recommended_space_views.len() == 1 {
+                            recommended_space_views.clear();
+                        }
+
+                        // If this is explicitly a 2D subspace (such as from a pinhole), or there were no
+                        // other image-bucketed recommendations, create a space at the root of the subspace.
+                        if subspace.dimensionality == SubSpaceDimensionality::TwoD
+                            || recommended_space_views.is_empty()
+                        {
+                            recommended_space_views.push(RecommendedSpaceView {
                                 root: subspace.origin.clone(),
                                 query_filter: EntityPathFilter::subtree_entity_filter(
                                     &subspace.origin,
                                 ),
-                            }]
-                        } else {
-                            #[allow(clippy::iter_kv_map)] // Not doing `values()` saves a path copy!
-                            images_by_bucket
-                                .into_iter()
-                                .map(|(_, entity_bucket)| {
-                                    // Pick a shared parent as origin, mostly because it looks nicer in the ui.
-                                    let root = EntityPath::common_ancestor_of(entity_bucket.iter());
-
-                                    let mut query_filter = EntityPathFilter::default();
-                                    for image in &entity_bucket {
-                                        // This might lead to overlapping subtrees and break the same image size bucketing again.
-                                        // We just take that risk, the heuristic doesn't need to be perfect.
-                                        query_filter.add_subtree(image.clone());
-                                    }
-
-                                    RecommendedSpaceView { root, query_filter }
-                                })
-                                .collect()
+                            });
                         }
+
+                        recommended_space_views
                     })
                     .collect(),
-            }
-        })
-        .unwrap_or_default();
+            })
+            .unwrap_or_default();
 
         // Find all entities that are not yet covered by the recommended space views and create a recommended
         // space-view for each one at that specific entity path.
@@ -286,6 +297,96 @@ impl SpaceViewClass for SpatialSpaceView2D {
     }
 }
 
+fn count_non_nested_entities_with_component(
+    entity_bucket: &IntSet<EntityPath>,
+    subtree: &EntityTree,
+    component_name: &ComponentName,
+) -> usize {
+    if entity_bucket.contains(&subtree.path) {
+        // bool true -> 1
+        subtree.entity.components.contains_key(component_name) as usize
+    } else if !entity_bucket
+        .iter()
+        .any(|e| e.is_descendant_of(&subtree.path))
+    {
+        0
+    } else {
+        subtree
+            .children
+            .values()
+            .map(|child| {
+                count_non_nested_entities_with_component(entity_bucket, child, component_name)
+            })
+            .sum()
+    }
+}
+
+fn add_recommended_space_views_for_bucket(
+    ctx: &ViewerContext<'_>,
+    entity_bucket: &IntSet<EntityPath>,
+    recommended: &mut Vec<RecommendedSpaceView>,
+) {
+    // TODO(jleibs): Converting entity_bucket to a Trie would probably make some of this easier.
+    let tree = ctx.entity_db.tree();
+
+    // Find the common ancestor of the bucket
+    let root = EntityPath::common_ancestor_of(entity_bucket.iter());
+
+    // If the root of this bucket contains an image itself, this means the rest of the content
+    // is nested under some kind of 2d-visualizable thing. We expect the user meant to create
+    // a layered 2d space.
+    if entity_bucket.contains(&root) {
+        recommended.push(RecommendedSpaceView {
+            root: root.clone(),
+            query_filter: EntityPathFilter::subtree_entity_filter(&root),
+        });
+        return;
+    }
+
+    // Alternatively we want to split this bucket into a group for each child-space.
+    let Some(subtree) = tree.subtree(&root) else {
+        if cfg!(debug_assertions) {
+            re_log::warn_once!("Ancestor of entity not found in entity tree.");
+        }
+        return;
+    };
+
+    let image_count = count_non_nested_entities_with_component(
+        entity_bucket,
+        subtree,
+        &Image::indicator().name(),
+    );
+
+    let depth_count = count_non_nested_entities_with_component(
+        entity_bucket,
+        subtree,
+        &DepthImage::indicator().name(),
+    );
+
+    // If there's no more than 1 image and 1 depth image at any of the top-level of the sub-buckets, we can still
+    // recommend the root.
+    if image_count <= 1 && depth_count <= 1 {
+        recommended.push(RecommendedSpaceView {
+            root: root.clone(),
+            query_filter: EntityPathFilter::subtree_entity_filter(&root),
+        });
+        return;
+    }
+
+    // Otherwise, split the space and recurse
+    for child in subtree.children.values() {
+        let sub_bucket: IntSet<_> = entity_bucket
+            .iter()
+            .filter(|e| e.starts_with(&child.path))
+            .cloned()
+            .collect();
+
+        if !sub_bucket.is_empty() {
+            add_recommended_space_views_for_bucket(ctx, &sub_bucket, recommended);
+        }
+    }
+}
+
 /// Groups all images in the subspace by size and draw order.
 fn bucket_images_in_subspace(
     ctx: &ViewerContext<'_>,

tests/python/release_checklist/check_2d_heuristics.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import os
+from argparse import Namespace
+from uuid import uuid4
+
+import numpy as np
+import rerun as rr
+
+README = """
+# 2D Heuristics
+
+This checks whether the heuristics do the right thing with images.
+
+Reset the blueprint to make sure you are viewing new heuristics and not a cached blueprint.
+
+### Action
+You should see 4 space-views. Depending on timing you may end up with a 5th space-view at the root.
+This should go away when you reset.
+
+The four remaining space-views should be:
+ - `image1` with a red square
+ - `image2` with a green square
+ - `image3` with a blue square and overlapping green square (rendered teal)
+ - `segmented` with a red square and overlapping green square (rendered yellow)
+"""
+
+
+def log_image(path: str, height: int, width: int, color: tuple[int, int, int]) -> None:
+    image = np.zeros((height, width, 3), dtype=np.uint8)
+    image[:, :, :] = color
+    rr.log(path, rr.Image(image))
+
+
+def log_image_nested(path: str, height: int, width: int, color: tuple[int, int, int]) -> None:
+    image = np.zeros((height, width, 3), dtype=np.uint8)
+    image[int(height / 4) : int(height - height / 4), int(width / 4) : int(width - width / 4), :] = color
+    rr.log(path, rr.Image(image))
+
+
+def log_annotation_context() -> None:
+    rr.log("/", rr.AnnotationContext([(1, "red", (255, 0, 0)), (2, "green", (0, 255, 0))]), timeless=True)
+
+
+def log_segmentation(path: str, height: int, width: int, class_id: int) -> None:
+    image = np.zeros((height, width, 1), dtype=np.uint8)
+    image[int(height / 4) : int(height - height / 4), int(width / 4) : int(width - width / 4), 0] = class_id
+    rr.log(path, rr.SegmentationImage(image))
+
+
+def log_readme() -> None:
+    rr.log("readme", rr.TextDocument(README, media_type=rr.MediaType.MARKDOWN), timeless=True)
+
+
+def log_images() -> None:
+    log_annotation_context()
+    log_image("image1", 20, 30, (255, 0, 0))
+    log_image("image2", 20, 30, (0, 255, 0))
+    log_image("image3", 20, 30, (0, 0, 255))
+    log_image_nested("image3/nested", 20, 30, (0, 255, 0))
+    log_image("segmented/image4", 20, 30, (255, 0, 0))
+    log_segmentation("segmented/seg", 20, 30, 2)
+
+
+def run(args: Namespace) -> None:
+    # TODO(cmc): I have no idea why this works without specifying a `recording_id`, but
+    # I'm not gonna rely on it anyway.
+    rr.script_setup(args, f"{os.path.basename(__file__)}", recording_id=uuid4())
+
+    log_readme()
+    log_images()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Interactive release checklist")
+    rr.script_add_args(parser)
+    args = parser.parse_args()
+    run(args)