move parts of annotation context machinery to re_space_view, remove u…

…nnecessary keypoint queries from many visualizers, add colors from annotations to map view

Cargo.lock

@@ -6083,8 +6083,10 @@ dependencies = [
 name = "re_space_view"
 version = "0.20.0-alpha.4+dev"
 dependencies = [
+ "ahash",
  "bytemuck",
  "egui",
+ "glam",
  "itertools 0.13.0",
  "nohash-hasher",
  "re_chunk_store",
@@ -6094,6 +6096,7 @@ dependencies = [
  "re_query",
  "re_renderer",
  "re_tracing",
+ "re_types",
  "re_types_core",
  "re_ui",
  "re_viewer_context",
@@ -6145,6 +6148,7 @@ dependencies = [
 name = "re_space_view_map"
 version = "0.20.0-alpha.4+dev"
 dependencies = [
+ "bytemuck",
  "egui",
  "glam",
  "itertools 0.13.0",

crates/store/re_types/definitions/rerun/archetypes/geo_points.fbs

@@ -30,6 +30,11 @@ table GeoPoints (
   /// \py As either 0-1 floats or 0-255 integers, with separate alpha.
   colors: [rerun.components.Color] ("attr.rerun.component_recommended", nullable, order: 2100);
 
+  /// Optional class Ids for the points.
+  ///
+  /// The [components.ClassId] provides colors if not specified explicitly.
+  class_ids: [rerun.components.ClassId] ("attr.rerun.component_optional", nullable, order: 3200);
+
   //TODO(ab): add `Label` and  `ShowLabels` components
   //TODO(ab): add `Altitude` component
 }

crates/viewer/re_space_view/Cargo.toml

@@ -30,11 +30,14 @@ re_query.workspace = true
 re_renderer.workspace = true
 re_tracing.workspace = true
 re_types_core.workspace = true
+re_types.workspace = true
 re_ui.workspace = true
 re_viewer_context.workspace = true
 re_viewport_blueprint.workspace = true
 
+ahash.workspace = true
 bytemuck.workspace = true
 egui.workspace = true
+glam.workspace = true
 itertools.workspace = true
 nohash-hasher.workspace = true

crates/viewer/re_space_view/src/annotation_context_utils.rs

@@ -0,0 +1,137 @@
+use ahash::HashMap;
+
+use re_types::components::Color;
+use re_viewer_context::{Annotations, QueryContext, ResolvedAnnotationInfos};
+
+use crate::clamped_or_nothing;
+
+#[inline]
+fn to_egui_color(color: &Color) -> egui::Color32 {
+    let [r, g, b, a] = color.to_array();
+    egui::Color32::from_rgba_unmultiplied(r, g, b, a)
+}
+
+/// Process [`Color`] components using annotations and default colors.
+pub fn process_color_slice<'a>(
+    ctx: &QueryContext<'_>,
+    fallback_provider: &'a dyn re_viewer_context::TypedComponentFallbackProvider<Color>,
+    num_instances: usize,
+    annotation_infos: &'a ResolvedAnnotationInfos,
+    colors: &'a [Color],
+) -> Vec<egui::Color32> {
+    // NOTE: Do not put tracing scopes here, this is called for every entity/timestamp in a frame.
+
+    if let Some(last_color) = colors.last() {
+        // If we have colors we can ignore the annotation infos/contexts.
+
+        if colors.len() == num_instances {
+            // Common happy path
+            colors.iter().map(to_egui_color).collect()
+        } else if colors.len() == 1 {
+            // Common happy path
+            vec![to_egui_color(last_color); num_instances]
+        } else {
+            let colors = clamped_or_nothing(colors, num_instances);
+            colors.map(to_egui_color).collect()
+        }
+    } else {
+        match annotation_infos {
+            ResolvedAnnotationInfos::Same(count, annotation_info) => {
+                re_tracing::profile_scope!("no colors, same annotation");
+                let color = annotation_info
+                    .color()
+                    .unwrap_or_else(|| fallback_provider.fallback_for(ctx).into());
+                vec![color; *count]
+            }
+            ResolvedAnnotationInfos::Many(annotation_info) => {
+                re_tracing::profile_scope!("no-colors, many annotations");
+                let fallback = fallback_provider.fallback_for(ctx).into();
+                annotation_info
+                    .iter()
+                    .map(|annotation_info| annotation_info.color().unwrap_or(fallback))
+                    .collect()
+            }
+        }
+    }
+}
+
+pub type Keypoints = HashMap<
+    (re_types::components::ClassId, i64),
+    HashMap<re_types::datatypes::KeypointId, glam::Vec3>,
+>;
+
+/// Resolves all annotations and keypoints for the given entity view.
+pub fn process_annotation_and_keypoint_slices(
+    latest_at: re_log_types::TimeInt,
+    num_instances: usize,
+    positions: impl Iterator<Item = glam::Vec3>,
+    keypoint_ids: &[re_types::components::KeypointId],
+    class_ids: &[re_types::components::ClassId],
+    annotations: &Annotations,
+) -> (ResolvedAnnotationInfos, Keypoints) {
+    re_tracing::profile_function!();
+
+    let mut keypoints: Keypoints = HashMap::default();
+
+    // No need to process annotations if we don't have class-ids
+    if class_ids.is_empty() {
+        let resolved_annotation = annotations
+            .resolved_class_description(None)
+            .annotation_info();
+
+        return (
+            ResolvedAnnotationInfos::Same(num_instances, resolved_annotation),
+            keypoints,
+        );
+    };
+
+    let class_ids = clamped_or_nothing(class_ids, num_instances);
+
+    if keypoint_ids.is_empty() {
+        let annotation_info = class_ids
+            .map(|&class_id| {
+                let class_description = annotations.resolved_class_description(Some(class_id));
+                class_description.annotation_info()
+            })
+            .collect();
+
+        (
+            ResolvedAnnotationInfos::Many(annotation_info),
+            Default::default(),
+        )
+    } else {
+        let keypoint_ids = clamped_or_nothing(keypoint_ids, num_instances);
+        let annotation_info = itertools::izip!(positions, keypoint_ids, class_ids)
+            .map(|(position, keypoint_id, &class_id)| {
+                let class_description = annotations.resolved_class_description(Some(class_id));
+
+                keypoints
+                    .entry((class_id, latest_at.as_i64()))
+                    .or_default()
+                    .insert(keypoint_id.0, position);
+                class_description.annotation_info_with_keypoint(keypoint_id.0)
+            })
+            .collect();
+
+        (ResolvedAnnotationInfos::Many(annotation_info), keypoints)
+    }
+}
+
+/// Resolves all annotations for the given entity view.
+pub fn process_annotation_slices(
+    latest_at: re_log_types::TimeInt,
+    num_instances: usize,
+    class_ids: &[re_types::components::ClassId],
+    annotations: &Annotations,
+) -> ResolvedAnnotationInfos {
+    let (annotations, _keypoints) = process_annotation_and_keypoint_slices(
+        latest_at,
+        num_instances,
+        std::iter::empty(), // positions are only needed for keypoint lookup
+        &[],
+        class_ids,
+        annotations,
+    );
+
+    annotations
+}

crates/viewer/re_space_view/src/lib.rs

@@ -4,6 +4,8 @@
 
 pub mod controls;
 
+mod annotation_context_utils;
+mod annotation_scene_context;
 mod heuristics;
 mod instance_hash_conversions;
 mod outlines;
@@ -12,6 +14,10 @@ mod results_ext;
 mod screenshot;
 mod view_property_ui;
 
+pub use annotation_context_utils::{
+    process_annotation_and_keypoint_slices, process_annotation_slices, process_color_slice,
+};
+pub use annotation_scene_context::AnnotationSceneContext;
 pub use heuristics::suggest_space_view_for_each_entity;
 pub use instance_hash_conversions::{
     instance_path_hash_from_picking_layer_id, picking_layer_id_from_instance_path_hash,
@@ -53,3 +59,20 @@ pub fn diff_component_filter<T: re_types_core::Component>(
                 .any(|instances| instances.iter().any(filter))
         })
 }
+
+/// Clamp the last value in `values` in order to reach a length of `clamped_len`.
+///
+/// Returns an empty iterator if values is empty.
+#[inline]
+pub fn clamped_or_nothing<T>(values: &[T], clamped_len: usize) -> impl Iterator<Item = &T> + Clone {
+    let Some(last) = values.last() else {
+        return itertools::Either::Left(std::iter::empty());
+    };
+
+    itertools::Either::Right(
+        values
+            .iter()
+            .chain(std::iter::repeat(last))
+            .take(clamped_len),
+    )
+}

crates/viewer/re_space_view_map/Cargo.toml

@@ -33,6 +33,7 @@ re_ui.workspace = true
 re_viewer_context.workspace = true
 re_viewport_blueprint.workspace = true
 
+bytemuck.workspace = true
 egui.workspace = true
 glam.workspace = true
 itertools.workspace = true

crates/viewer/re_space_view_map/src/map_space_view.rs

@@ -1,4 +1,5 @@
 use egui::{Context, NumExt as _, Rect, Response};
+use re_space_view::AnnotationSceneContext;
 use walkers::{HttpTiles, Map, MapMemory, Tiles};
 
 use re_data_ui::{item_ui, DataUi};
@@ -119,7 +120,11 @@ Displays geospatial primitives on a map.
         system_registry: &mut SpaceViewSystemRegistrator<'_>,
     ) -> Result<(), SpaceViewClassRegistryError> {
         system_registry.register_visualizer::<GeoPointsVisualizer>()?;
-        system_registry.register_visualizer::<GeoLineStringsVisualizer>()
+        system_registry.register_visualizer::<GeoLineStringsVisualizer>()?;
+
+        system_registry.register_context_system::<AnnotationSceneContext>()?;
+
+        Ok(())
     }
 
     fn new_state(&self) -> Box<dyn SpaceViewState> {