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
use re_log_types::component_types::{Tensor, TensorDimension, TensorId, TensorTrait};
#[derive(thiserror::Error, Clone, Debug)]
pub enum TensorDecodeError {
// TODO(jleibs): It would be nice to just transparently wrap
// `image::ImageError` and `tensor::TensorImageError` but neither implements
// `Clone`, which we need if we want to cache the Result.
#[error("Failed to decode bytes as tensor: {0}")]
CouldNotDecode(String),
#[error("Failed to interpret image as tensor: {0}")]
InvalidImage(String),
#[error("The encoded tensor did not match its metadata {expected:?} != {found:?}")]
InvalidMetaData {
expected: Vec<TensorDimension>,
found: Vec<TensorDimension>,
},
}
#[derive(Clone)]
struct DecodedTensor {
/// Cached `Result` from decoding the `Tensor`
tensor: Result<Tensor, TensorDecodeError>,
/// Total memory used by this `Tensor`.
memory_used: u64,
/// Which [`DecodeCache::generation`] was this `Tensor` last used?
last_use_generation: u64,
}
/// A cache of decoded [`Tensor`] entities, indexed by `TensorId`.
#[derive(Default)]
pub struct DecodeCache {
images: nohash_hasher::IntMap<TensorId, DecodedTensor>,
memory_used: u64,
generation: u64,
}
#[allow(clippy::map_err_ignore)]
impl DecodeCache {
/// Decode a [`Tensor`] if necessary and cache the result.
///
/// This is a no-op for Tensors that are not compressed.
///
/// Currently supports JPEG encoded tensors.
pub fn try_decode_tensor_if_necessary(
&mut self,
maybe_encoded_tensor: Tensor,
) -> Result<Tensor, TensorDecodeError> {
crate::profile_function!();
match &maybe_encoded_tensor.data {
re_log_types::component_types::TensorData::JPEG(buf) => {
let lookup = self
.images
.entry(maybe_encoded_tensor.id())
.or_insert_with(|| {
use image::io::Reader as ImageReader;
let mut reader = ImageReader::new(std::io::Cursor::new(buf.0.as_slice()));
reader.set_format(image::ImageFormat::Jpeg);
let img = {
crate::profile_scope!("decode_jpeg");
reader.decode()
};
let tensor = match img {
Ok(img) => match Tensor::from_image(img) {
Ok(tensor) => {
if tensor.shape() == maybe_encoded_tensor.shape() {
Ok(tensor)
} else {
Err(TensorDecodeError::InvalidMetaData {
expected: maybe_encoded_tensor.shape().into(),
found: tensor.shape().into(),
})
}
}
Err(err) => Err(TensorDecodeError::InvalidImage(err.to_string())),
},
Err(err) => Err(TensorDecodeError::CouldNotDecode(err.to_string())),
};
let memory_used = match &tensor {
Ok(tensor) => tensor.size_in_bytes() as u64,
Err(_) => 0,
};
self.memory_used += memory_used;
let last_use_generation = 0;
DecodedTensor {
tensor,
memory_used,
last_use_generation,
}
});
lookup.last_use_generation = self.generation;
lookup.tensor.clone()
}
_ => Ok(maybe_encoded_tensor),
}
}
/// Call once per frame to (potentially) flush the cache.
pub fn begin_frame(&mut self, max_memory_use: u64) {
// TODO(jleibs): a more incremental purging mechanism, maybe switching to an LRU Cache
// would likely improve the behavior.
if self.memory_used > max_memory_use {
self.purge_memory();
}
self.generation += 1;
}
/// Attempt to free up memory.
pub fn purge_memory(&mut self) {
crate::profile_function!();
// Very aggressively flush everything not used in this frame
let before = self.memory_used;
self.images.retain(|_, ci| {
let retain = ci.last_use_generation == self.generation;
if !retain {
self.memory_used -= ci.memory_used;
}
retain
});
re_log::debug!(
"Flushed tensor decode cache. Before: {:.2} GB. After: {:.2} GB",
before as f64 / 1e9,
self.memory_used as f64 / 1e9,
);
}
}