Skip to main content

re_viewer_context/cache/
image_histogram_cache.rs

1use crate::cache::filter_blob_removed_events;
2use crate::image_info::StoredBlobCacheKey;
3use crate::{Cache, CacheEntryAccess, ImageInfo};
4use ahash::HashMap;
5use re_byte_size::SizeBytes as _;
6use re_byte_size::{MemUsageTree, MemUsageTreeCapture};
7use re_chunk_store::ChunkStoreEvent;
8use re_entity_db::EntityDb;
9use std::sync::Arc;
10
11/// Per-channel histogram of an 8-bit `RGB` image.
12///
13/// Each channel has 256 bins.
14#[derive(Clone, Debug, re_byte_size::SizeBytes)]
15pub struct Rgb8Histogram {
16    /// One 256-bin histogram per channel (R, G, B).
17    pub bins: [[u64; 256]; 3],
18}
19
20impl Rgb8Histogram {
21    /// Compute the per-channel histogram of an 8-bit `RGB` buffer.
22    pub fn from_rgb8(rgb: &[u8]) -> Self {
23        re_tracing::profile_function!();
24
25        let mut bins_r = [0_u64; 256];
26        let mut bins_g = [0_u64; 256];
27        let mut bins_b = [0_u64; 256];
28
29        let (chunks, _remainder) = rgb.as_chunks::<3>();
30        for &[r, g, b] in chunks {
31            bins_r[r as usize] += 1;
32            bins_g[g as usize] += 1;
33            bins_b[b as usize] += 1;
34        }
35
36        Self {
37            bins: [bins_r, bins_g, bins_b],
38        }
39    }
40}
41
42/// Caches per-channel histograms for 8-bit RGB images, keyed by image content.
43#[derive(Default)]
44pub struct ImageHistogramCache(HashMap<StoredBlobCacheKey, Arc<Rgb8Histogram>>);
45
46impl ImageHistogramCache {
47    /// Get the histogram for the given 8-bit `RGB` image, computing and caching it on first access.
48    ///
49    /// The caller is responsible for only passing in 8-bit RGB images.
50    pub fn entry(&mut self, image: &ImageInfo) -> Arc<Rgb8Histogram> {
51        self.0
52            .entry(image.buffer_content_hash)
53            .or_insert_with(|| Arc::new(Rgb8Histogram::from_rgb8(&image.buffer)))
54            .clone()
55    }
56}
57
58impl CacheEntryAccess<ImageInfo, Arc<Rgb8Histogram>> for ImageHistogramCache {
59    fn read(&self, image: &ImageInfo) -> Option<Arc<Rgb8Histogram>> {
60        self.0.get(&image.buffer_content_hash).cloned()
61    }
62
63    fn compute(&mut self, image: &ImageInfo) -> Arc<Rgb8Histogram> {
64        self.entry(image)
65    }
66}
67
68impl Cache for ImageHistogramCache {
69    fn name(&self) -> &'static str {
70        "ImageHistogramCache"
71    }
72
73    fn purge_memory(&mut self) {
74        // [[u64; 256]; 3] ≈ 6 KiB per cached image — small enough that we
75        // leave it to store-event invalidation rather than periodic purging.
76    }
77
78    fn on_store_events(&mut self, events: &[&ChunkStoreEvent], _entity_db: &EntityDb) {
79        let removed = filter_blob_removed_events(events);
80        if removed.is_empty() {
81            return;
82        }
83        self.0.retain(|key, _| !removed.contains(key));
84    }
85}
86
87impl MemUsageTreeCapture for ImageHistogramCache {
88    fn capture_mem_usage_tree(&self) -> MemUsageTree {
89        MemUsageTree::Bytes(self.0.total_size_bytes())
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn empty_buffer_yields_zero_bins() {
99        let hist = Rgb8Histogram::from_rgb8(&[]);
100        for channel in &hist.bins {
101            assert!(channel.iter().all(|&c| c == 0));
102        }
103    }
104
105    #[test]
106    fn single_pixel_increments_one_bin_per_channel() {
107        let hist = Rgb8Histogram::from_rgb8(&[10, 20, 30]);
108        assert_eq!(hist.bins[0][10], 1);
109        assert_eq!(hist.bins[1][20], 1);
110        assert_eq!(hist.bins[2][30], 1);
111        let total: u64 = hist.bins.iter().flatten().sum();
112        assert_eq!(total, 3);
113    }
114
115    #[test]
116    fn trailing_bytes_are_ignored() {
117        // 4 trailing bytes; only the first complete pixel should contribute.
118        let hist = Rgb8Histogram::from_rgb8(&[1, 2, 3, 99]);
119        assert_eq!(hist.bins[0][1], 1);
120        assert_eq!(hist.bins[1][2], 1);
121        assert_eq!(hist.bins[2][3], 1);
122        assert_eq!(hist.bins[0][99], 0);
123        assert_eq!(hist.bins[1][99], 0);
124        assert_eq!(hist.bins[2][99], 0);
125    }
126
127    #[test]
128    fn many_pixels_count_correctly() {
129        // 100 pixels, all (5, 6, 7).
130        let buffer: Vec<u8> = (0..100).flat_map(|_| [5, 6, 7]).collect();
131        let hist = Rgb8Histogram::from_rgb8(&buffer);
132        assert_eq!(hist.bins[0][5], 100);
133        assert_eq!(hist.bins[1][6], 100);
134        assert_eq!(hist.bins[2][7], 100);
135        // No other bins should be set.
136        for channel in 0..3 {
137            for bin in 0..256 {
138                let expected = match (channel, bin) {
139                    (0, 5) | (1, 6) | (2, 7) => 100,
140                    _ => 0,
141                };
142                assert_eq!(
143                    hist.bins[channel][bin], expected,
144                    "channel {channel} bin {bin}"
145                );
146            }
147        }
148    }
149}