environment/training-source/source-snapshot/src/source.rs

//! Where frames come from.
//!
//! The passthrough camera is not wired up yet, so this exists mainly as the
//! seam it will slot into: the XR loop pulls from a [`FrameSource`] and does
//! not care whether the pixels came from a camera or were made up.
//!
//! Until then [`TestPattern`] provides something with real spatial
//! structure. That matters more than it sounds — a flat or noisy image
//! makes every patch feature statistically identical, so the PCA colouring
//! would look plausible while proving nothing. A scene with distinct
//! regions shows immediately whether features are tracking content.

/// A source of square RGB8 frames at the encoder's input resolution.
pub trait FrameSource {
    /// Side length in pixels; must match `Config::image_size`.
    fn size(&self) -> usize;

    /// The next frame as interleaved RGB8, `[size, size, 3]`, or `None`
    /// when nothing new is available.
    fn next_frame(&mut self) -> Option<&[u8]>;
}

/// A moving synthetic scene: coloured discs drifting over a gradient, with
/// a checkerboard patch for high-frequency contrast.
///
/// Deliberately built from a few large, distinctly-coloured regions, since
/// that is what DINO features separate well and therefore what makes the
/// PCA view legible while bringing the pipeline up.
pub struct TestPattern {
    size: usize,
    buf: Vec<u8>,
    frame: u32,
}

impl TestPattern {
    pub fn new(size: usize) -> Self {
        Self {
            size,
            buf: vec![0; size * size * 3],
            frame: 0,
        }
    }

    fn render(&mut self) {
        let size = self.size;
        let t = self.frame as f32 * 0.02;
        // Three discs on circular paths, each a saturated primary so the
        // top principal components have something unambiguous to separate.
        let discs = [
            (0.30 + 0.18 * t.cos(), 0.30 + 0.18 * t.sin(), 0.16, [230u8, 60, 50]),
            (0.70 + 0.15 * (t * 0.7 + 2.0).cos(), 0.35 + 0.15 * (t * 0.7).sin(), 0.13, [60, 200, 90]),
            (0.50 + 0.20 * (t * 0.5 + 4.0).sin(), 0.72 + 0.10 * (t * 0.9).cos(), 0.15, [70, 110, 235]),
        ];

        for y in 0..size {
            let v = y as f32 / size as f32;
            for x in 0..size {
                let u = x as f32 / size as f32;

                // Background: a slow vertical gradient.
                let mut rgb = [
                    (40.0 + 60.0 * v) as u8,
                    (50.0 + 40.0 * (1.0 - v)) as u8,
                    (70.0 + 50.0 * v) as u8,
                ];

                // A checkerboard corner, for a region whose texture differs
                // from everything else without its colour doing so.
                if u > 0.72 && v > 0.72 {
                    let cell = ((x / 8) + (y / 8)) % 2;
                    let shade = if cell == 0 { 200 } else { 90 };
                    rgb = [shade, shade, shade];
                }

                for &(cx, cy, r, color) in &discs {
                    let dx = u - cx;
                    let dy = v - cy;
                    if dx * dx + dy * dy < r * r {
                        rgb = color;
                    }
                }

                let i = (y * size + x) * 3;
                self.buf[i] = rgb[0];
                self.buf[i + 1] = rgb[1];
                self.buf[i + 2] = rgb[2];
            }
        }
    }
}

impl FrameSource for TestPattern {
    fn size(&self) -> usize {
        self.size
    }

    fn next_frame(&mut self) -> Option<&[u8]> {
        self.render();
        self.frame = self.frame.wrapping_add(1);
        Some(&self.buf)
    }
}

/// Centre-crop an interleaved RGBA image to a square and resample it down
/// to `size`, dropping alpha.
///
/// Box-averaging rather than nearest: a patch embedding sees 16×16 pixels,
/// and point-sampling a 2560-wide screen down to 224 would alias hard
/// enough to change what the features encode. Cheap here — it runs once per
/// captured frame, not per patch.
pub fn square_downscale_rgba(rgba: &[u8], width: usize, height: usize, size: usize) -> Vec<u8> {
    let side = width.min(height);
    let x0 = (width - side) / 2;
    let y0 = (height - side) / 2;

    let mut out = vec![0u8; size * size * 3];
    for oy in 0..size {
        let sy0 = y0 + oy * side / size;
        let sy1 = (y0 + (oy + 1) * side / size).max(sy0 + 1);
        for ox in 0..size {
            let sx0 = x0 + ox * side / size;
            let sx1 = (x0 + (ox + 1) * side / size).max(sx0 + 1);

            let mut acc = [0u32; 3];
            let mut n = 0u32;
            for sy in sy0..sy1.min(height) {
                for sx in sx0..sx1.min(width) {
                    let i = (sy * width + sx) * 4;
                    acc[0] += rgba[i] as u32;
                    acc[1] += rgba[i + 1] as u32;
                    acc[2] += rgba[i + 2] as u32;
                    n += 1;
                }
            }
            let n = n.max(1);
            let o = (oy * size + ox) * 3;
            for c in 0..3 {
                out[o + c] = (acc[c] / n) as u8;
            }
        }
    }
    out
}

/// Live screen capture.
///
/// Exists mostly so the whole pipeline can be developed and demonstrated
/// without a headset: it is the same `FrameSource` contract the Quest
/// passthrough camera will implement, so the capture loop, colour handling,
/// and downscale are all exercised here first.
#[cfg(feature = "capture")]
pub struct ScreenCapture {
    monitor: xcap::Monitor,
    size: usize,
    buf: Vec<u8>,
}

#[cfg(feature = "capture")]
impl ScreenCapture {
    /// Capture the monitor at `index`, or the primary one if out of range.
    pub fn new(size: usize, index: usize) -> Result<Self, Box<dyn std::error::Error>> {
        let monitors = xcap::Monitor::all()?;
        if monitors.is_empty() {
            return Err("no monitors found".into());
        }
        for (i, m) in monitors.iter().enumerate() {
            log::info!(
                "monitor {i}: {}x{}{}",
                m.width().unwrap_or(0),
                m.height().unwrap_or(0),
                if i == index { " (selected)" } else { "" }
            );
        }
        let monitor = monitors
            .into_iter()
            .nth(index)
            .ok_or("monitor index out of range")?;
        Ok(Self {
            monitor,
            size,
            buf: vec![128; size * size * 3],
        })
    }
}

#[cfg(feature = "capture")]
impl FrameSource for ScreenCapture {
    fn size(&self) -> usize {
        self.size
    }

    fn next_frame(&mut self) -> Option<&[u8]> {
        // A dropped frame is not worth failing over — the previous one is
        // still displayable, and capture hiccups when windows change.
        match self.monitor.capture_image() {
            Ok(image) => {
                let (w, h) = (image.width() as usize, image.height() as usize);
                self.buf = square_downscale_rgba(&image.into_raw(), w, h, self.size);
                Some(&self.buf)
            }
            Err(e) => {
                log::warn!("screen capture failed: {e}");
                None
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn downscale_centre_crops_and_averages() {
        // A 40×20 image: left half red, right half blue. The centre crop is
        // the middle 20×20, which straddles the boundary evenly.
        let (w, h) = (40usize, 20usize);
        let mut rgba = vec![0u8; w * h * 4];
        for y in 0..h {
            for x in 0..w {
                let i = (y * w + x) * 4;
                let c = if x < w / 2 { [255, 0, 0] } else { [0, 0, 255] };
                rgba[i..i + 3].copy_from_slice(&c);
                rgba[i + 3] = 255;
            }
        }

        let out = square_downscale_rgba(&rgba, w, h, 4);
        assert_eq!(out.len(), 4 * 4 * 3);
        // Left column should be red, right column blue.
        let px = |x: usize, y: usize| {
            let i = (y * 4 + x) * 3;
            [out[i], out[i + 1], out[i + 2]]
        };
        assert_eq!(px(0, 0), [255, 0, 0], "left edge should be red");
        assert_eq!(px(3, 0), [0, 0, 255], "right edge should be blue");
    }

    #[test]
    fn downscale_averages_rather_than_point_samples() {
        // Alternating single-pixel columns must average to grey, not pick
        // one extreme. Point sampling would give 0 or 255.
        let (w, h) = (64usize, 64usize);
        let mut rgba = vec![255u8; w * h * 4];
        for y in 0..h {
            for x in 0..w {
                let v = if x % 2 == 0 { 0 } else { 255 };
                let i = (y * w + x) * 4;
                rgba[i..i + 3].copy_from_slice(&[v, v, v]);
            }
        }
        let out = square_downscale_rgba(&rgba, w, h, 8);
        for px in out.chunks_exact(3) {
            assert!(
                (100..=155).contains(&px[0]),
                "expected mid-grey from averaging, got {}",
                px[0]
            );
        }
    }

    #[test]
    fn test_pattern_has_spatial_structure() {
        let mut src = TestPattern::new(224);
        let frame = src.next_frame().unwrap().to_vec();
        assert_eq!(frame.len(), 224 * 224 * 3);

        // A uniform image would make the whole exercise meaningless, so
        // check there is real variation to encode.
        let mean = frame.iter().map(|&b| b as f64).sum::<f64>() / frame.len() as f64;
        let var = frame
            .iter()
            .map(|&b| (b as f64 - mean).powi(2))
            .sum::<f64>()
            / frame.len() as f64;
        assert!(var > 400.0, "test pattern is too flat: variance {var:.1}");
    }

    #[test]
    fn test_pattern_animates() {
        let mut src = TestPattern::new(64);
        let a = src.next_frame().unwrap().to_vec();
        for _ in 0..20 {
            src.next_frame();
        }
        let b = src.next_frame().unwrap().to_vec();
        assert_ne!(a, b, "frames are identical; the scene is not moving");
    }
}