Skip to main content

eventcv_core/representation/
redblue.rs

1use super::{
2    countmask::percentile_linear, frame_len, polarity::polarity_counts_on, EventFrame,
3    EventFrameData, Representation, RepresentationError, RepresentationKind,
4};
5use crate::{accel::Device, EventStream};
6
7/// White-background RGB counts with independent polarity scales and positive ties.
8#[derive(Clone, Copy, Debug)]
9pub struct RedBlue {
10    pct: f64,
11}
12
13impl RedBlue {
14    pub fn new(pct: f64) -> Self {
15        Self { pct }
16    }
17}
18
19impl Default for RedBlue {
20    fn default() -> Self {
21        Self::new(99.0)
22    }
23}
24
25// Repeated float32 additions of 1 stop changing at 2^24, as in the NumPy reference.
26fn reference_count(count: u64) -> f32 {
27    count.min(1 << 24) as f32
28}
29
30fn normalize(counts: &[u64], pct: f64) -> Vec<f32> {
31    let mut nonzero: Vec<f64> = counts
32        .iter()
33        .filter(|&&count| count > 0)
34        .map(|&count| f64::from(reference_count(count)))
35        .collect();
36    let threshold = if nonzero.is_empty() {
37        1.0
38    } else {
39        percentile_linear(&mut nonzero, pct) as f32
40    };
41    counts
42        .iter()
43        .map(|&count| reference_count(count).min(threshold) / threshold)
44        .collect()
45}
46
47impl Representation for RedBlue {
48    type Output = EventFrame;
49
50    fn generate(&self, stream: &EventStream) -> Result<EventFrame, RepresentationError> {
51        self.generate_on(stream, Device::Cpu)
52    }
53
54    fn generate_on(
55        &self,
56        stream: &EventStream,
57        device: Device,
58    ) -> Result<EventFrame, RepresentationError> {
59        if !self.pct.is_finite() || !(0.0..=100.0).contains(&self.pct) {
60            return Err(RepresentationError::InvalidParameter("pct"));
61        }
62        let (width, height, length) = frame_len(stream, 3)?;
63        let plane = width * height;
64        let (_, _, counts) = polarity_counts_on(stream, device)?;
65        let pos = normalize(&counts[..plane], self.pct);
66        let neg = normalize(&counts[plane..], self.pct);
67        let mut data = vec![255; length];
68        for i in 0..plane {
69            let positive = pos[i] >= neg[i];
70            let intensity = if positive { pos[i] } else { neg[i] };
71            let faded = ((1.0 - intensity).clamp(0.0, 1.0) * 255.0) as u8;
72            data[plane + i] = faded;
73            data[if positive { 2 * plane + i } else { i }] = faded;
74        }
75        Ok(EventFrame {
76            data: EventFrameData::U8(data),
77            channels: 3,
78            width,
79            height,
80            kind: RepresentationKind::RedBlue,
81            channel_names: vec!["red".into(), "green".into(), "blue".into()],
82        })
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use crate::viz::{render_frame, Colormap};
90    use ndarray::{array, Array2};
91
92    #[test]
93    fn preserves_rgb_when_rendered() {
94        let stream = EventStream::from_array2(array![[0, 0, 0, 1], [1, 0, 0, 0]], 3, 1, 0.001);
95        let frame = RedBlue::default().generate(&stream).unwrap();
96        for normalize in [false, true] {
97            assert_eq!(
98                render_frame(&frame, Colormap::Turbo, normalize).pixels,
99                vec![255, 0, 0, 0, 0, 255, 255, 255, 255]
100            );
101        }
102        assert_eq!(RepresentationKind::from_tag("redblue"), Some(frame.kind()));
103        assert_eq!(frame.kind().as_str(), "redblue");
104    }
105
106    #[test]
107    fn matches_float32_count_saturation() {
108        assert_eq!(reference_count((1 << 24) + 1), 16_777_216.0);
109        assert_eq!(reference_count(u64::MAX), 16_777_216.0);
110        assert_eq!(normalize(&[1 << 24, u64::MAX], 99.0), vec![1.0, 1.0]);
111    }
112
113    #[test]
114    fn rejects_invalid_coordinates_and_dimensions() {
115        let outside = EventStream::from_array2(array![[2, 0, 0, 1]], 2, 1, 0.001);
116        assert!(matches!(
117            RedBlue::default().generate(&outside),
118            Err(RepresentationError::EventOutOfBounds { .. })
119        ));
120        let oversized = EventStream::from_array2(Array2::zeros((0, 4)), usize::MAX, 2, 0.001);
121        assert_eq!(
122            RedBlue::default().generate(&oversized).unwrap_err(),
123            RepresentationError::SizeOverflow
124        );
125    }
126}