From d8f6b06355bdcb69aa87536d85768840010db3b2 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 8 Feb 2024 20:48:54 +0100 Subject: [PATCH 1/7] New 2d heuristic -- split any bucket with more than 1 image --- .../src/space_view_2d.rs | 167 ++++++++++++++---- 1 file changed, 136 insertions(+), 31 deletions(-) diff --git a/crates/re_space_view_spatial/src/space_view_2d.rs b/crates/re_space_view_spatial/src/space_view_2d.rs index 86e35d1fdcd0..bf6caa261d9d 100644 --- a/crates/re_space_view_spatial/src/space_view_2d.rs +++ b/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::::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 = + bucket_entities.iter().cloned().collect(); + + add_recommended_space_views_for_bucket( + ctx, + &bucket_entities, + &mut recommended_space_views, + ); + } + + // 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 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 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,100 @@ impl SpaceViewClass for SpatialSpaceView2D { } } +fn count_non_nested_entities_with_component( + entity_bucket: &IntSet, + subtree: &EntityTree, + component_name: &ComponentName, +) -> usize { + if entity_bucket.contains(&subtree.path) { + if subtree.entity.components.contains_key(component_name) { + 1 + } else { + 0 + } + } else if !entity_bucket + .iter() + .any(|e| e.is_descendant_of(&subtree.path)) + { + 0 + } else { + subtree + .children + .values() + .map(|child| { + // TODO(jleibs): Early terminate if know the subtree has nothing from the bucket + count_non_nested_entities_with_component(entity_bucket, child, component_name) + }) + .sum() + } +} + +fn add_recommended_space_views_for_bucket( + ctx: &ViewerContext<'_>, + entity_bucket: &IntSet, + recommended: &mut Vec, +) { + // 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<'_>, From 38ae1b88e22d363060004de0c3e84f3cf20b9cf5 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 8 Feb 2024 21:07:25 +0100 Subject: [PATCH 2/7] Add release checklist for some simple 2d heuristics --- .../release_checklist/check_2d_heuristics.py | 80 +++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 tests/python/release_checklist/check_2d_heuristics.py diff --git a/tests/python/release_checklist/check_2d_heuristics.py b/tests/python/release_checklist/check_2d_heuristics.py new file mode 100644 index 000000000000..9b25b5481073 --- /dev/null +++ b/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) From 736676eb639095a71aaa3cdb8010a17348e93817 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 8 Feb 2024 21:19:01 +0100 Subject: [PATCH 3/7] lint --- crates/re_space_view_spatial/src/space_view_2d.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/crates/re_space_view_spatial/src/space_view_2d.rs b/crates/re_space_view_spatial/src/space_view_2d.rs index bf6caa261d9d..41be9629ae01 100644 --- a/crates/re_space_view_spatial/src/space_view_2d.rs +++ b/crates/re_space_view_spatial/src/space_view_2d.rs @@ -303,11 +303,8 @@ fn count_non_nested_entities_with_component( component_name: &ComponentName, ) -> usize { if entity_bucket.contains(&subtree.path) { - if subtree.entity.components.contains_key(component_name) { - 1 - } else { - 0 - } + // 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)) @@ -318,7 +315,6 @@ fn count_non_nested_entities_with_component( .children .values() .map(|child| { - // TODO(jleibs): Early terminate if know the subtree has nothing from the bucket count_non_nested_entities_with_component(entity_bucket, child, component_name) }) .sum() From ef96839cdd61fc30dbd9dfd946237361f498f96a Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 8 Feb 2024 22:28:16 +0100 Subject: [PATCH 4/7] lint --- crates/re_space_view_spatial/src/space_view_2d.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/re_space_view_spatial/src/space_view_2d.rs b/crates/re_space_view_spatial/src/space_view_2d.rs index 41be9629ae01..696266d450ec 100644 --- a/crates/re_space_view_spatial/src/space_view_2d.rs +++ b/crates/re_space_view_spatial/src/space_view_2d.rs @@ -346,7 +346,7 @@ fn add_recommended_space_views_for_bucket( // 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.") + re_log::warn_once!("Ancestor of entity not found in entity tree."); } return; }; From 4277f43a635d3664397dbbec5a2c5e3d447fc3e7 Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt Date: Fri, 9 Feb 2024 09:15:37 +0100 Subject: [PATCH 5/7] Bucket into sets from the start --- .../re_space_view_spatial/src/space_view_2d.rs | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/crates/re_space_view_spatial/src/space_view_2d.rs b/crates/re_space_view_spatial/src/space_view_2d.rs index 696266d450ec..8b7497b14863 100644 --- a/crates/re_space_view_spatial/src/space_view_2d.rs +++ b/crates/re_space_view_spatial/src/space_view_2d.rs @@ -195,13 +195,8 @@ impl SpaceViewClass for SpatialSpaceView2D { let mut recommended_space_views = Vec::::new(); for bucket_entities in - bucket_images_in_subspace(ctx, subspace, image_entities).values() + bucket_images_in_subspace(ctx, subspace, image_entities) { - // Collect the bucket as a set - // TODO(jleibs): might as well do this when bucketing in the first place - let bucket_entities: IntSet = - bucket_entities.iter().cloned().collect(); - add_recommended_space_views_for_bucket( ctx, &bucket_entities, @@ -392,7 +387,7 @@ fn bucket_images_in_subspace( ctx: &ViewerContext<'_>, subspace: &SubSpace, image_entities: &ApplicableEntities, -) -> HashMap<(u64, u64), Vec> { +) -> Vec> { re_tracing::profile_function!(); let store = ctx.entity_db.store(); @@ -407,11 +402,11 @@ fn bucket_images_in_subspace( // Very common case, early out before we get into the more expensive query code. return image_entities .into_iter() - .map(|e| ((0, 0), vec![e.clone()])) + .map(|e| std::iter::once(e.clone()).collect()) .collect(); } - let mut images_by_bucket = HashMap::<(u64, u64), Vec>::default(); + let mut images_by_bucket = HashMap::<(u64, u64), IntSet>::default(); for image_entity in image_entities { // TODO(andreas): We really don't want to do a latest at query here since this means the heuristic can have different results depending on the // current query, but for this we'd have to store the max-size over time somewhere using another store subscriber (?). @@ -424,11 +419,11 @@ fn bucket_images_in_subspace( images_by_bucket .entry((height, width)) .or_default() - .push(image_entity.clone()); + .insert(image_entity.clone()); } } } } - images_by_bucket + images_by_bucket.drain().map(|(_, v)| v).collect() } From 0618d817e1e7c1bf2c45980919aed61d430a4602 Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt Date: Fri, 9 Feb 2024 09:18:42 +0100 Subject: [PATCH 6/7] Mention that the function is for images= --- crates/re_space_view_spatial/src/space_view_2d.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/re_space_view_spatial/src/space_view_2d.rs b/crates/re_space_view_spatial/src/space_view_2d.rs index 8b7497b14863..061c1f91a0cf 100644 --- a/crates/re_space_view_spatial/src/space_view_2d.rs +++ b/crates/re_space_view_spatial/src/space_view_2d.rs @@ -197,7 +197,7 @@ impl SpaceViewClass for SpatialSpaceView2D { for bucket_entities in bucket_images_in_subspace(ctx, subspace, image_entities) { - add_recommended_space_views_for_bucket( + add_recommended_space_views_for_image_bucket( ctx, &bucket_entities, &mut recommended_space_views, @@ -316,7 +316,8 @@ fn count_non_nested_entities_with_component( } } -fn add_recommended_space_views_for_bucket( +/// Given a bucket of image entities. +fn add_recommended_space_views_for_image_bucket( ctx: &ViewerContext<'_>, entity_bucket: &IntSet, recommended: &mut Vec, @@ -377,7 +378,7 @@ fn add_recommended_space_views_for_bucket( .collect(); if !sub_bucket.is_empty() { - add_recommended_space_views_for_bucket(ctx, &sub_bucket, recommended); + add_recommended_space_views_for_image_bucket(ctx, &sub_bucket, recommended); } } } From b6b855478dda18f83afde6724bdeb915936a2751 Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt Date: Fri, 9 Feb 2024 09:24:07 +0100 Subject: [PATCH 7/7] Try to explain --- crates/re_space_view_spatial/src/space_view_2d.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/crates/re_space_view_spatial/src/space_view_2d.rs b/crates/re_space_view_spatial/src/space_view_2d.rs index 061c1f91a0cf..0f86892954a8 100644 --- a/crates/re_space_view_spatial/src/space_view_2d.rs +++ b/crates/re_space_view_spatial/src/space_view_2d.rs @@ -292,6 +292,12 @@ impl SpaceViewClass for SpatialSpaceView2D { } } +/// Count how many occurrences of the given component are in +/// the intersection of the subtree and the entity bucket, +/// but stops recursion if there is a hit. +/// +/// * So if the subtree root is in the bucket: return 1 if it has the component, 0 otherwise. +/// * If not: recurse on children and sum the results. fn count_non_nested_entities_with_component( entity_bucket: &IntSet, subtree: &EntityTree, @@ -304,7 +310,7 @@ fn count_non_nested_entities_with_component( .iter() .any(|e| e.is_descendant_of(&subtree.path)) { - 0 + 0 // early-out optimization } else { subtree .children