Upload folder using huggingface_hub (part 2)

eae424a verified 3 months ago

10.3 kB

//! 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");
    }
}

	//! 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");
	}
	}