Download environment/training-source/source-snapshot/src/render.rs from mad-bot/dinovision: direct link, hf CLI and curl.
- Browser
- Download file 13.2 kB
-
https://huggingface.co/mad-bot/dinovision/resolve/main/environment/training-source/source-snapshot/src/render.rs
- Command line
-
hf download hf://mad-bot/dinovision/environment/training-source/source-snapshot/src/render.rs
-
curl -L -o render.rs https://huggingface.co/mad-bot/dinovision/resolve/main/environment/training-source/source-snapshot/src/render.rs
13.2 kB
| //! Drawing the coloured feature grid to an eye buffer. | |
| //! | |
| //! Deliberately tiny: a full-screen triangle and a `grid x grid x 3` float | |
| //! buffer. The renderer runs at the headset's refresh rate regardless of | |
| //! how slowly inference produces new grids, so the expensive part of the | |
| //! frame is never on the display path. | |
| use blade_graphics as gpu; | |
| use crate::pca::COMPONENTS; | |
| /// Maximum eyes we allocate parameter slots for. | |
| pub const MAX_EYES: usize = 2; | |
| struct Params { | |
| scale: [f32; 2], | |
| offset: [f32; 2], | |
| grid: u32, | |
| pad: [u32; 3], | |
| } | |
| /// Where a flat overlay lands in one eye's clip space. | |
| /// | |
| /// Purely 2D. The only thing an eye's frustum contributes to a head-locked | |
| /// image is where its axis sits and how wide the image should be — a scale | |
| /// and an offset — so that is all this carries. | |
| pub struct EyeTransform { | |
| pub scale: [f32; 2], | |
| pub offset: [f32; 2], | |
| } | |
| impl EyeTransform { | |
| /// Fills the eye buffer exactly. Right for a flat window, wrong for a | |
| /// headset, where it stretches the image over the whole field of view | |
| /// and ignores the frustum's asymmetry. | |
| pub const FULLSCREEN: Self = Self { | |
| scale: [1.0, 1.0], | |
| offset: [0.0, 0.0], | |
| }; | |
| /// Fill an eye's field of view, centred on a given direction rather | |
| /// than on the eye's axis. | |
| /// | |
| /// `centre_tan` is where the image's middle should sit, in tangent | |
| /// units: `(x/-z, y/-z)` of the target direction expressed in the eye's | |
| /// space. Zero is straight ahead and reproduces [`filling`]. | |
| /// | |
| /// This is how the image gets to stay put while the head moves. The | |
| /// content is as old as the last completed inference — a tenth of a | |
| /// second here — so labelling it with the current head pose tells the | |
| /// compositor to compensate for nothing, and it slides around with the | |
| /// head. Feeding the direction the camera was *pointing when the frame | |
| /// was captured* moves the image opposite to head motion, which is what | |
| /// world-locked looks like. | |
| pub fn filling_at(fov: [f32; 4], centre_tan: [f32; 2]) -> Self { | |
| let [left, right, up, down] = fov; | |
| let (tan_l, tan_r) = (left.tan(), right.tan()); | |
| let (tan_u, tan_d) = (up.tan(), down.tan()); | |
| let (span_x, span_y) = (tan_r - tan_l, tan_u - tan_d); | |
| Self { | |
| // Same size as filling the buffer. | |
| scale: [1.0, 1.0], | |
| offset: [ | |
| (2.0 * centre_tan[0] - (tan_r + tan_l)) / span_x, | |
| (2.0 * centre_tan[1] - (tan_u + tan_d)) / span_y, | |
| ], | |
| } | |
| } | |
| /// Fill an eye's field of view, recentred on its axis. | |
| /// | |
| /// The right default for passthrough parity: the system's own | |
| /// passthrough maps these cameras across your whole view, so matching it | |
| /// means spanning the frustum rather than inventing an angular size. | |
| /// | |
| /// Note this is *not* the same as filling the eye buffer. Headset frusta | |
| /// are asymmetric and mirror between the eyes, so an image painted | |
| /// edge-to-edge has its centre at a different angle in each eye, and the | |
| /// two will not fuse. Spanning the frustum symmetrically about the axis | |
| /// is what makes them agree. | |
| pub fn filling(fov: [f32; 4]) -> Self { | |
| let [left, right, up, down] = fov; | |
| let half_tan = [ | |
| (right.tan() - left.tan()) * 0.5, | |
| (up.tan() - down.tan()) * 0.5, | |
| ]; | |
| Self::for_eye(fov, half_tan) | |
| } | |
| /// Place an image into an eye whose frustum is `[left, right, up, down]` | |
| /// radians, as OpenXR reports it. | |
| /// | |
| /// `half_tan` is the tangent of the image's half-angle on each axis. | |
| /// Separate axes matter: the camera frame is squashed into a square for | |
| /// the encoder, and giving x and y their true angular extents here is | |
| /// what unsquashes it, so the view keeps the camera's full field | |
| /// without looking stretched. | |
| /// | |
| /// A ray at angle θ from the view axis lands at clip | |
| /// `(2·tanθ − (tanR + tanL)) / (tanR − tanL)`, so the quad's ±1 edges | |
| /// map to a scale of `2·half_tan / (tanR − tanL)` about an offset of | |
| /// `−(tanR + tanL) / (tanR − tanL)`. | |
| /// | |
| /// Headset frusta are asymmetric and differ between eyes, which is why | |
| /// the offset is not zero and not shared: dropping it is what leaves the | |
| /// two views unfusable. | |
| pub fn for_eye(fov: [f32; 4], half_tan: [f32; 2]) -> Self { | |
| let [left, right, up, down] = fov; | |
| let (tan_l, tan_r) = (left.tan(), right.tan()); | |
| let (tan_u, tan_d) = (up.tan(), down.tan()); | |
| let (span_x, span_y) = (tan_r - tan_l, tan_u - tan_d); | |
| Self { | |
| scale: [ | |
| 2.0 * half_tan[0] / span_x, | |
| 2.0 * half_tan[1] / span_y, | |
| ], | |
| offset: [-(tan_r + tan_l) / span_x, -(tan_u + tan_d) / span_y], | |
| } | |
| } | |
| } | |
| struct ViewData { | |
| params: gpu::BufferPiece, | |
| cells: gpu::BufferPiece, | |
| } | |
| /// Renders the latest feature grid full-screen. | |
| pub struct GridView { | |
| pipeline: gpu::RenderPipeline, | |
| params_buf: gpu::Buffer, | |
| cells_buf: gpu::Buffer, | |
| grid: usize, | |
| } | |
| impl GridView { | |
| /// `color_format` must match the eye buffer being rendered into — | |
| /// on an XR surface, `XrSurface::format()`. | |
| pub fn new(context: &gpu::Context, color_format: gpu::TextureFormat, grid: usize) -> Self { | |
| let shader = context.create_shader(gpu::ShaderDesc { | |
| source: include_str!("shaders/grid_view.wgsl"), | |
| naga_module: None, | |
| }); | |
| let data_layout = <ViewData as gpu::ShaderData>::layout(); | |
| let pipeline = context.create_render_pipeline(gpu::RenderPipelineDesc { | |
| name: "grid-view", | |
| data_layouts: &[&data_layout], | |
| vertex: shader.at("vs_main"), | |
| vertex_fetches: &[], | |
| primitive: gpu::PrimitiveState::default(), | |
| // The view covers every pixel, so there is nothing to depth-test | |
| // against and no depth buffer to allocate. | |
| depth_stencil: None, | |
| fragment: Some(shader.at("fs_main")), | |
| color_targets: &[gpu::ColorTargetState::from(color_format)], | |
| multisample_state: gpu::MultisampleState::default(), | |
| }); | |
| // One slot per eye: the projection differs between them, and both | |
| // are in flight within a single frame's submission. | |
| let params_buf = context.create_buffer(gpu::BufferDesc { | |
| name: "grid-view-params", | |
| size: (std::mem::size_of::<Params>() * MAX_EYES) as u64, | |
| memory: gpu::Memory::Shared, | |
| }); | |
| let cells_buf = context.create_buffer(gpu::BufferDesc { | |
| name: "grid-view-cells", | |
| size: (grid * grid * COMPONENTS * std::mem::size_of::<f32>()) as u64, | |
| memory: gpu::Memory::Shared, | |
| }); | |
| let mut view = Self { | |
| pipeline, | |
| params_buf, | |
| cells_buf, | |
| grid, | |
| }; | |
| // Full-screen until the caller supplies per-eye transforms, which | |
| // is right for a window and wrong for a headset. | |
| for eye in 0..MAX_EYES { | |
| view.set_transform(context, eye, EyeTransform::FULLSCREEN); | |
| } | |
| // A mid-grey field, so a failure to ever produce a result looks | |
| // like "no data" rather than like a working black-and-nothing view. | |
| view.upload(context, &vec![0.5; grid * grid * COMPONENTS]); | |
| view | |
| } | |
| /// Replace the displayed grid. `rgb` is `[grid * grid, 3]`, row-major. | |
| pub fn upload(&mut self, context: &gpu::Context, rgb: &[f32]) { | |
| let expected = self.grid * self.grid * COMPONENTS; | |
| debug_assert_eq!(rgb.len(), expected, "grid payload is the wrong size"); | |
| let n = rgb.len().min(expected); | |
| unsafe { | |
| std::ptr::copy_nonoverlapping(rgb.as_ptr(), self.cells_buf.data() as *mut f32, n); | |
| } | |
| context.sync_buffer(self.cells_buf); | |
| } | |
| /// Set where the overlay lands in one eye. | |
| pub fn set_transform(&mut self, context: &gpu::Context, eye: usize, t: EyeTransform) { | |
| debug_assert!(eye < MAX_EYES); | |
| unsafe { | |
| let slot = (self.params_buf.data() as *mut Params).add(eye); | |
| *slot = Params { | |
| scale: t.scale, | |
| offset: t.offset, | |
| grid: self.grid as u32, | |
| pad: [0; 3], | |
| }; | |
| } | |
| context.sync_buffer(self.params_buf); | |
| } | |
| /// Draw into an already-started render pass, using `eye`'s transform. | |
| pub fn draw(&self, pass: &mut gpu::RenderCommandEncoder, eye: usize) { | |
| let offset = (eye.min(MAX_EYES - 1) * std::mem::size_of::<Params>()) as u64; | |
| let mut encoder = pass.with(&self.pipeline); | |
| encoder.bind( | |
| 0, | |
| &ViewData { | |
| params: self.params_buf.at(offset), | |
| cells: self.cells_buf.into(), | |
| }, | |
| ); | |
| encoder.draw(0, 6, 0, 1); | |
| } | |
| pub fn destroy(mut self, context: &gpu::Context) { | |
| context.destroy_buffer(self.cells_buf); | |
| context.destroy_buffer(self.params_buf); | |
| context.destroy_render_pipeline(&mut self.pipeline); | |
| } | |
| } | |
| /// Rotate a vector by a quaternion `[x, y, z, w]`. | |
| fn rotate(q: [f32; 4], v: [f32; 3]) -> [f32; 3] { | |
| let (qx, qy, qz, qw) = (q[0], q[1], q[2], q[3]); | |
| // t = 2 * (q_vec × v); v' = v + qw * t + q_vec × t | |
| let tx = 2.0 * (qy * v[2] - qz * v[1]); | |
| let ty = 2.0 * (qz * v[0] - qx * v[2]); | |
| let tz = 2.0 * (qx * v[1] - qy * v[0]); | |
| [ | |
| v[0] + qw * tx + qy * tz - qz * ty, | |
| v[1] + qw * ty + qz * tx - qx * tz, | |
| v[2] + qw * tz + qx * ty - qy * tx, | |
| ] | |
| } | |
| /// Where a direction that was straight ahead under `then` appears under | |
| /// `now`, in tangent units suitable for [`EyeTransform::filling_at`]. | |
| /// | |
| /// Both quaternions are `[x, y, z, w]` in the same reference space. The | |
| /// result is the image's centre after the head has moved, so a rotation to | |
| /// the right pushes the content left — the image staying put in the world | |
| /// while the view sweeps across it. | |
| /// | |
| /// Returns `None` when the old direction has swung behind the viewer, where | |
| /// a tangent-plane shift stops meaning anything. | |
| pub fn reprojection_offset(then: [f32; 4], now: [f32; 4]) -> Option<[f32; 2]> { | |
| // Forward under the capture pose, brought into the current eye's space | |
| // by the inverse of the current pose. | |
| let forward = rotate(then, [0.0, 0.0, -1.0]); | |
| let inverse_now = [-now[0], -now[1], -now[2], now[3]]; | |
| let v = rotate(inverse_now, forward); | |
| if v[2] >= -1e-3 { | |
| return None; | |
| } | |
| Some([v[0] / -v[2], v[1] / -v[2]]) | |
| } | |
| mod tests { | |
| use super::*; | |
| /// Straight ahead stays straight ahead when the head has not moved. | |
| fn no_motion_means_no_shift() { | |
| let identity = [0.0, 0.0, 0.0, 1.0]; | |
| let offset = reprojection_offset(identity, identity).unwrap(); | |
| assert!(offset[0].abs() < 1e-6 && offset[1].abs() < 1e-6); | |
| } | |
| /// Turning the head right must push the content left, so the image | |
| /// appears to stay where it was in the world. | |
| fn turning_right_pushes_content_left() { | |
| let identity = [0.0, 0.0, 0.0, 1.0]; | |
| // +15 degrees about Y is a leftward yaw in a right-handed system | |
| // looking down -Z, so -15 turns the view to the right. | |
| let half = (-15.0f32).to_radians() * 0.5; | |
| let turned_right = [0.0, half.sin(), 0.0, half.cos()]; | |
| let offset = reprojection_offset(identity, turned_right).unwrap(); | |
| assert!( | |
| offset[0] < -0.1, | |
| "expected the image to move left, got {offset:?}" | |
| ); | |
| // The magnitude should be tan(15 deg) = 0.268. | |
| assert!((offset[0] + 15.0f32.to_radians().tan()).abs() < 1e-3); | |
| assert!(offset[1].abs() < 1e-6, "yaw should not shift vertically"); | |
| } | |
| /// Looking up must push the content down. | |
| fn looking_up_pushes_content_down() { | |
| let identity = [0.0, 0.0, 0.0, 1.0]; | |
| let half = 10.0f32.to_radians() * 0.5; | |
| let looked_up = [half.sin(), 0.0, 0.0, half.cos()]; | |
| let offset = reprojection_offset(identity, looked_up).unwrap(); | |
| assert!(offset[1] < -0.1, "expected downward shift, got {offset:?}"); | |
| assert!(offset[0].abs() < 1e-6); | |
| } | |
| /// A view swung right round has nothing sensible to show. | |
| fn facing_away_has_no_offset() { | |
| let identity = [0.0, 0.0, 0.0, 1.0]; | |
| let behind = [0.0, 1.0, 0.0, 0.0]; | |
| assert!(reprojection_offset(identity, behind).is_none()); | |
| } | |
| } | |