📄 replay/src/synth.rs
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 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479
//! The synthetic depth-camera: a software Kinect.
//!
//! Every frame, [`SyntheticSource`] raycasts the scene from the rig's
//! center viewpoint — one ray per depth pixel — and emits a raw
//! [`DepthFrame`]: a millimetre depth map plus RGB, like a Kinect v2
//! reports them. `splatgen` then turns those frames into surfels, so the
//! dev/CI path exercises the real depth-processing pipeline.
//!
//! Realism knobs: the source injects deterministic sensor noise
//! ([`NoiseParams`]) — isolated fliers and per-pixel jitter — so the
//! flier removal and bilateral filter have real work to do. The noise is
//! a pure function of `(frame, pixel)`, never the wall clock: playback
//! stays deterministic.
//!
//! The subject sways a little, also deterministically, so temporal
//! behaviour is testable.

use anyhow::Result;
use glam::{UVec2, Vec3};

use rig::{Camera, CameraIntrinsics, CameraPose, DepthFrame, DepthSource};

use crate::scene::{self, Scene, Surface};

/// Where the synthetic subject sits: 2 m in front of the rig's center
/// camera at the origin.
pub const SUBJECT: Vec3 = Vec3::new(0.0, 0.0, 2.0);

/// Opacity used for synthetic surfels (PLAN.md: ~0.9–0.99). Applied by
/// `splatgen`, kept here as the dataset's declared intent.
pub const SURFEL_OPACITY: f32 = 0.95;

/// Subject sway amplitude, metres, and period in frames.
const SWAY_AMPLITUDE_M: f32 = 0.04;
const SWAY_PERIOD_FRAMES: u64 = 90;

/// The depth range a Kinect v2 reports; synthetic noise stays in it.
const KINECT_MIN_MM: u16 = 500;
const KINECT_MAX_MM: u16 = 4000;

/// Sensor noise injected into the synthetic depth maps.
///
/// Deterministic: a pure function of the frame index and pixel, so
/// playback is reproducible.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct NoiseParams {
    /// Fraction of valid pixels replaced with a random depth (fliers),
    /// in `[0, 1]`.
    pub flier_rate: f32,
    /// Per-pixel depth jitter, ± millimetres.
    pub jitter_mm: u32,
}

impl Default for NoiseParams {
    fn default() -> Self {
        // Mild but realistic: scattered fliers, a few mm of jitter.
        Self {
            flier_rate: 0.005,
            jitter_mm: 4,
        }
    }
}

impl NoiseParams {
    /// Clean frames: no fliers, no jitter.
    pub const OFF: NoiseParams = NoiseParams {
        flier_rate: 0.0,
        jitter_mm: 0,
    };
}

/// Deterministic uniform noise in `[0, 1)` for a pixel: a hash of
/// `(frame, x, y, salt)` — no RNG state, so playback is reproducible.
fn noise01(frame: u64, x: u32, y: u32, salt: u32) -> f32 {
    let mut h = (frame as u32)
        ^ x.wrapping_mul(0x9E37_79B1)
        ^ y.wrapping_mul(0x85EB_CA77)
        ^ salt.wrapping_mul(0xC2B2_AE3D);
    h ^= h >> 16;
    h = h.wrapping_mul(0x7FEB_352D);
    h ^= h >> 15;
    h = h.wrapping_mul(0x846C_A68B);
    h ^= h >> 16;
    h as f32 / u32::MAX as f32
}

/// The development rig: a Kinect-like center camera flanked by two
/// webcams 0.8 m apart, converged on [`SUBJECT`].
///
/// The center camera approximates the Kinect v2 depth sensor (70° × 60°
/// field of view); the side cameras are typical 60° webcams. Pass smaller
/// resolutions to keep the same fields of view with fewer pixels — what
/// the test suites use to stay fast.
pub fn dev_rig_with_resolutions(center: UVec2, side: UVec2) -> Rig {
    // Kinect v2 depth: ~70° x 60° field of view.
    let center = Camera {
        id: "kinect".into(),
        pose: CameraPose::looking_at(Vec3::ZERO, SUBJECT),
        intrinsics: CameraIntrinsics {
            fx: CENTER_FX * center.x as f32,
            fy: CENTER_FY * center.y as f32,
            cx: center.x as f32 / 2.0,
            cy: center.y as f32 / 2.0,
        },
        resolution: center,
    };
    // Typical webcams with a 60 degree horizontal FOV.
    let make_side = |x: f32| Camera {
        id: if x < 0.0 { "left" } else { "right" }.to_owned(),
        pose: CameraPose::looking_at(Vec3::new(x, 0.0, 0.0), SUBJECT),
        intrinsics: CameraIntrinsics {
            fx: SIDE_FX * side.x as f32,
            fy: SIDE_FX * side.y as f32,
            cx: side.x as f32 / 2.0,
            cy: side.y as f32 / 2.0,
        },
        resolution: side,
    };
    Rig {
        left: make_side(-0.4),
        center,
        right: make_side(0.4),
    }
}

/// Center-camera focal ratios (per resolution pixel): a ~70° horizontal
/// and 60° vertical field of view, like the Kinect v2 depth sensor.
const CENTER_FX: f32 = 0.714_074; // 1 / (2 * tan(35°)), FOV_x = 70°
const CENTER_FY: f32 = 0.866_025_4; // 1 / (2 * tan(30°)), FOV_y = 60°
/// Side-camera focal ratio: a 60° horizontal FOV, square pixels.
const SIDE_FX: f32 = 0.866_025_4; // 1 / (2 * tan(30°))

use rig::Rig;

/// The development rig at full size: Kinect-like depth at 256x212
/// (~54k depth pixels, inside the plan's 50–150k splat budget) and
/// 640x480 webcams.
pub fn dev_rig() -> Rig {
    dev_rig_with_resolutions(UVec2::new(256, 212), UVec2::new(640, 480))
}

/// A synthetic source over the default development scene (see
/// [`dev_rig`]), at `fps` frames per second.
pub fn dev_source(fps: f32) -> SyntheticSource {
    dev_source_with_center_camera(dev_rig().center, fps)
}

/// Like [`dev_source`], but with a custom center (depth) camera and
/// subject position — what the configurable app uses.
pub fn dev_source_with(center: Camera, subject: Vec3, fps: f32) -> SyntheticSource {
    let (head, shoulders) = scene::head_and_shoulders(subject);
    SyntheticSource::new(
        Scene::new().with(scene::backdrop(subject)),
        vec![head, shoulders],
        center,
        fps,
    )
}

/// Like [`dev_source`], but with a custom center (depth) camera — e.g.
/// `dev_rig_with_resolutions(..).center` at a small resolution for tests.
pub fn dev_source_with_center_camera(center: Camera, fps: f32) -> SyntheticSource {
    dev_source_with(center, SUBJECT, fps)
}

/// A synthetic depth-camera frame source.
pub struct SyntheticSource {
    /// Static surroundings (backdrop).
    scene: Scene,
    /// The animated subject surfaces, translated by the sway offset each
    /// frame.
    subject: Vec<Surface>,
    /// The depth viewpoint: the rig's center camera.
    camera: Camera,
    fps: f32,
    noise: NoiseParams,
    frame_index: u64,
}

impl SyntheticSource {
    pub fn new(scene: Scene, subject: Vec<Surface>, camera: Camera, fps: f32) -> Self {
        Self {
            scene,
            subject,
            camera,
            fps: fps.max(1.0),
            noise: NoiseParams::default(),
            frame_index: 0,
        }
    }

    /// Override the sensor noise; [`NoiseParams::OFF`] for clean frames.
    pub fn with_noise(mut self, noise: NoiseParams) -> Self {
        self.noise = noise;
        self
    }

    /// The number of frames produced so far.
    pub fn frame_index(&self) -> u64 {
        self.frame_index
    }

    /// The depth viewpoint: the rig's center camera.
    pub fn center(&self) -> &Camera {
        &self.camera
    }

    /// The deterministic subject sway offset for a given frame.
    fn sway_offset(&self, frame: u64) -> Vec3 {
        let phase = 2.0 * std::f32::consts::PI * (frame % SWAY_PERIOD_FRAMES) as f32
            / SWAY_PERIOD_FRAMES as f32;
        Vec3::new(SWAY_AMPLITUDE_M * phase.sin(), 0.0, 0.0)
    }

    /// The full scene for a given frame: static surroundings plus the
    /// swayed subject.
    fn frame_scene(&self, frame: u64) -> Scene {
        let offset = self.sway_offset(frame);
        let mut scene = self.scene.clone();
        for surface in &self.subject {
            scene.surfaces.push(surface.translated(offset));
        }
        scene
    }

    /// Raycast one frame: the clean depth map and RGB the software
    /// Kinect sees.
    fn cast_frame(&self, frame: u64) -> (Vec<u16>, Vec<u8>) {
        let scene = self.frame_scene(frame);
        let (width, height) = (self.camera.resolution.x, self.camera.resolution.y);
        let camera_from_world = self.camera.pose.camera_from_world();
        let mut depth = vec![0_u16; (width * height) as usize];
        let mut rgb = vec![0_u8; (width * height * 3) as usize];

        for y in 0..height {
            for x in 0..width {
                // Sample through pixel centres.
                let (origin, dir) = self.camera.pixel_ray(x as f32 + 0.5, y as f32 + 0.5);
                let Some(hit) = scene.cast_ray(origin, dir) else {
                    continue; // no return: depth stays 0
                };
                // Planar depth along the viewing axis, in millimetres —
                // what a Kinect depth map reports.
                let z = camera_from_world.transform_point3(hit.position).z;
                let idx = (y * width + x) as usize;
                depth[idx] = (z * 1000.0).round().clamp(1.0, u16::MAX as f32) as u16;
                rgb[idx * 3] = (hit.color[0] * 255.0) as u8;
                rgb[idx * 3 + 1] = (hit.color[1] * 255.0) as u8;
                rgb[idx * 3 + 2] = (hit.color[2] * 255.0) as u8;
            }
        }
        (depth, rgb)
    }

    /// Apply the deterministic sensor noise to a clean frame.
    fn add_noise(&self, depth: &mut [u16], frame: u64) {
        if self.noise == NoiseParams::OFF {
            return;
        }
        let width = self.camera.resolution.x;
        for (i, d) in depth.iter_mut().enumerate() {
            if *d == 0 {
                continue; // invalid pixels stay invalid
            }
            let x = (i as u32) % width;
            let y = (i as u32) / width;
            if self.noise.jitter_mm > 0 {
                let jitter = noise01(frame, x, y, 1) * 2.0 - 1.0;
                let delta = (jitter * self.noise.jitter_mm as f32).round() as i32;
                *d = (*d as i32 + delta).clamp(1, u16::MAX as i32) as u16;
            }
            if self.noise.flier_rate > 0.0 && noise01(frame, x, y, 2) < self.noise.flier_rate {
                let span = (KINECT_MAX_MM - KINECT_MIN_MM) as f32;
                *d = KINECT_MIN_MM + (noise01(frame, x, y, 3) * span) as u16;
            }
        }
    }

    /// Build the raw depth frame for `frame`.
    fn build_frame(&self, frame: u64) -> DepthFrame {
        let (mut depth, rgb) = self.cast_frame(frame);
        self.add_noise(&mut depth, frame);
        DepthFrame {
            timestamp_secs: frame as f64 / self.fps as f64,
            depth_mm: depth,
            rgb,
        }
    }
}

impl DepthSource for SyntheticSource {
    fn camera(&self) -> &Camera {
        &self.camera
    }

    fn next_depth_frame(&mut self) -> Result<DepthFrame> {
        let frame = self.frame_index;
        self.frame_index += 1;
        Ok(self.build_frame(frame))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::scene;
    use glam::uvec2;
    use rig::DepthSource;

    /// A small-resolution source so tests stay fast: 32x24 depth pixels.
    fn tiny_source() -> SyntheticSource {
        let (head, shoulders) = scene::head_and_shoulders(SUBJECT);
        let mut center = dev_rig().center;
        // Keep the center camera's FOV, shrink the resolution.
        let scale = 32.0 / center.resolution.x as f32;
        center.intrinsics.fx *= scale;
        center.intrinsics.fy *= scale;
        center.intrinsics.cx = 16.0;
        center.intrinsics.cy = 12.0;
        center.resolution = uvec2(32, 24);
        SyntheticSource::new(
            Scene::new().with(scene::backdrop(SUBJECT)),
            vec![head, shoulders],
            center,
            30.0,
        )
    }

    /// The straight-ahead ray passes below the head and hits the front
    /// of the shoulders: ~1.85 m out.
    const SHOULDER_FRONT_MM: f32 = 1852.0;

    #[test]
    fn frames_are_depth_maps_like_a_kinect() {
        let mut src = tiny_source();
        let frame = src.next_depth_frame().expect("frame");
        assert_eq!(frame.depth_mm.len(), 32 * 24);
        assert_eq!(frame.rgb.len(), 32 * 24 * 3);

        // Every ray hits something (the backdrop is infinite): all
        // depths valid.
        assert!(frame.depth_mm.iter().all(|&d| d > 0));

        // The pixel looking straight at the subject sees the front of
        // the shoulders (± noise), in the shoulders' colour.
        let mid = 12 * 32 + 16;
        assert!(
            (frame.depth_mm[mid] as f32 - SHOULDER_FRONT_MM).abs() < 40.0,
            "center depth {} mm",
            frame.depth_mm[mid]
        );
        let shoulders = &frame.rgb[mid * 3..mid * 3 + 3];
        assert!(
            (shoulders[0] as i32 - 51).abs() < 6
                && (shoulders[1] as i32 - 89).abs() < 6
                && (shoulders[2] as i32 - 140).abs() < 6,
            "shoulders rgb {shoulders:?}"
        );

        // A top-corner pixel sees the backdrop behind the subject.
        let corner = frame.depth_mm[32 + 2] as f32;
        assert!((corner - 2800.0).abs() < 200.0, "corner depth {corner}");
    }

    #[test]
    fn timestamps_advance_with_index() {
        let mut src = tiny_source();
        assert_eq!(src.next_depth_frame().unwrap().timestamp_secs, 0.0);
        assert!((src.next_depth_frame().unwrap().timestamp_secs - 1.0 / 30.0).abs() < 1e-9);
        assert!((src.next_depth_frame().unwrap().timestamp_secs - 2.0 / 30.0).abs() < 1e-9);
    }

    #[test]
    fn playback_is_deterministic() {
        let mut a = tiny_source();
        let mut b = tiny_source();
        for _ in 0..5 {
            assert_eq!(a.next_depth_frame().unwrap(), b.next_depth_frame().unwrap());
        }
    }

    #[test]
    fn subject_sways_and_repeats() {
        // The sway is periodic in the frame index; the sensor noise is
        // not, so exact repetition is asserted on clean frames.
        let mut src = tiny_source().with_noise(NoiseParams::OFF);
        let frame0 = src.next_depth_frame().unwrap();
        // After a quarter period the sway is near its maximum.
        let mut quarter = None;
        while src.frame_index() <= SWAY_PERIOD_FRAMES / 4 {
            quarter = Some(src.next_depth_frame().unwrap());
        }
        let quarter = quarter.expect("quarter frame produced");
        assert_ne!(frame0.depth_mm, quarter.depth_mm, "no visible sway");

        // The frame at index SWAY_PERIOD_FRAMES repeats frame 0 exactly.
        let mut period = None;
        while src.frame_index() <= SWAY_PERIOD_FRAMES {
            period = Some(src.next_depth_frame().unwrap());
        }
        let period = period.expect("period frame produced");
        assert_eq!(
            frame0.depth_mm, period.depth_mm,
            "sway must be periodic in the frame index"
        );
        assert_eq!(frame0.rgb, period.rgb);

        // With noise on, the sway still moves the same pixels but the
        // exact map differs — noise must not be frame-periodic.
        let mut noisy = tiny_source();
        let n0 = noisy.next_depth_frame().unwrap();
        let mut quarter = None;
        while noisy.frame_index() <= SWAY_PERIOD_FRAMES / 4 {
            quarter = Some(noisy.next_depth_frame().unwrap());
        }
        assert_ne!(n0.depth_mm, quarter.unwrap().depth_mm);
    }

    #[test]
    fn sway_offset_is_deterministic_and_bounded() {
        let src = tiny_source();
        assert_eq!(src.sway_offset(0), Vec3::ZERO);
        assert_eq!(src.sway_offset(SWAY_PERIOD_FRAMES), Vec3::ZERO);
        let max = src.sway_offset(SWAY_PERIOD_FRAMES / 4);
        assert!(max.x > 0.99 * SWAY_AMPLITUDE_M, "sway max was {max:?}");
        for frame in 0..SWAY_PERIOD_FRAMES * 2 {
            let o = src.sway_offset(frame);
            assert!(o.x.abs() <= SWAY_AMPLITUDE_M + 1e-5);
        }
    }

    #[test]
    fn noise_is_deterministic_and_can_be_disabled() {
        // Noise must not break determinism, and OFF must give clean maps.
        let mut noisy = tiny_source();
        let mut clean = tiny_source().with_noise(NoiseParams::OFF);

        let noisy_frame = noisy.next_depth_frame().unwrap();
        let clean_frame = clean.next_depth_frame().unwrap();
        assert_ne!(noisy_frame.depth_mm, clean_frame.depth_mm);

        // Clean frames repeat exactly across sources.
        let mut clean2 = tiny_source().with_noise(NoiseParams::OFF);
        assert_eq!(clean_frame, clean2.next_depth_frame().unwrap());

        // And noisy ones too.
        let mut noisy2 = tiny_source();
        assert_eq!(noisy_frame, noisy2.next_depth_frame().unwrap());

        // Jitter is bounded: every noisy depth within ±jitter of the
        // clean one (fliers excepted).
        let mut fliers = 0;
        for (n, c) in noisy_frame.depth_mm.iter().zip(&clean_frame.depth_mm) {
            if (*n as i32 - *c as i32).abs() > NoiseParams::default().jitter_mm as i32 + 1 {
                fliers += 1;
            }
        }
        assert!(
            fliers as f32 <= 3.0 * NoiseParams::default().flier_rate * 32.0 * 24.0,
            "too many changed pixels: {fliers}"
        );
    }

    #[test]
    fn dev_rig_is_symmetric() {
        let rig = dev_rig();
        assert_eq!(rig.left.resolution, uvec2(640, 480));
        assert_eq!(rig.right.resolution, uvec2(640, 480));
        assert_eq!(rig.center.resolution, uvec2(256, 212));
        // The center camera sits exactly between the side cameras.
        let mid = rig.virtual_pose(0.5);
        assert!(mid.position.distance(rig.center.pose.position) < 1e-6);
        // And it looks at the subject.
        let forward = mid.rotation * Vec3::Z;
        let to_subject = (SUBJECT - mid.position).normalize();
        assert!(forward.dot(to_subject) > 0.999);
    }
}