iris/tracking/
subtractor.rs

1use crate::error::Result;
2use crate::image::Image;
3use burn::tensor::backend::Backend;
4
5/// Background subtractor pipeline.
6pub struct BackgroundSubtractor<B: Backend> {
7    pub learning_rate: f32,
8    pub threshold: f32,
9    background: Option<Image<B>>,
10}
11
12impl<B: Backend> BackgroundSubtractor<B> {
13    /// Creates a new background subtractor.
14    #[must_use]
15    pub fn new(learning_rate: f32, threshold: f32) -> Self {
16        Self {
17            learning_rate,
18            threshold,
19            background: None,
20        }
21    }
22
23    /// Processes a new frame, returning the foreground binary mask.
24    pub fn apply(&mut self, frame: &Image<B>) -> Result<Image<B>> {
25        let frame_gray = frame.grayscale()?;
26
27        let bg = if let Some(bg_img) = &self.background {
28            // Update running background model: B_t = (1 - alpha)*B_{t-1} + alpha*I_t
29            let updated = bg_img
30                .tensor
31                .clone()
32                .mul_scalar(1.0 - self.learning_rate)
33                .add(frame_gray.tensor.clone().mul_scalar(self.learning_rate));
34            let bg_new = Image::new(updated);
35            self.background = Some(bg_new.clone());
36            bg_new
37        } else {
38            self.background = Some(frame_gray.clone());
39            frame_gray.clone()
40        };
41
42        // Foreground mask: F = |Frame - Background| > threshold
43        let diff = frame_gray.absdiff(&bg)?;
44        let mask = diff.threshold(self.threshold, 1.0, crate::threshold::ThresholdType::Binary)?;
45        Ok(mask)
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use crate::test_helpers::{TestBackend, test_device};
53    use burn::tensor::{Tensor, TensorData};
54
55    #[test]
56    fn test_background_subtractor() {
57        let device = test_device();
58        let flat_data1 = vec![0.5f32; 3 * 8 * 8];
59        let flat_data2 = vec![0.6f32; 3 * 8 * 8];
60
61        let img1 = Image::new(Tensor::<TestBackend, 3>::from_data(
62            TensorData::new(flat_data1, [3, 8, 8]),
63            &device,
64        ));
65        let img2 = Image::new(Tensor::<TestBackend, 3>::from_data(
66            TensorData::new(flat_data2, [3, 8, 8]),
67            &device,
68        ));
69
70        let mut bs = BackgroundSubtractor::new(0.1, 0.05);
71        let mask1 = bs.apply(&img1).unwrap();
72        assert_eq!(mask1.shape(), [1, 8, 8]);
73
74        let mask2 = bs.apply(&img2).unwrap();
75        assert_eq!(mask2.shape(), [1, 8, 8]);
76    }
77}
iris/tracking/subtractor.rs

iris/tracking/
subtractor.rs
