1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

use itertools::Itertools as _;

use re_chunk_store::RowId;
use re_log_types::EntityPath;
use re_types::datatypes::TensorData;
use re_ui::UiExt as _;
use re_viewer_context::{TensorStats, TensorStatsCache, UiLayout, ViewerContext};

use super::EntityDataUi;

fn format_tensor_shape_single_line(tensor: &TensorData) -> String {
    let has_names = tensor
        .shape
        .iter()
        .enumerate()
        .any(|(dim_idx, _)| tensor.dim_name(dim_idx).is_some());
    tensor
        .shape
        .iter()
        .enumerate()
        .map(|(dim_idx, dim_len)| {
            format!(
                "{}{}",
                dim_len,
                if let Some(name) = tensor.dim_name(dim_idx) {
                    format!(" ({name})")
                } else {
                    String::new()
                }
            )
        })
        .join(if has_names { " × " } else { "×" })
}

impl EntityDataUi for re_types::components::TensorData {
    fn entity_data_ui(
        &self,
        ctx: &ViewerContext<'_>,
        ui: &mut egui::Ui,
        ui_layout: UiLayout,
        _entity_path: &EntityPath,
        row_id: Option<re_chunk_store::RowId>,
        _query: &re_chunk_store::LatestAtQuery,
        _db: &re_entity_db::EntityDb,
    ) {
        re_tracing::profile_function!();

        let tensor_data_row_id = row_id.unwrap_or(RowId::ZERO);
        tensor_ui(ctx, ui, ui_layout, tensor_data_row_id, &self.0);
    }
}

pub fn tensor_ui(
    ctx: &ViewerContext<'_>,
    ui: &mut egui::Ui,
    ui_layout: UiLayout,
    tensor_data_row_id: RowId,
    tensor: &TensorData,
) {
    // See if we can convert the tensor to a GPU texture.
    // Even if not, we will show info about the tensor.
    let tensor_stats = ctx
        .cache
        .entry(|c: &mut TensorStatsCache| c.entry(tensor_data_row_id, tensor));

    if ui_layout.is_single_line() {
        ui.horizontal(|ui| {
            let text = format!(
                "{}, {}",
                tensor.dtype(),
                format_tensor_shape_single_line(tensor)
            );
            ui_layout.label(ui, text).on_hover_ui(|ui| {
                tensor_summary_ui(ui, tensor, &tensor_stats);
            });
        });
    } else {
        ui.vertical(|ui| {
            ui.set_min_width(100.0);
            tensor_summary_ui(ui, tensor, &tensor_stats);
        });
    }
}

pub fn tensor_summary_ui_grid_contents(
    ui: &mut egui::Ui,
    tensor: &TensorData,
    tensor_stats: &TensorStats,
) {
    let TensorData { shape, names, .. } = tensor;

    ui.grid_left_hand_label("Data type")
        .on_hover_text("Data type used for all individual elements within the tensor");
    ui.label(tensor.dtype().to_string());
    ui.end_row();

    ui.grid_left_hand_label("Shape")
        .on_hover_text("Extent of every dimension");
    ui.vertical(|ui| {
        // For unnamed tensor dimension more than a single line usually doesn't make sense!
        // But what if some are named and some are not?
        // -> If more than 1 is named, make it a column!
        if let Some(names) = names {
            for (name, size) in itertools::izip!(names, shape) {
                ui.label(format!("{name}={size}"));
            }
        } else {
            ui.label(format_tensor_shape_single_line(tensor));
        }
    });
    ui.end_row();

    let TensorStats {
        range,
        finite_range,
    } = tensor_stats;

    if let Some((min, max)) = range {
        ui.label("Data range")
            .on_hover_text("All values of the tensor range within these bounds");
        ui.monospace(format!(
            "[{} - {}]",
            re_format::format_f64(*min),
            re_format::format_f64(*max)
        ));
        ui.end_row();
    }
    // Show finite range only if it is different from the actual range.
    if range != &Some(*finite_range) {
        ui.label("Finite data range").on_hover_text(
            "The finite values (ignoring all NaN & -Inf/+Inf) of the tensor range within these bounds"
        );
        let (min, max) = finite_range;
        ui.monospace(format!(
            "[{} - {}]",
            re_format::format_f64(*min),
            re_format::format_f64(*max)
        ));
        ui.end_row();
    }
}

pub fn tensor_summary_ui(ui: &mut egui::Ui, tensor: &TensorData, tensor_stats: &TensorStats) {
    egui::Grid::new("tensor_summary_ui")
        .num_columns(2)
        .show(ui, |ui| {
            tensor_summary_ui_grid_contents(ui, tensor, tensor_stats);
        });
}