Remove the ability to display multiple tensors in a single space view

crates/re_space_view_tensor/src/space_view_class.rs

@@ -32,10 +32,20 @@ type ViewType = re_types::blueprint::views::TensorView;
 
 #[derive(Default)]
 pub struct ViewTensorState {
-    /// Selects in [`Self::state_tensors`].
-    pub selected_tensor: Option<EntityPath>,
+    /// What slice are we viewing?
+    ///
+    /// This get automatically reset if/when the current tensor shape changes.
+    pub(crate) slice: SliceSelection,
 
-    pub state_tensors: ahash::HashMap<EntityPath, PerTensorState>,
+    /// How we map values to colors.
+    pub(crate) color_mapping: ColorMapping,
+
+    /// Scaling, filtering, aspect ratio, etc for the rendered texture.
+    texture_settings: TextureSettings,
+
+    /// Last viewed tensor, copied each frame.
+    /// Used for the selection view.
+    tensor: Option<(RowId, DecodedTensor)>,
 }
 
 impl SpaceViewState for ViewTensorState {
@@ -58,88 +68,6 @@ pub struct SliceSelection {
     pub selector_values: BTreeMap<usize, u64>,
 }
 
-pub struct PerTensorState {
-    /// What slice are we vieiwing?
-    slice: SliceSelection,
-
-    /// How we map values to colors.
-    color_mapping: ColorMapping,
-
-    /// Scaling, filtering, aspect ratio, etc for the rendered texture.
-    texture_settings: TextureSettings,
-
-    /// Last viewed tensor, copied each frame.
-    /// Used for the selection view.
-    tensor: Option<(RowId, DecodedTensor)>,
-}
-
-impl PerTensorState {
-    pub fn create(tensor_data_row_id: RowId, tensor: &DecodedTensor) -> Self {
-        Self {
-            slice: SliceSelection {
-                dim_mapping: DimensionMapping::create(tensor.shape()),
-                selector_values: Default::default(),
-            },
-            color_mapping: ColorMapping::default(),
-            texture_settings: TextureSettings::default(),
-            tensor: Some((tensor_data_row_id, tensor.clone())),
-        }
-    }
-
-    pub fn slice(&self) -> &SliceSelection {
-        &self.slice
-    }
-
-    pub fn color_mapping(&self) -> &ColorMapping {
-        &self.color_mapping
-    }
-
-    pub fn ui(&mut self, ctx: &ViewerContext<'_>, ui: &mut egui::Ui) {
-        let Some((tensor_data_row_id, tensor)) = &self.tensor else {
-            ui.label("No Tensor shown in this Space View.");
-            return;
-        };
-
-        let tensor_stats = ctx
-            .cache
-            .entry(|c: &mut TensorStatsCache| c.entry(*tensor_data_row_id, tensor));
-        ctx.re_ui
-            .selection_grid(ui, "tensor_selection_ui")
-            .show(ui, |ui| {
-                // We are in a bare Tensor view -- meaning / meter is unknown.
-                let meaning = TensorDataMeaning::Unknown;
-                let meter = None;
-                tensor_summary_ui_grid_contents(
-                    ctx.re_ui,
-                    ui,
-                    tensor,
-                    tensor,
-                    meaning,
-                    meter,
-                    &tensor_stats,
-                );
-                self.texture_settings.ui(ctx.re_ui, ui);
-                self.color_mapping.ui(ctx.render_ctx, ctx.re_ui, ui);
-            });
-
-        ui.separator();
-        ui.strong("Dimension Mapping");
-        dimension_mapping_ui(ctx.re_ui, ui, &mut self.slice.dim_mapping, tensor.shape());
-        let default_mapping = DimensionMapping::create(tensor.shape());
-        if ui
-            .add_enabled(
-                self.slice.dim_mapping != default_mapping,
-                egui::Button::new("Reset mapping"),
-            )
-            .on_disabled_hover_text("The default is already set up")
-            .on_hover_text("Reset dimension mapping to the default")
-            .clicked()
-        {
-            self.slice.dim_mapping = DimensionMapping::create(tensor.shape());
-        }
-    }
-}
-
 impl SpaceViewClass for TensorSpaceView {
     fn identifier() -> SpaceViewClassIdentifier {
         ViewType::identifier()
@@ -207,11 +135,51 @@ impl SpaceViewClass for TensorSpaceView {
         _root_entity_properties: &mut EntityProperties,
     ) -> Result<(), SpaceViewSystemExecutionError> {
         let state = state.downcast_mut::<ViewTensorState>()?;
-        if let Some(selected_tensor) = &state.selected_tensor {
-            if let Some(state_tensor) = state.state_tensors.get_mut(selected_tensor) {
-                state_tensor.ui(ctx, ui);
+
+        ctx.re_ui
+            .selection_grid(ui, "tensor_selection_ui")
+            .show(ui, |ui| {
+                if let Some((tensor_data_row_id, tensor)) = &state.tensor {
+                    let tensor_stats = ctx
+                        .cache
+                        .entry(|c: &mut TensorStatsCache| c.entry(*tensor_data_row_id, tensor));
+
+                    // We are in a bare Tensor view -- meaning / meter is unknown.
+                    let meaning = TensorDataMeaning::Unknown;
+                    let meter = None;
+                    tensor_summary_ui_grid_contents(
+                        ctx.re_ui,
+                        ui,
+                        tensor,
+                        tensor,
+                        meaning,
+                        meter,
+                        &tensor_stats,
+                    );
+                }
+
+                state.texture_settings.ui(ctx.re_ui, ui);
+                state.color_mapping.ui(ctx.render_ctx, ctx.re_ui, ui);
+            });
+
+        if let Some((_, tensor)) = &state.tensor {
+            ui.separator();
+            ui.strong("Dimension Mapping");
+            dimension_mapping_ui(ctx.re_ui, ui, &mut state.slice.dim_mapping, tensor.shape());
+            let default_mapping = DimensionMapping::create(tensor.shape());
+            if ui
+                .add_enabled(
+                    state.slice.dim_mapping != default_mapping,
+                    egui::Button::new("Reset mapping"),
+                )
+                .on_disabled_hover_text("The default is already set up")
+                .on_hover_text("Reset dimension mapping to the default")
+                .clicked()
+            {
+                state.slice.dim_mapping = DimensionMapping::create(tensor.shape());
             }
         }
+
         Ok(())
     }
 
@@ -238,40 +206,26 @@ impl SpaceViewClass for TensorSpaceView {
 
         let tensors = &system_output.view_systems.get::<TensorSystem>()?.tensors;
 
-        if tensors.is_empty() {
-            ui.centered_and_justified(|ui| ui.label("(empty)"));
-            state.selected_tensor = None;
-        } else {
-            if let Some(selected_tensor) = &state.selected_tensor {
-                if !tensors.contains_key(selected_tensor) {
-                    state.selected_tensor = None;
-                }
-            }
-            if state.selected_tensor.is_none() {
-                state.selected_tensor = Some(tensors.iter().next().unwrap().0.clone());
-            }
-
-            if tensors.len() > 1 {
-                // Show radio buttons for the different tensors we have in this view - better than nothing!
-                ui.horizontal(|ui| {
-                    for instance_path in tensors.keys() {
-                        let is_selected = state.selected_tensor.as_ref() == Some(instance_path);
-                        if ui.radio(is_selected, instance_path.to_string()).clicked() {
-                            state.selected_tensor = Some(instance_path.clone());
-                        }
-                    }
-                });
-            }
+        if tensors.len() > 1 {
+            state.tensor = None;
 
-            if let Some(selected_tensor) = &state.selected_tensor {
-                if let Some((tensor_data_row_id, tensor)) = tensors.get(selected_tensor) {
-                    let state_tensor = state
-                        .state_tensors
-                        .entry(selected_tensor.clone())
-                        .or_insert_with(|| PerTensorState::create(*tensor_data_row_id, tensor));
-                    view_tensor(ctx, ui, state_tensor, *tensor_data_row_id, tensor);
-                }
+            egui::Frame {
+                inner_margin: re_ui::ReUi::view_padding().into(),
+                ..egui::Frame::default()
             }
+            .show(ui, |ui| {
+                ui.label(format!(
+                    "Can only show one tensor at a time; was given {}. Update the query so that it \
+                    returns a single tensor entity and create additional views for the others.",
+                    tensors.len()
+                ));
+            });
+        } else if let Some((tensor_data_row_id, tensor)) = tensors.first() {
+            state.tensor = Some((*tensor_data_row_id, tensor.clone()));
+            view_tensor(ctx, ui, state, *tensor_data_row_id, tensor);
+        } else {
+            state.tensor = None;
+            ui.centered_and_justified(|ui| ui.label("(empty)"));
         }
 
         Ok(())
@@ -281,14 +235,12 @@ impl SpaceViewClass for TensorSpaceView {
 fn view_tensor(
     ctx: &ViewerContext<'_>,
     ui: &mut egui::Ui,
-    state: &mut PerTensorState,
+    state: &mut ViewTensorState,
     tensor_data_row_id: RowId,
     tensor: &DecodedTensor,
 ) {
     re_tracing::profile_function!();
 
-    state.tensor = Some((tensor_data_row_id, tensor.clone()));
-
     if !state.slice.dim_mapping.is_valid(tensor.num_dim()) {
         state.slice.dim_mapping = DimensionMapping::create(tensor.shape());
     }
@@ -339,7 +291,7 @@ fn view_tensor(
 fn tensor_slice_ui(
     ctx: &ViewerContext<'_>,
     ui: &mut egui::Ui,
-    state: &PerTensorState,
+    state: &ViewTensorState,
     tensor_data_row_id: RowId,
     tensor: &DecodedTensor,
     dimension_labels: [(String, bool); 2],
@@ -358,7 +310,7 @@ fn tensor_slice_ui(
 fn paint_tensor_slice(
     ctx: &ViewerContext<'_>,
     ui: &mut egui::Ui,
-    state: &PerTensorState,
+    state: &ViewTensorState,
     tensor_data_row_id: RowId,
     tensor: &DecodedTensor,
 ) -> anyhow::Result<(egui::Response, egui::Painter, egui::Rect)> {
@@ -750,7 +702,7 @@ fn paint_axis_names(
     }
 }
 
-fn selectors_ui(ui: &mut egui::Ui, state: &mut PerTensorState, tensor: &TensorData) {
+fn selectors_ui(ui: &mut egui::Ui, state: &mut ViewTensorState, tensor: &TensorData) {
     for selector in &state.slice.dim_mapping.selectors {
         if !selector.visible {
             continue;

crates/re_space_view_tensor/src/tensor_slice_to_gpu.rs

@@ -12,7 +12,7 @@ use re_viewer_context::{
     TensorStats,
 };
 
-use crate::space_view_class::{selected_tensor_slice, PerTensorState, SliceSelection};
+use crate::space_view_class::{selected_tensor_slice, SliceSelection, ViewTensorState};
 
 #[derive(thiserror::Error, Debug, PartialEq)]
 pub enum TensorUploadError {
@@ -31,24 +31,22 @@ pub fn colormapped_texture(
     tensor_data_row_id: RowId,
     tensor: &DecodedTensor,
     tensor_stats: &TensorStats,
-    state: &PerTensorState,
+    state: &ViewTensorState,
 ) -> Result<ColormappedTexture, TextureManager2DError<TensorUploadError>> {
     re_tracing::profile_function!();
 
     let range = tensor_data_range_heuristic(tensor_stats, tensor.dtype())
         .map_err(|err| TextureManager2DError::DataCreation(err.into()))?;
     let texture =
-        upload_texture_slice_to_gpu(render_ctx, tensor_data_row_id, tensor, state.slice())?;
-
-    let color_mapping = state.color_mapping();
+        upload_texture_slice_to_gpu(render_ctx, tensor_data_row_id, tensor, &state.slice)?;
 
     Ok(ColormappedTexture {
         texture,
         range,
         decode_srgb: false,
         multiply_rgb_with_alpha: false,
-        gamma: color_mapping.gamma,
-        color_mapper: re_renderer::renderer::ColorMapper::Function(color_mapping.map),
+        gamma: state.color_mapping.gamma,
+        color_mapper: re_renderer::renderer::ColorMapper::Function(state.color_mapping.map),
         shader_decoding: match tensor.buffer {
             TensorBuffer::Nv12(_) => Some(ShaderDecoding::Nv12),
             TensorBuffer::Yuy2(_) => Some(ShaderDecoding::Yuy2),

crates/re_space_view_tensor/src/visualizer_system.rs

@@ -9,7 +9,7 @@ use re_viewer_context::{
 
 #[derive(Default)]
 pub struct TensorSystem {
-    pub tensors: std::collections::BTreeMap<EntityPath, (RowId, DecodedTensor)>,
+    pub tensors: Vec<(RowId, DecodedTensor)>,
 }
 
 impl IdentifiedViewSystem for TensorSystem {
@@ -64,8 +64,7 @@ impl TensorSystem {
             .entry(|c: &mut TensorDecodeCache| c.entry(row_id, tensor.value.0))
         {
             Ok(decoded_tensor) => {
-                self.tensors
-                    .insert(ent_path.clone(), (row_id, decoded_tensor));
+                self.tensors.push((row_id, decoded_tensor));
             }
             Err(err) => {
                 re_log::warn_once!("Failed to decode decoding tensor at path {ent_path}: {err}");

crates/re_space_view_text_document/src/space_view_class.rs

@@ -177,7 +177,9 @@ impl SpaceViewClass for TextDocumentSpaceView {
                         } else {
                             // TODO(jleibs): better handling for multiple results
                             ui.label(format!(
-                                "Can only show one text document at a time; was given {}.",
+                                "Can only show one text document at a time; was given {}. Update \
+                                the query so that it returns a single text document and create \
+                                additional views for the others.",
                                 text_document.text_entries.len()
                             ));
                         }

docs/snippets/all/views/tensor.py

@@ -6,12 +6,8 @@
 
 rr.init("rerun_example_tensor", spawn=True)
 
-tensor_one = np.random.randint(0, 256, (8, 6, 3, 5), dtype=np.uint8)
-rr.log("tensors/one", rr.Tensor(tensor_one, dim_names=("width", "height", "channel", "batch")))
-tensor_two = np.random.random_sample((10, 20, 30))
-rr.log("tensors/two", rr.Tensor(tensor_two))
-
-# Create a tensor view that displays both tensors (you can switch between them inside the view).
-blueprint = rrb.Blueprint(rrb.TensorView(origin="/tensors", name="Tensors"), collapse_panels=True)
+tensor = np.random.randint(0, 256, (8, 6, 3, 5), dtype=np.uint8)
+rr.log("tensor", rr.Tensor(tensor, dim_names=("width", "height", "channel", "batch")))
 
+blueprint = rrb.Blueprint(rrb.TensorView(origin="tensor", name="Tensor"), collapse_panels=True)
 rr.send_blueprint(blueprint)