tensorboard_rs/
summary.rs

1#![allow(clippy::too_many_arguments)]
2use tensorboard_proto::summary::{Summary, Summary_Value, Summary_Image, SummaryMetadata, SummaryMetadata_PluginData, HistogramProto};
3use tensorboard_proto::layout::{Layout, Category};
4use protobuf::RepeatedField;
5
6use image::{RgbImage, DynamicImage, ImageOutputFormat};
7
8pub fn scalar(name: &str, scalar_value: f32) -> Summary {
9
10    let mut value = Summary_Value::new();
11    value.set_tag(name.to_string());
12    value.set_simple_value(scalar_value);
13
14    let values = RepeatedField::from(vec![value]);
15    let mut summary = Summary::new();
16    summary.set_value(values);
17
18    summary
19}
20
21pub fn histogram_raw(name: &str,
22                     min: f64, max: f64,
23                     num: f64,
24                     sum: f64, sum_squares: f64,
25                     bucket_limits: &[f64],
26                     bucket_counts: &[f64],
27) -> Summary {
28    let mut hist = HistogramProto::new();
29    hist.set_min(min);
30    hist.set_max(max);
31    hist.set_num(num);
32    hist.set_sum(sum);
33    hist.set_sum_squares(sum_squares);
34    hist.set_bucket_limit(bucket_limits.to_vec());
35    hist.set_bucket(bucket_counts.to_vec());
36    
37    let mut value = Summary_Value::new();
38    value.set_tag(name.to_string());
39    value.set_histo(hist);
40
41    let values = RepeatedField::from(vec![value]);
42    let mut summary = Summary::new();
43    summary.set_value(values);
44
45    summary
46}
47
48/// dim is in CHW
49pub fn image(tag: &str, data: &[u8], dim: &[usize]) -> Summary {
50    if dim.len() != 3 {
51        panic!("format:CHW");
52    }
53    if dim[0] != 3 {
54        panic!("needs rgb");
55    }
56    if data.len() != dim[0]*dim[1]*dim[2] {
57        panic!("length of data should matches with dim.");
58    }
59    
60    let mut img = RgbImage::new(dim[1] as u32, dim[2] as u32);
61    img.clone_from_slice(data);
62    let dimg = DynamicImage::ImageRgb8(img);
63    let mut output_buf = Vec::<u8>::new();
64    dimg.write_to(&mut output_buf, ImageOutputFormat::Png).expect("");
65
66    let mut output_image = Summary_Image::new();
67    output_image.set_height(dim[1] as i32);
68    output_image.set_width(dim[2] as i32);
69    output_image.set_colorspace(3);
70    output_image.set_encoded_image_string(output_buf);
71    let mut value = Summary_Value::new();
72    value.set_tag(tag.to_string());
73    value.set_image(output_image);
74    let values = RepeatedField::from(vec![value]);
75    let mut summary = Summary::new();
76    summary.set_value(values);
77
78    summary
79}
80
81pub fn custom_scalars(_layout: f32) {
82    let mut layout = Layout::new();
83    let value = Category::new();
84    let values = RepeatedField::from(vec![value]);
85    layout.set_category(values);
86
87    let mut plugin_data = SummaryMetadata_PluginData::new();
88    plugin_data.set_plugin_name("custom_scalars".to_string());
89    let mut smd = SummaryMetadata::new();
90    smd.set_plugin_data(plugin_data);
91
92    
93}