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 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 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 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 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}