Skip to main content

re_data_ui/
tensor_ui.rs

1use itertools::Itertools as _;
2use re_log_types::EntityPath;
3use re_log_types::hash::Hash64;
4use re_sdk_types::datatypes::TensorData;
5use re_sdk_types::{ComponentDescriptor, RowId};
6use re_ui::UiExt as _;
7use re_viewer_context::{StoreViewContext, TensorStats, TensorStatsCache, UiLayout};
8
9use super::EntityDataUi;
10
11fn format_tensor_shape_single_line(tensor: &TensorData) -> String {
12    let has_names = tensor
13        .shape
14        .iter()
15        .enumerate()
16        .any(|(dim_idx, _)| tensor.dim_name(dim_idx).is_some());
17    tensor
18        .shape
19        .iter()
20        .enumerate()
21        .map(|(dim_idx, dim_len)| {
22            format!(
23                "{}{}",
24                dim_len,
25                if let Some(name) = tensor.dim_name(dim_idx) {
26                    format!(" ({name})")
27                } else {
28                    String::new()
29                }
30            )
31        })
32        .join(if has_names { " × " } else { "×" })
33}
34
35impl EntityDataUi for re_sdk_types::components::TensorData {
36    fn entity_data_ui(
37        &self,
38        ctx: &StoreViewContext<'_>,
39        ui: &mut egui::Ui,
40        ui_layout: UiLayout,
41        _entity_path: &EntityPath,
42        _component_descriptor: &ComponentDescriptor,
43        row_id: Option<RowId>,
44    ) {
45        re_tracing::profile_function!();
46        // RowId is enough for cache keying the tensor stats right now since you can't have more than one per row.
47        let tensor_cache_key = row_id.map_or(Hash64::ZERO, Hash64::hash);
48        tensor_ui(ctx, ui, ui_layout, tensor_cache_key, &self.0);
49    }
50}
51
52pub fn tensor_ui(
53    ctx: &StoreViewContext<'_>,
54    ui: &mut egui::Ui,
55    ui_layout: UiLayout,
56    tensor_cache_key: Hash64,
57    tensor: &TensorData,
58) {
59    // See if we can convert the tensor to a GPU texture.
60    // Even if not, we will show info about the tensor.
61    let tensor_stats = ctx.memoizer(|c: &mut TensorStatsCache| c.entry(tensor_cache_key, tensor));
62
63    if ui_layout.is_single_line() {
64        ui.horizontal(|ui| {
65            let text = format!(
66                "{}, {}",
67                tensor.dtype(),
68                format_tensor_shape_single_line(tensor)
69            );
70            ui_layout.label(ui, text).on_hover_ui(|ui| {
71                tensor_summary_ui(ui, tensor, &tensor_stats);
72            });
73        });
74    } else {
75        ui.vertical(|ui| {
76            ui.set_min_width(100.0);
77            tensor_summary_ui(ui, tensor, &tensor_stats);
78        });
79    }
80}
81
82pub fn tensor_summary_ui_grid_contents(
83    ui: &mut egui::Ui,
84    tensor: &TensorData,
85    tensor_stats: &TensorStats,
86) {
87    let TensorData { shape, names, .. } = tensor;
88
89    ui.grid_left_hand_label("Data type")
90        .on_hover_text("Data type used for all individual elements within the tensor");
91    ui.label(tensor.dtype().to_string());
92    ui.end_row();
93
94    ui.grid_left_hand_label("Shape")
95        .on_hover_text("Extent of every dimension");
96    ui.vertical(|ui| {
97        // For unnamed tensor dimension more than a single line usually doesn't make sense!
98        // But what if some are named and some are not?
99        // -> If more than 1 is named, make it a column!
100        if let Some(names) = names {
101            for (name, size) in itertools::izip!(names, shape) {
102                ui.label(format!("{name}={size}"));
103            }
104        } else {
105            ui.label(format_tensor_shape_single_line(tensor));
106        }
107    });
108    ui.end_row();
109
110    let TensorStats {
111        range,
112        finite_range,
113    } = tensor_stats;
114
115    if let Some((min, max)) = range {
116        ui.label("Data range")
117            .on_hover_text("All values of the tensor range within these bounds");
118        ui.monospace(format!(
119            "[{} - {}]",
120            re_format::format_f64(*min),
121            re_format::format_f64(*max)
122        ));
123        ui.end_row();
124    }
125    // Show finite range only if it is different from the actual range.
126    if range != &Some(*finite_range) {
127        ui.label("Finite data range").on_hover_text(
128            "The finite values (ignoring all NaN & -Inf/+Inf) of the tensor range within these bounds"
129        );
130        let (min, max) = finite_range;
131        ui.monospace(format!(
132            "[{} - {}]",
133            re_format::format_f64(*min),
134            re_format::format_f64(*max)
135        ));
136        ui.end_row();
137    }
138}
139
140pub fn tensor_summary_ui(ui: &mut egui::Ui, tensor: &TensorData, tensor_stats: &TensorStats) {
141    egui::Grid::new("tensor_summary_ui")
142        .num_columns(2)
143        .show(ui, |ui| {
144            tensor_summary_ui_grid_contents(ui, tensor, tensor_stats);
145        });
146}