File size: 5,638 Bytes
eae424a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | //! Run the full pipeline on desktop and write the result as an image.
//!
//! Exercises exactly what the headset does — inference worker, PCA fit,
//! colour projection — minus OpenXR, and writes a side-by-side PPM so the
//! output can actually be looked at. Rendering a feature view you cannot
//! see is a poor way to find out it is wrong.
//!
//! ```text
//! cargo run --release --example preview -- [out.ppm] [model.safetensors]
//! ```
//!
//! Without a checkpoint it runs on synthetic weights: the pipeline is
//! exercised end to end, but the colours mean nothing. With real weights,
//! regions that belong to the same object should share a colour.
use std::path::PathBuf;
use std::time::{Duration, Instant};
use dinovision::dinov3::Config;
use dinovision::inference::{self, Weights};
use dinovision::pca::COMPONENTS;
use dinovision::source::{FrameSource, TestPattern};
/// Nearest-neighbour upscale of the feature grid, so each patch reads as a
/// solid block. Bilinear would look prettier and hide the true resolution.
fn upscale(grid_rgb: &[f32], grid: usize, out: usize) -> Vec<u8> {
let mut px = vec![0u8; out * out * 3];
for y in 0..out {
let gy = (y * grid) / out;
for x in 0..out {
let gx = (x * grid) / out;
let src = (gy * grid + gx) * COMPONENTS;
let dst = (y * out + x) * 3;
for c in 0..3 {
px[dst + c] = (grid_rgb[src + c].clamp(0.0, 1.0) * 255.0) as u8;
}
}
}
px
}
/// Source on the left, feature view on the right.
fn write_side_by_side(path: &PathBuf, left: &[u8], right: &[u8], size: usize) {
let mut joined = vec![0u8; size * size * 6];
for y in 0..size {
let src = y * size * 3;
let dst = y * size * 6;
joined[dst..dst + size * 3].copy_from_slice(&left[src..src + size * 3]);
joined[dst + size * 3..dst + size * 6].copy_from_slice(&right[src..src + size * 3]);
}
image::RgbImage::from_raw(size as u32 * 2, size as u32, joined)
.expect("image dimensions do not match buffer")
.save(path)
.expect("failed to write image");
}
fn main() {
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
let mut args = std::env::args().skip(1);
let out_path = PathBuf::from(args.next().unwrap_or_else(|| "preview.png".into()));
let weights = match args.next().map(PathBuf::from) {
Some(p) if p.exists() => {
log::info!("using weights from {}", p.display());
Weights::SafeTensors(p)
}
Some(p) => panic!("{} does not exist", p.display()),
None => Weights::Synthetic,
};
let config = Config::vits16().at_resolution(256);
let gpu = dinovision::init_context(None).expect("failed to initialize GPU context");
let worker = inference::spawn(gpu, config.clone(), weights, None, 1, dinovision::inference::Display::PcaColour);
// `DINOVISION_SOURCE=screen` grabs the desktop instead of the synthetic
// scene, which is the quickest way to see the colouring against real
// imagery without opening a window.
let mut source: Box<dyn FrameSource> = match std::env::var("DINOVISION_SOURCE").as_deref() {
#[cfg(feature = "capture")]
Ok("screen") => match dinovision::source::ScreenCapture::new(config.image_size, 0) {
Ok(s) => Box::new(s),
Err(e) => panic!("screen capture unavailable: {e}"),
},
#[cfg(not(feature = "capture"))]
Ok("screen") => panic!("rebuild with --features capture for screen capture"),
_ => Box::new(TestPattern::new(config.image_size)),
};
let mut last_frame: Vec<u8> = Vec::new();
// Let the session build, then run a few frames so the PCA basis is
// refitted from real features rather than the placeholder.
let deadline = Instant::now() + Duration::from_secs(180);
let mut completed = 0;
while completed < 5 && Instant::now() < deadline {
if worker.is_ready() {
let rgb = source.next_frame().unwrap();
last_frame = rgb.to_vec();
let patches = dinovision::preprocess::patches_from_rgb8(rgb, &config);
worker.submit(patches);
}
std::thread::sleep(Duration::from_millis(20));
completed = worker.generation();
}
let grid = worker
.latest()
.expect("inference produced no result before the deadline");
log::info!(
"{} completed encodes, last took {:.1} ms ({}x{} grid)",
completed,
grid.latency_ms,
grid.grid,
grid.grid
);
let size = config.image_size;
let features = upscale(grid.patch_rgb(), grid.grid, size);
write_side_by_side(&out_path, &last_frame, &features, size);
// A view that is one flat colour means the projection collapsed, which
// is easy to miss by eye in a small image.
let spread = {
let mut lo = [255u8; 3];
let mut hi = [0u8; 3];
for px in features.chunks_exact(3) {
for c in 0..3 {
lo[c] = lo[c].min(px[c]);
hi[c] = hi[c].max(px[c]);
}
}
(0..3).map(|c| hi[c] as i32 - lo[c] as i32).max().unwrap()
};
println!("wrote {} (source | features)", out_path.display());
println!("colour spread: {spread}/255");
if spread < 16 {
println!("WARNING: the feature view is nearly flat — projection may have collapsed");
}
}
|