From a9a466fb6d6c0d8dfe70da4089aa511502d9bba9 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Mon, 12 Feb 2024 11:48:32 +0100 Subject: [PATCH] Refactor the 2D heuristic again (#5166) ### What - Builds on top of https://github.com/rerun-io/rerun/pull/5164 Adds a store subscriber for tracking the dimensions of image entities. Gets rid of the previous bucketing logic and replaces it with a new implementation that starts at the top of the subspace and incrementally splits when it finds conflicting image entities within the space. ### Checklist * [x] I have read and agree to [Contributor Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and the [Code of Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md) * [x] I've included a screenshot or gif (if applicable) * [x] I have tested the web demo (if applicable): * Using newly built examples: [app.rerun.io](https://app.rerun.io/pr/5166/index.html) * Using examples from latest `main` build: [app.rerun.io](https://app.rerun.io/pr/5166/index.html?manifest_url=https://app.rerun.io/version/main/examples_manifest.json) * Using full set of examples from `nightly` build: [app.rerun.io](https://app.rerun.io/pr/5166/index.html?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json) * [x] The PR title and labels are set such as to maximize their usefulness for the next release's CHANGELOG * [x] If applicable, add a new check to the [release checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)! - [PR Build Summary](https://build.rerun.io/pr/5166) - [Docs preview](https://rerun.io/preview/da9a0d206235016d14d9826e8c42ed8481d29643/docs) - [Examples preview](https://rerun.io/preview/da9a0d206235016d14d9826e8c42ed8481d29643/examples) - [Recent benchmark results](https://build.rerun.io/graphs/crates.html) - [Wasm size tracking](https://build.rerun.io/graphs/sizes.html) --------- Co-authored-by: Andreas Reich --- crates/re_space_view_spatial/src/lib.rs | 1 + .../src/max_image_dimension_subscriber.rs | 96 +++++ .../src/space_view_2d.rs | 334 ++++++++---------- 3 files changed, 247 insertions(+), 184 deletions(-) create mode 100644 crates/re_space_view_spatial/src/max_image_dimension_subscriber.rs diff --git a/crates/re_space_view_spatial/src/lib.rs b/crates/re_space_view_spatial/src/lib.rs index 33edb33a2b85..f7db2a3228de 100644 --- a/crates/re_space_view_spatial/src/lib.rs +++ b/crates/re_space_view_spatial/src/lib.rs @@ -6,6 +6,7 @@ mod contexts; mod eye; mod heuristics; mod instance_hash_conversions; +mod max_image_dimension_subscriber; mod mesh_cache; mod mesh_loader; mod picking; diff --git a/crates/re_space_view_spatial/src/max_image_dimension_subscriber.rs b/crates/re_space_view_spatial/src/max_image_dimension_subscriber.rs new file mode 100644 index 000000000000..1dd5cc4d8ec9 --- /dev/null +++ b/crates/re_space_view_spatial/src/max_image_dimension_subscriber.rs @@ -0,0 +1,96 @@ +use ahash::HashMap; +use nohash_hasher::IntMap; +use once_cell::sync::OnceCell; +use re_data_store::{StoreSubscriber, StoreSubscriberHandle}; +use re_log_types::{EntityPath, StoreId}; +use re_types::{components::TensorData, Loggable}; + +#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)] +pub struct ImageDimensions { + pub height: u64, + pub width: u64, + pub channels: u64, +} + +#[derive(Default, Debug, Clone)] +pub struct MaxImageDimensions(IntMap); + +impl MaxImageDimensions { + /// Accesses the image dimension information for a given store + pub fn access( + store_id: &StoreId, + f: impl FnOnce(&IntMap) -> T, + ) -> Option { + re_data_store::DataStore::with_subscriber_once( + MaxImageDimensionSubscriber::subscription_handle(), + move |subscriber: &MaxImageDimensionSubscriber| { + subscriber.max_dimensions.get(store_id).map(|v| &v.0).map(f) + }, + ) + .flatten() + } +} + +#[derive(Default)] +struct MaxImageDimensionSubscriber { + max_dimensions: HashMap, +} + +impl MaxImageDimensionSubscriber { + /// Accesses the global store subscriber. + /// + /// Lazily registers the subscriber if it hasn't been registered yet. + pub fn subscription_handle() -> StoreSubscriberHandle { + static SUBSCRIPTION: OnceCell = OnceCell::new(); + *SUBSCRIPTION.get_or_init(|| { + re_data_store::DataStore::register_subscriber( + Box::::default(), + ) + }) + } +} + +impl StoreSubscriber for MaxImageDimensionSubscriber { + fn name(&self) -> String { + "MaxImageDimensionStoreSubscriber".to_owned() + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn on_events(&mut self, events: &[re_data_store::StoreEvent]) { + re_tracing::profile_function!(); + + for event in events { + if event.diff.kind != re_data_store::StoreDiffKind::Addition { + // Max image dimensions are strictly additive + continue; + } + + if let Some(cell) = event.diff.cells.get(&TensorData::name()) { + if let Ok(Some(tensor_data)) = cell.try_to_native_mono::() { + if let Some([height, width, channels]) = + tensor_data.image_height_width_channels() + { + let dimensions = self + .max_dimensions + .entry(event.store_id.clone()) + .or_default() + .0 + .entry(event.diff.entity_path.clone()) + .or_default(); + + dimensions.height = dimensions.height.max(height); + dimensions.width = dimensions.width.max(width); + dimensions.channels = dimensions.channels.max(channels); + } + } + } + } + } +} 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 0f86892954a8..cbab4e35f8a0 100644 --- a/crates/re_space_view_spatial/src/space_view_2d.rs +++ b/crates/re_space_view_spatial/src/space_view_2d.rs @@ -1,19 +1,16 @@ -use ahash::HashMap; -use itertools::Itertools; -use nohash_hasher::IntSet; +use ahash::HashSet; +use nohash_hasher::{IntMap, IntSet}; use re_entity_db::{EntityProperties, EntityTree}; use re_log_types::{EntityPath, EntityPathFilter}; -use re_tracing::profile_scope; use re_types::{ archetypes::{DepthImage, Image}, - components::TensorData, Archetype, ComponentName, }; use re_viewer_context::{ - ApplicableEntities, IdentifiedViewSystem as _, PerSystemEntities, RecommendedSpaceView, - SpaceViewClass, SpaceViewClassRegistryError, SpaceViewId, SpaceViewSpawnHeuristics, - SpaceViewSystemExecutionError, ViewQuery, ViewerContext, VisualizableFilterContext, + PerSystemEntities, RecommendedSpaceView, SpaceViewClass, SpaceViewClassRegistryError, + SpaceViewId, SpaceViewSpawnHeuristics, SpaceViewSystemExecutionError, ViewQuery, ViewerContext, + VisualizableFilterContext, }; use crate::{ @@ -21,10 +18,11 @@ use crate::{ heuristics::{ default_visualized_entities_for_visualizer_kind, update_object_property_heuristics, }, - spatial_topology::{SpatialTopology, SubSpace, SubSpaceDimensionality}, + max_image_dimension_subscriber::{ImageDimensions, MaxImageDimensions}, + spatial_topology::{SpatialTopology, SubSpaceDimensionality}, ui::SpatialSpaceViewState, view_kind::SpatialSpaceViewKind, - visualizers::{register_2d_spatial_visualizers, ImageVisualizer}, + visualizers::register_2d_spatial_visualizers, }; #[derive(Default)] @@ -171,91 +169,58 @@ impl SpaceViewClass for SpatialSpaceView2D { SpatialSpaceViewKind::TwoD, ); - let image_entities_fallback = ApplicableEntities::default(); - let image_entities = ctx - .applicable_entities_per_visualizer - .get(&ImageVisualizer::identifier()) - .unwrap_or(&image_entities_fallback); + let image_dimensions = + MaxImageDimensions::access(ctx.entity_db.store_id(), |image_dimensions| { + image_dimensions.clone() + }) + .unwrap_or_default(); // 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 { - recommended_space_views: topo - .iter_subspaces() - .flat_map(|subspace| { - if subspace.dimensionality == SubSpaceDimensionality::ThreeD - || subspace.entities.is_empty() - || indicated_entities.is_disjoint(&subspace.entities) - { - return Vec::new(); - } - - let mut recommended_space_views = Vec::::new(); - - for bucket_entities in - bucket_images_in_subspace(ctx, subspace, image_entities) - { - add_recommended_space_views_for_image_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 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, - ), - }); - } - - recommended_space_views - }) - .collect(), - }) - .unwrap_or_default(); + SpatialTopology::access(ctx.entity_db.store_id(), |topo| SpaceViewSpawnHeuristics { + recommended_space_views: topo + .iter_subspaces() + .flat_map(|subspace| { + if subspace.dimensionality == SubSpaceDimensionality::ThreeD + || subspace.entities.is_empty() + || indicated_entities.is_disjoint(&subspace.entities) + { + return Vec::new(); + } - // 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. - // TODO(jleibs): This is expensive. Would be great to track this as we build up the covering instead. - { - profile_scope!("space_view_2d: find uncovered entities"); - let remaining_entities = indicated_entities - .iter() - .filter(|entity| { - heuristics - .recommended_space_views + // Collect just the 2d-relevant entities in this subspace + let relevant_entities: IntSet = subspace + .entities .iter() - .all(|r| !r.query_filter.is_included(entity)) + .filter(|e| indicated_entities.contains(e)) + .cloned() + .collect(); + + // For explicit 2D spaces (typically indicated by a pinhole or similar) start at the origin, otherwise start at the common ancestor. + // This generally avoids the `/` root entity unless it's required as a common ancestor. + let recommended_root = + if subspace.dimensionality == SubSpaceDimensionality::TwoD { + subspace.origin.clone() + } else { + EntityPath::common_ancestor_of(relevant_entities.iter()) + }; + + let mut recommended_space_views = Vec::::new(); + + recommended_space_views_with_image_splits( + ctx, + &image_dimensions, + &recommended_root, + &relevant_entities, + &mut recommended_space_views, + ); + + recommended_space_views }) - .collect::>(); - - for entity in remaining_entities { - heuristics - .recommended_space_views - .push(RecommendedSpaceView { - root: entity.clone(), - query_filter: EntityPathFilter::single_entity_filter(entity), - }); - } - } - - heuristics + .collect(), + }) + .unwrap_or_default() } fn selection_ui( @@ -292,18 +257,15 @@ 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( +// Count the number of image entities with the given component exist that aren't +// children of other entities in the bucket. +fn count_non_nested_images_with_component( + image_dimensions: &IntMap, entity_bucket: &IntSet, subtree: &EntityTree, component_name: &ComponentName, ) -> usize { - if entity_bucket.contains(&subtree.path) { + if image_dimensions.contains_key(&subtree.path) { // bool true -> 1 subtree.entity.components.contains_key(component_name) as usize } else if !entity_bucket @@ -316,121 +278,125 @@ fn count_non_nested_entities_with_component( .children .values() .map(|child| { - count_non_nested_entities_with_component(entity_bucket, child, component_name) + count_non_nested_images_with_component( + image_dimensions, + entity_bucket, + child, + component_name, + ) }) .sum() } } -/// Given a bucket of image entities. -fn add_recommended_space_views_for_image_bucket( - ctx: &ViewerContext<'_>, +// Find the image dimensions of every image-entity in the bucket that is not +// not nested under another image. +// +// We track a set of just height/width as different channels could be allowed to +// stack. +fn find_non_nested_image_dimensions( + image_dimensions: &IntMap, entity_bucket: &IntSet, - recommended: &mut Vec, + subtree: &EntityTree, + found_image_dimensions: &mut HashSet<[u64; 2]>, ) { - // TODO(jleibs): Converting entity_bucket to a Trie would probably make some of this easier. - let tree = ctx.entity_db.tree(); + if let Some(dimensions) = image_dimensions.get(&subtree.path) { + // If we found an image entity, add its dimensions to the set. + found_image_dimensions.insert([dimensions.height, dimensions.width]); + } else if entity_bucket + .iter() + .any(|e| e.is_descendant_of(&subtree.path)) + { + // Otherwise recurse + for child in subtree.children.values() { + find_non_nested_image_dimensions( + image_dimensions, + entity_bucket, + child, + found_image_dimensions, + ); + } + } +} - // Find the common ancestor of the bucket - let root = EntityPath::common_ancestor_of(entity_bucket.iter()); +fn recommended_space_views_with_image_splits( + ctx: &ViewerContext<'_>, + image_dimensions: &IntMap, + recommended_root: &EntityPath, + entities: &IntSet, + recommended: &mut Vec, +) { + re_tracing::profile_function!(); - // 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; - } + let tree = ctx.entity_db.tree(); - // Alternatively we want to split this bucket into a group for each child-space. - let Some(subtree) = tree.subtree(&root) else { + let Some(subtree) = tree.subtree(recommended_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, + let mut found_image_dimensions = Default::default(); + + find_non_nested_image_dimensions( + image_dimensions, + entities, + subtree, + &mut found_image_dimensions, + ); + + let image_count = count_non_nested_images_with_component( + image_dimensions, + entities, subtree, &Image::indicator().name(), ); - let depth_count = count_non_nested_entities_with_component( - entity_bucket, + let depth_count = count_non_nested_images_with_component( + image_dimensions, + entities, 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_image_bucket(ctx, &sub_bucket, recommended); + // If there are images of multiple dimensions, more than 1 image, or more than 1 depth image + // then split the space. + if found_image_dimensions.len() > 1 || image_count > 1 || depth_count > 1 { + // Otherwise, split the space and recurse + + // If the root also had a visualizable entity, give it its own space. + // TODO(jleibs): Maybe merge this entity into each child + if entities.contains(recommended_root) { + recommended.push(RecommendedSpaceView { + root: recommended_root.clone(), + query_filter: EntityPathFilter::single_entity_filter(recommended_root), + }); } - } -} -/// Groups all images in the subspace by size and draw order. -fn bucket_images_in_subspace( - ctx: &ViewerContext<'_>, - subspace: &SubSpace, - image_entities: &ApplicableEntities, -) -> Vec> { - re_tracing::profile_function!(); - - let store = ctx.entity_db.store(); - - let image_entities = subspace - .entities - .iter() - .filter(|e| image_entities.contains(e)) - .collect_vec(); - - if image_entities.len() <= 1 { - // Very common case, early out before we get into the more expensive query code. - return image_entities - .into_iter() - .map(|e| std::iter::once(e.clone()).collect()) - .collect(); - } - - 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 (?). - if let Some(tensor) = - store.query_latest_component::(image_entity, &ctx.current_query()) - { - if let Some([height, width, _]) = tensor.image_height_width_channels() { - // 1D tensors are typically handled by tensor or bar chart views and make generally for poor image buckets! - if height > 1 && width > 1 { - images_by_bucket - .entry((height, width)) - .or_default() - .insert(image_entity.clone()); - } + // And then recurse into the children + for child in subtree.children.values() { + let sub_bucket: IntSet<_> = entities + .iter() + .filter(|e| e.starts_with(&child.path)) + .cloned() + .collect(); + + if !sub_bucket.is_empty() { + recommended_space_views_with_image_splits( + ctx, + image_dimensions, + &child.path, + &sub_bucket, + recommended, + ); } } + } else { + // Otherwise we can use the space as it is + recommended.push(RecommendedSpaceView { + root: recommended_root.clone(), + query_filter: EntityPathFilter::subtree_entity_filter(recommended_root), + }); } - - images_by_bucket.drain().map(|(_, v)| v).collect() }