Split Image visualizer into Image/DepthImage/SegmentationImage (#6654)

### What

* follow-up to #6644
* Fixes #6548
* Fixes #6660

Lots of copy & paste between the three visualizers left. Didn't want to
remove it entirely to not make it too hard to have them diverge more in
the future. But thanks to some clean-up passes it's not too bad (imho).


![image](https://github.com/rerun-io/rerun/assets/1220815/fe2b92eb-01ce-4864-b32d-7613355c763c)

![image](https://github.com/rerun-io/rerun/assets/1220815/1ed948d7-5cbf-45cd-9204-244aec4cd3f7)

![image](https://github.com/rerun-io/rerun/assets/1220815/3fe3c087-b39c-49ec-b795-6f3507059604)


### 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 examples from latest `main` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/6654?manifest_url=https://app.rerun.io/version/main/examples_manifest.json)
* Using full set of examples from `nightly` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/6654?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/6654)
- [Recent benchmark results](https://build.rerun.io/graphs/crates.html)
- [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)

To run all checks from `main`, comment on the PR with `@rerun-bot
full-check`.

crates/re_data_ui/src/image.rs

@@ -12,7 +12,7 @@ use re_viewer_context::{
     ViewerContext,
 };
 
-use crate::image_meaning_for_entity;
+use crate::image_meaning::image_meaning_for_entity;
 
 use super::EntityDataUi;
 

crates/re_data_ui/src/lib.rs

@@ -37,7 +37,6 @@ pub use crate::image::{
 };
 pub use component::EntityLatestAtResults;
 pub use component_ui_registry::{add_to_registry, create_component_ui_registry};
-pub use image_meaning::image_meaning_for_entity;
 
 /// Sort components for display in the UI.
 pub fn sorted_component_list_for_ui<'a>(

crates/re_space_view_spatial/src/lib.rs

@@ -12,6 +12,7 @@ mod instance_hash_conversions;
 mod max_image_dimension_subscriber;
 mod mesh_cache;
 mod mesh_loader;
+mod pickable_image;
 mod picking;
 mod scene_bounding_boxes;
 mod space_camera_3d;
@@ -25,13 +26,16 @@ mod view_3d;
 mod view_3d_properties;
 mod visualizers;
 
-use re_space_view::DataResultQuery as _;
-use re_viewer_context::{ViewContext, ViewerContext};
 pub use view_2d::SpatialSpaceView2D;
 pub use view_3d::SpatialSpaceView3D;
 
+pub(crate) use pickable_image::PickableImageRect;
+
 // ---
 
+use re_space_view::DataResultQuery as _;
+use re_viewer_context::{ViewContext, ViewerContext};
+
 use re_renderer::RenderContext;
 use re_types::blueprint::components::BackgroundKind;
 use re_types::components::{Color, Resolution, TensorData};

crates/re_space_view_spatial/src/pickable_image.rs

@@ -0,0 +1,11 @@
+use re_log_types::EntityPath;
+use re_renderer::renderer::TexturedRect;
+
+/// Image rectangle that can be picked in the view.
+pub struct PickableImageRect {
+    /// Path to the image (note image instance ids would refer to pixels!)
+    pub ent_path: EntityPath,
+
+    /// Textured rectangle used by the renderer.
+    pub textured_rect: TexturedRect,
+}

crates/re_space_view_spatial/src/picking.rs

@@ -1,11 +1,12 @@
 //! Handles picking in 2D and 3D spaces.
 
-use ahash::HashSet;
+use std::collections::HashSet;
+
 use re_entity_db::InstancePathHash;
 use re_log_types::Instance;
 use re_renderer::PickingLayerProcessor;
 
-use crate::visualizers::ViewerImage;
+use crate::PickableImageRect;
 use crate::{eye::Eye, instance_hash_conversions::instance_path_hash_from_picking_layer_id};
 
 #[derive(Clone, PartialEq, Eq)]
@@ -102,12 +103,12 @@ impl PickingContext {
     }
 
     /// Performs picking for a given scene.
-    pub fn pick(
+    pub fn pick<'a>(
         &self,
         render_ctx: &re_renderer::RenderContext,
         gpu_readback_identifier: re_renderer::GpuReadbackIdentifier,
         previous_picking_result: &Option<PickingResult>,
-        images: &[ViewerImage],
+        images: impl Iterator<Item = &'a PickableImageRect>,
         ui_rects: &[PickableUiRect],
     ) -> PickingResult {
         re_tracing::profile_function!();
@@ -240,11 +241,16 @@ fn picking_gpu(
     }
 }
 
-fn picking_textured_rects(context: &PickingContext, images: &[ViewerImage]) -> Vec<PickingRayHit> {
+fn picking_textured_rects<'a>(
+    context: &PickingContext,
+    images: impl Iterator<Item = &'a PickableImageRect>,
+) -> Vec<PickingRayHit> {
     re_tracing::profile_function!();
 
     let mut hits = Vec::new();
 
+    let mut hit_image_rect_entities = HashSet::new();
+
     for image in images {
         let rect = &image.textured_rect;
         let Some(normal) = rect.extent_u.cross(rect.extent_v).try_normalize() else {
@@ -266,18 +272,24 @@ fn picking_textured_rects(context: &PickingContext, images: &[ViewerImage]) -> V
 
         if (0.0..=1.0).contains(&u) && (0.0..=1.0).contains(&v) {
             let [width, height] = rect.colormapped_texture.width_height();
-            hits.push(PickingRayHit {
-                instance_path_hash: InstancePathHash {
-                    entity_path_hash: image.ent_path.hash(),
-                    instance: Instance::from_2d_image_coordinate(
-                        [(u * width as f32) as u32, (v * height as f32) as u32],
-                        width as u64,
-                    ),
-                },
-                space_position: intersection_world,
-                hit_type: PickingHitType::TexturedRect,
-                depth_offset: rect.options.depth_offset,
-            });
+
+            // Ignore the image if we hit the same entity already as an image.
+            // This happens if the same entity has multiple textured rects.
+            let entity_path_hash = image.ent_path.hash();
+            if hit_image_rect_entities.insert(entity_path_hash) {
+                hits.push(PickingRayHit {
+                    instance_path_hash: InstancePathHash {
+                        entity_path_hash,
+                        instance: Instance::from_2d_image_coordinate(
+                            [(u * width as f32) as u32, (v * height as f32) as u32],
+                            width as u64,
+                        ),
+                    },
+                    space_position: intersection_world,
+                    hit_type: PickingHitType::TexturedRect,
+                    depth_offset: rect.options.depth_offset,
+                });
+            }
         }
     }
 

crates/re_space_view_spatial/src/ui.rs

@@ -4,8 +4,7 @@ use egui::{epaint::util::OrderedFloat, text::TextWrapping, NumExt, WidgetText};
 use macaw::BoundingBox;
 
 use re_data_ui::{
-    image_meaning_for_entity, item_ui, show_zoomed_image_region,
-    show_zoomed_image_region_area_outline, DataUi,
+    item_ui, show_zoomed_image_region, show_zoomed_image_region_area_outline, DataUi,
 };
 use re_format::format_f32;
 use re_log_types::Instance;
@@ -29,7 +28,10 @@ use crate::{
     contexts::AnnotationSceneContext,
     picking::{PickableUiRect, PickingContext, PickingHitType, PickingResult},
     view_kind::SpatialSpaceViewKind,
-    visualizers::{CamerasVisualizer, ImageVisualizer, UiLabel, UiLabelTarget},
+    visualizers::{
+        CamerasVisualizer, DepthImageVisualizer, ImageVisualizer, SegmentationImageVisualizer,
+        UiLabel, UiLabelTarget,
+    },
 };
 use crate::{contexts::PrimitiveCounter, scene_bounding_boxes::SceneBoundingBoxes};
 use crate::{eye::EyeMode, heuristics::auto_size_world_heuristic};
@@ -108,13 +110,10 @@ impl SpatialSpaceViewState {
             .num_primitives
             .load(std::sync::atomic::Ordering::Relaxed);
 
-        self.num_non_segmentation_images_last_frame = system_output
-            .view_systems
-            .get::<ImageVisualizer>()?
-            .images
-            .iter()
-            .filter(|i| i.meaning != TensorDataMeaning::ClassId)
-            .count();
+        let view_systems = &system_output.view_systems;
+        let num_images = view_systems.get::<ImageVisualizer>()?.images.len();
+        let num_depth_images = view_systems.get::<DepthImageVisualizer>()?.images.len();
+        self.num_non_segmentation_images_last_frame = num_images + num_depth_images;
 
         Ok(())
     }
@@ -499,12 +498,18 @@ pub fn picking(
 
     let annotations = view_ctx.get::<AnnotationSceneContext>()?;
     let images = visualizers.get::<ImageVisualizer>()?;
+    let depth_images = visualizers.get::<DepthImageVisualizer>()?;
+    let segmentation_images = visualizers.get::<SegmentationImageVisualizer>()?;
 
     let picking_result = picking_context.pick(
         render_ctx,
         query.space_view_id.gpu_readback_id(),
         &state.previous_picking_result,
-        &images.images,
+        images
+            .images
+            .iter()
+            .chain(depth_images.images.iter())
+            .chain(segmentation_images.images.iter()),
         ui_rects,
     );
     state.previous_picking_result = Some(picking_result.clone());
@@ -546,10 +551,8 @@ pub fn picking(
             continue;
         }
 
-        let store = ctx.recording_store();
-
         // Special hover ui for images.
-        let is_depth_cloud = images
+        let is_depth_cloud = depth_images
             .depth_cloud_entities
             .contains(&instance_path.entity_path.hash());
 
@@ -563,11 +566,22 @@ pub fn picking(
         }
 
         let picked_image = if hit.hit_type == PickingHitType::TexturedRect || is_depth_cloud {
-            let meaning = image_meaning_for_entity(
-                &instance_path.entity_path,
-                &query.latest_at_query(),
-                store,
-            );
+            let meaning = if segmentation_images
+                .images
+                .iter()
+                .any(|i| i.ent_path == instance_path.entity_path)
+            {
+                TensorDataMeaning::ClassId
+            } else if is_depth_cloud
+                || depth_images
+                    .images
+                    .iter()
+                    .any(|i| i.ent_path == instance_path.entity_path)
+            {
+                TensorDataMeaning::Depth
+            } else {
+                TensorDataMeaning::Unknown
+            };
 
             let query_shadowed_defaults = false;
             let results = latest_at_with_blueprint_resolved_data(