spatialrust_vision/
video.rs

1//! Dense motion, background modeling, tracking, and video-source contracts.
2
3#[cfg(feature = "video-adapters")]
4use std::collections::VecDeque;
5
6#[cfg(any(feature = "video-adapters", test))]
7use spatialrust_image::Image;
8use spatialrust_image::ImageView;
9
10use crate::{BinaryMask, BoundingBox2, Detection, FlowField, VisionError, VisionResult};
11
12/// Integer block-matching controls for dense grayscale optical flow.
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub struct DenseFlowOptions {
15    /// Half-size of the square comparison block.
16    pub block_radius: usize,
17    /// Maximum integer displacement searched in each axis.
18    pub search_radius: usize,
19    /// Minimum improvement over zero motion required for a valid vector.
20    pub minimum_improvement: u64,
21}
22
23impl Default for DenseFlowOptions {
24    fn default() -> Self {
25        Self { block_radius: 2, search_radius: 4, minimum_improvement: 1 }
26    }
27}
28
29/// Computes a dense integer flow field with deterministic local block matching.
30///
31/// Pixels whose full block/search window leaves the image are marked with NaN.
32pub fn dense_flow_block_match(
33    previous: ImageView<'_, u8, 1>,
34    next: ImageView<'_, u8, 1>,
35    options: DenseFlowOptions,
36) -> VisionResult<FlowField> {
37    if (previous.width(), previous.height()) != (next.width(), next.height()) {
38        return Err(VisionError::ShapeMismatch(
39            "dense-flow frames must share dimensions".to_owned(),
40        ));
41    }
42    if options.block_radius == 0 || options.search_radius == 0 {
43        return Err(VisionError::InvalidParameter(
44            "dense-flow block/search radii must be positive".to_owned(),
45        ));
46    }
47    let margin = options.block_radius.saturating_add(options.search_radius);
48    if previous.width() <= margin * 2 || previous.height() <= margin * 2 {
49        return Err(VisionError::InvalidDimensions(
50            "dense-flow frames are too small for the configured search".to_owned(),
51        ));
52    }
53    let mut flow = vec![f32::NAN; previous.width() * previous.height() * 2];
54    for y in margin..previous.height() - margin {
55        for x in margin..previous.width() - margin {
56            let zero_cost = block_cost(previous, next, x, y, 0, 0, options.block_radius);
57            let mut best_cost = zero_cost;
58            let mut best = (0_isize, 0_isize);
59            let search = options.search_radius as isize;
60            for dy in -search..=search {
61                for dx in -search..=search {
62                    let cost = block_cost(previous, next, x, y, dx, dy, options.block_radius);
63                    if cost < best_cost
64                        || (cost == best_cost
65                            && (dy.abs() + dx.abs(), dy, dx)
66                                < (best.1.abs() + best.0.abs(), best.1, best.0))
67                    {
68                        best_cost = cost;
69                        best = (dx, dy);
70                    }
71                }
72            }
73            if zero_cost.saturating_sub(best_cost) >= options.minimum_improvement {
74                let offset = (y * previous.width() + x) * 2;
75                flow[offset] = best.0 as f32;
76                flow[offset + 1] = best.1 as f32;
77            }
78        }
79    }
80    FlowField::try_new(previous.width(), previous.height(), flow)
81}
82
83fn block_cost(
84    previous: ImageView<'_, u8, 1>,
85    next: ImageView<'_, u8, 1>,
86    x: usize,
87    y: usize,
88    dx: isize,
89    dy: isize,
90    radius: usize,
91) -> u64 {
92    let mut cost = 0_u64;
93    let radius = radius as isize;
94    for by in -radius..=radius {
95        for bx in -radius..=radius {
96            let px = (x as isize + bx) as usize;
97            let py = (y as isize + by) as usize;
98            let nx = (x as isize + bx + dx) as usize;
99            let ny = (y as isize + by + dy) as usize;
100            let left = previous.get(px, py).expect("validated flow window")[0];
101            let right = next.get(nx, ny).expect("validated flow window")[0];
102            let difference = i32::from(left) - i32::from(right);
103            cost = cost.saturating_add((difference * difference) as u64);
104        }
105    }
106    cost
107}
108
109/// Adaptive Gaussian background-model controls.
110#[derive(Clone, Copy, Debug, PartialEq)]
111pub struct BackgroundModelOptions {
112    /// Exponential update factor in `(0, 1]`.
113    pub learning_rate: f32,
114    /// Foreground threshold measured in standard deviations.
115    pub sigma_threshold: f32,
116    /// Initial per-pixel variance.
117    pub initial_variance: f32,
118}
119
120impl Default for BackgroundModelOptions {
121    fn default() -> Self {
122        Self { learning_rate: 0.02, sigma_threshold: 3.0, initial_variance: 225.0 }
123    }
124}
125
126/// Stateful single-Gaussian grayscale background model.
127#[derive(Clone, Debug, PartialEq)]
128pub struct AdaptiveBackgroundModel {
129    width: usize,
130    height: usize,
131    mean: Vec<f32>,
132    variance: Vec<f32>,
133    options: BackgroundModelOptions,
134    frames_seen: u64,
135}
136
137impl AdaptiveBackgroundModel {
138    /// Initializes the model from its first frame.
139    pub fn try_new(
140        frame: ImageView<'_, u8, 1>,
141        options: BackgroundModelOptions,
142    ) -> VisionResult<Self> {
143        validate_background_options(options)?;
144        let mean = packed_gray(frame).into_iter().map(f32::from).collect();
145        Ok(Self {
146            width: frame.width(),
147            height: frame.height(),
148            mean,
149            variance: vec![options.initial_variance; frame.width() * frame.height()],
150            options,
151            frames_seen: 1,
152        })
153    }
154
155    /// Segments foreground and updates background statistics in one explicit step.
156    pub fn apply(&mut self, frame: ImageView<'_, u8, 1>) -> VisionResult<BackgroundUpdate> {
157        if (frame.width(), frame.height()) != (self.width, self.height) {
158            return Err(VisionError::ShapeMismatch(
159                "background-model frame dimensions changed".to_owned(),
160            ));
161        }
162        let pixels = packed_gray(frame);
163        let mut mask = Vec::with_capacity(pixels.len());
164        for (index, &pixel) in pixels.iter().enumerate() {
165            let value = f32::from(pixel);
166            let delta = value - self.mean[index];
167            let threshold = self.options.sigma_threshold * self.variance[index].max(1.0).sqrt();
168            let foreground = delta.abs() > threshold;
169            mask.push(u8::from(foreground));
170            let alpha = if foreground {
171                self.options.learning_rate * 0.1
172            } else {
173                self.options.learning_rate
174            };
175            self.mean[index] += alpha * delta;
176            self.variance[index] =
177                ((1.0 - alpha) * self.variance[index] + alpha * delta * delta).max(1.0);
178        }
179        self.frames_seen = self.frames_seen.saturating_add(1);
180        let mask = BinaryMask::try_new(self.width, self.height, mask)?;
181        let foreground_ratio = mask.area() as f32 / (self.width * self.height).max(1) as f32;
182        Ok(BackgroundUpdate { mask, foreground_ratio, frame_index: self.frames_seen - 1 })
183    }
184
185    /// Returns the number of frames incorporated into the model.
186    #[must_use]
187    pub const fn frames_seen(&self) -> u64 {
188        self.frames_seen
189    }
190}
191
192/// Output of one background-model update.
193#[derive(Clone, Debug, PartialEq)]
194pub struct BackgroundUpdate {
195    /// Binary foreground segmentation.
196    pub mask: BinaryMask,
197    /// Foreground area divided by image area.
198    pub foreground_ratio: f32,
199    /// Zero-based sequence index of the processed frame.
200    pub frame_index: u64,
201}
202
203fn validate_background_options(options: BackgroundModelOptions) -> VisionResult<()> {
204    if !options.learning_rate.is_finite()
205        || !(0.0..=1.0).contains(&options.learning_rate)
206        || options.learning_rate == 0.0
207        || !options.sigma_threshold.is_finite()
208        || options.sigma_threshold <= 0.0
209        || !options.initial_variance.is_finite()
210        || options.initial_variance <= 0.0
211    {
212        return Err(VisionError::InvalidParameter(
213            "invalid background-model learning/threshold/variance options".to_owned(),
214        ));
215    }
216    Ok(())
217}
218
219fn packed_gray(frame: ImageView<'_, u8, 1>) -> Vec<u8> {
220    let mut data = Vec::with_capacity(frame.width() * frame.height());
221    for y in 0..frame.height() {
222        data.extend_from_slice(frame.row(y).expect("frame row in bounds"));
223    }
224    data
225}
226
227/// Track lifecycle state.
228#[derive(Clone, Copy, Debug, PartialEq, Eq)]
229pub enum TrackState {
230    /// Track has not accumulated the configured hit count.
231    Tentative,
232    /// Track has accumulated enough consecutive associations.
233    Confirmed,
234}
235
236/// One deterministic detection track.
237#[derive(Clone, Copy, Debug, PartialEq)]
238pub struct ObjectTrack {
239    /// Monotonic track identifier.
240    pub id: u64,
241    /// Most recent bounding box.
242    pub bbox: BoundingBox2,
243    /// Model-defined class identifier.
244    pub class_id: i64,
245    /// Most recent detection score.
246    pub score: f32,
247    /// Frames since track creation.
248    pub age: u32,
249    /// Total successful associations.
250    pub hits: u32,
251    /// Consecutive frames without an association.
252    pub missed: u32,
253    /// Tentative/confirmed state.
254    pub state: TrackState,
255}
256
257/// IoU tracker controls.
258#[derive(Clone, Copy, Debug, PartialEq)]
259pub struct MultiObjectTrackerOptions {
260    /// Minimum IoU for same-class association.
261    pub iou_threshold: f32,
262    /// Consecutive misses retained before removal.
263    pub max_missed: u32,
264    /// Hits required to confirm a track.
265    pub min_confirmed_hits: u32,
266}
267
268impl Default for MultiObjectTrackerOptions {
269    fn default() -> Self {
270        Self { iou_threshold: 0.3, max_missed: 3, min_confirmed_hits: 2 }
271    }
272}
273
274/// Deterministic same-class greedy IoU multi-object tracker.
275#[derive(Clone, Debug, PartialEq)]
276pub struct MultiObjectTracker {
277    options: MultiObjectTrackerOptions,
278    next_id: u64,
279    tracks: Vec<ObjectTrack>,
280}
281
282impl MultiObjectTracker {
283    /// Creates an empty tracker.
284    pub fn try_new(options: MultiObjectTrackerOptions) -> VisionResult<Self> {
285        if !options.iou_threshold.is_finite()
286            || !(0.0..=1.0).contains(&options.iou_threshold)
287            || options.min_confirmed_hits == 0
288        {
289            return Err(VisionError::InvalidParameter("invalid tracker options".to_owned()));
290        }
291        Ok(Self { options, next_id: 1, tracks: Vec::new() })
292    }
293
294    /// Associates one detection frame and returns current tracks ordered by id.
295    pub fn update(&mut self, detections: &[Detection]) -> VisionResult<&[ObjectTrack]> {
296        if detections
297            .iter()
298            .any(|detection| !detection.bbox.is_valid() || !detection.score.is_finite())
299        {
300            return Err(VisionError::InvalidParameter(
301                "tracker detections must contain valid boxes and finite scores".to_owned(),
302            ));
303        }
304        for track in &mut self.tracks {
305            track.age = track.age.saturating_add(1);
306            track.missed = track.missed.saturating_add(1);
307        }
308        let mut used = vec![false; detections.len()];
309        for track in &mut self.tracks {
310            let best = detections
311                .iter()
312                .enumerate()
313                .filter(|(index, detection)| !used[*index] && detection.class_id == track.class_id)
314                .map(|(index, detection)| (index, track.bbox.iou(detection.bbox)))
315                .filter(|(_, iou)| *iou >= self.options.iou_threshold)
316                .max_by(|left, right| {
317                    left.1.total_cmp(&right.1).then_with(|| right.0.cmp(&left.0))
318                });
319            if let Some((index, _)) = best {
320                used[index] = true;
321                let detection = detections[index];
322                track.bbox = detection.bbox;
323                track.score = detection.score;
324                track.hits = track.hits.saturating_add(1);
325                track.missed = 0;
326                if track.hits >= self.options.min_confirmed_hits {
327                    track.state = TrackState::Confirmed;
328                }
329            }
330        }
331        for (index, &detection) in detections.iter().enumerate() {
332            if !used[index] {
333                let hits = 1;
334                self.tracks.push(ObjectTrack {
335                    id: self.next_id,
336                    bbox: detection.bbox,
337                    class_id: detection.class_id,
338                    score: detection.score,
339                    age: 1,
340                    hits,
341                    missed: 0,
342                    state: if hits >= self.options.min_confirmed_hits {
343                        TrackState::Confirmed
344                    } else {
345                        TrackState::Tentative
346                    },
347                });
348                self.next_id = self.next_id.saturating_add(1);
349            }
350        }
351        self.tracks.retain(|track| track.missed <= self.options.max_missed);
352        self.tracks.sort_by_key(|track| track.id);
353        Ok(&self.tracks)
354    }
355
356    /// Borrows current tracks ordered by id.
357    #[must_use]
358    pub fn tracks(&self) -> &[ObjectTrack] {
359        &self.tracks
360    }
361}
362
363/// Timestamped owned video frame.
364#[cfg(feature = "video-adapters")]
365#[derive(Clone, Debug, PartialEq)]
366pub struct VideoFrame<T, const CHANNELS: usize> {
367    /// Monotonic source sequence.
368    pub sequence: u64,
369    /// Source timestamp in nanoseconds.
370    pub timestamp_ns: i64,
371    /// Owned typed image; adapters do not hide host/device copies.
372    pub image: Image<T, CHANNELS>,
373}
374
375/// Pull-based video source implemented by optional codec/camera adapters.
376#[cfg(feature = "video-adapters")]
377pub trait VideoFrameSource<T, const CHANNELS: usize> {
378    /// Returns the next frame, or `None` at end of stream.
379    fn next_frame(&mut self) -> VisionResult<Option<VideoFrame<T, CHANNELS>>>;
380}
381
382/// Deterministic in-memory source used for adapter conformance and replay.
383#[cfg(feature = "video-adapters")]
384#[derive(Clone, Debug, PartialEq)]
385pub struct MemoryVideoSource<T, const CHANNELS: usize> {
386    frames: VecDeque<VideoFrame<T, CHANNELS>>,
387}
388
389#[cfg(feature = "video-adapters")]
390impl<T, const CHANNELS: usize> MemoryVideoSource<T, CHANNELS> {
391    /// Creates a source and validates strictly increasing sequence/timestamps.
392    pub fn try_new(frames: Vec<VideoFrame<T, CHANNELS>>) -> VisionResult<Self> {
393        if frames.windows(2).any(|pair| {
394            pair[0].sequence >= pair[1].sequence
395                || pair[0].timestamp_ns >= pair[1].timestamp_ns
396                || (pair[0].image.width(), pair[0].image.height())
397                    != (pair[1].image.width(), pair[1].image.height())
398        }) {
399            return Err(VisionError::InvalidParameter(
400                "video frames need increasing sequence/time and stable dimensions".to_owned(),
401            ));
402        }
403        Ok(Self { frames: frames.into() })
404    }
405}
406
407#[cfg(feature = "video-adapters")]
408impl<T, const CHANNELS: usize> VideoFrameSource<T, CHANNELS> for MemoryVideoSource<T, CHANNELS> {
409    fn next_frame(&mut self) -> VisionResult<Option<VideoFrame<T, CHANNELS>>> {
410        Ok(self.frames.pop_front())
411    }
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417
418    fn translated_texture(
419        width: usize,
420        height: usize,
421        dx: usize,
422        dy: usize,
423    ) -> (Image<u8, 1>, Image<u8, 1>) {
424        let mut previous = vec![0_u8; width * height];
425        for y in 0..height {
426            for x in 0..width {
427                previous[y * width + x] = ((x * 17 + y * 29 + x * y * 3) % 251) as u8;
428            }
429        }
430        let mut next = vec![0_u8; width * height];
431        for y in 0..height - dy {
432            for x in 0..width - dx {
433                next[(y + dy) * width + x + dx] = previous[y * width + x];
434            }
435        }
436        (
437            Image::try_new(width, height, previous).unwrap(),
438            Image::try_new(width, height, next).unwrap(),
439        )
440    }
441
442    #[test]
443    fn dense_flow_recovers_integer_translation() {
444        let (previous, next) = translated_texture(40, 32, 2, 1);
445        let flow =
446            dense_flow_block_match(previous.view(), next.view(), DenseFlowOptions::default())
447                .unwrap();
448        let center = flow.image().get(20, 16).unwrap();
449        assert_eq!(center, &[2.0, 1.0]);
450        assert!(flow.image().get(0, 0).unwrap().iter().all(|value| value.is_nan()));
451    }
452
453    #[test]
454    fn background_model_detects_new_foreground() {
455        let base = Image::<u8, 1>::try_new(8, 8, vec![20; 64]).unwrap();
456        let mut model =
457            AdaptiveBackgroundModel::try_new(base.view(), BackgroundModelOptions::default())
458                .unwrap();
459        let mut changed = vec![20_u8; 64];
460        for y in 2..6 {
461            for x in 3..5 {
462                changed[y * 8 + x] = 200;
463            }
464        }
465        let frame = Image::<u8, 1>::try_new(8, 8, changed).unwrap();
466        let update = model.apply(frame.view()).unwrap();
467        assert_eq!(update.mask.area(), 8);
468        assert_eq!(update.frame_index, 1);
469    }
470
471    #[test]
472    fn tracker_preserves_id_and_expires_misses() {
473        let mut tracker = MultiObjectTracker::try_new(MultiObjectTrackerOptions {
474            max_missed: 1,
475            ..MultiObjectTrackerOptions::default()
476        })
477        .unwrap();
478        let detection = |x: f32| Detection {
479            bbox: BoundingBox2::try_new(x, 0.0, x + 10.0, 10.0).unwrap(),
480            score: 0.9,
481            class_id: 1,
482        };
483        assert_eq!(tracker.update(&[detection(0.0)]).unwrap()[0].id, 1);
484        let tracks = tracker.update(&[detection(1.0)]).unwrap();
485        assert_eq!(tracks[0].id, 1);
486        assert_eq!(tracks[0].state, TrackState::Confirmed);
487        assert_eq!(tracker.update(&[]).unwrap().len(), 1);
488        assert!(tracker.update(&[]).unwrap().is_empty());
489    }
490
491    #[cfg(feature = "video-adapters")]
492    #[test]
493    fn memory_source_preserves_sequence_and_end_of_stream() {
494        let frames = (0..3)
495            .map(|sequence| VideoFrame {
496                sequence,
497                timestamp_ns: sequence as i64 * 10,
498                image: Image::<u8, 1>::try_new(2, 2, vec![sequence as u8; 4]).unwrap(),
499            })
500            .collect();
501        let mut source = MemoryVideoSource::try_new(frames).unwrap();
502        for sequence in 0..3 {
503            assert_eq!(source.next_frame().unwrap().unwrap().sequence, sequence);
504        }
505        assert!(source.next_frame().unwrap().is_none());
506    }
507}
spatialrust_vision/video.rs

spatialrust_vision/
video.rs
