src/relight.rs
Build-time source snapshot · Plain text
1//! Surface-independent teaching renderer: known materials -> GGX -> splats -> HDR.2//! This is not an inverse renderer or a reproduction of a research method.3use bytemuck::{Pod, Zeroable};4use glam::{Mat3, Mat4, Vec3};5use wgpu::util::DeviceExt;67const COUNT: usize = 1536;8const BASE: Vec3 = Vec3::new(0.55, 0.14, 0.045); // linear, artist-authored9#[repr(C)]10#[derive(Clone, Copy, Pod, Zeroable)]11struct Globals {12 vp: [[f32; 4]; 4],13 eye: [f32; 4],14 light: [f32; 4],15 material: [f32; 4],16 viewport: [f32; 4],17}18// ANCHOR: asset19#[repr(C)]20#[derive(Clone, Copy, Pod, Zeroable)]21struct Splat {22 center: [f32; 4],23 u: [f32; 4],24 v: [f32; 4],25 w: [f32; 4],26 normal: [f32; 4],27 base: [f32; 4],28}29// END: asset3031#[derive(Clone, Copy, Debug)]32pub struct RelightFrame {33 pub azimuth: f32,34 pub elevation: f32,35 pub intensity: f32,36 pub roughness: f32,37 pub metallic: f32,38 pub time: f32,39 pub yaw: f32,40 pub mode: u32, // combined, diffuse, specular, normals, base color41 pub wrong_normals: bool,42}43impl Default for RelightFrame {44 fn default() -> Self {45 Self {46 azimuth: -0.6,47 elevation: 0.5,48 intensity: 3.0,49 roughness: 0.4,50 metallic: 0.0,51 time: 0.0,52 yaw: 0.0,53 mode: 0,54 wrong_normals: false,55 }56 }57}58impl RelightFrame {59 pub fn validate(self) -> Result<(), String> {60 if ![61 self.azimuth,62 self.elevation,63 self.intensity,64 self.roughness,65 self.metallic,66 self.time,67 self.yaw,68 ]69 .iter()70 .all(|v| v.is_finite())71 || self.azimuth.abs() > std::f32::consts::TAU72 || self.elevation.abs() > 1.5773 || self.yaw.abs() > std::f32::consts::TAU74 || !(0.0..=12.0).contains(&self.intensity)75 || !(0.12..=1.0).contains(&self.roughness)76 || !(0.0..=1.0).contains(&self.metallic)77 || !(0.0..=1.0).contains(&self.time)78 || self.mode > 479 {80 return Err("Invalid relighting parameter (roughness range: 0.12..1)".into());81 }82 Ok(())83 }84 pub fn light(self) -> Vec3 {85 Vec3::new(86 self.azimuth.sin() * self.elevation.cos(),87 self.elevation.sin(),88 self.azimuth.cos() * self.elevation.cos(),89 )90 }91 pub fn eye(self) -> Vec3 {92 Vec3::new(4.2 * self.yaw.sin(), 0.5, 4.2 * self.yaw.cos())93 }94 pub fn transform(self) -> Mat3 {95 let phase = self.time * std::f32::consts::TAU;96 Mat3::from_rotation_y(phase)97 * Mat3::from_diagonal(Vec3::new(98 1.0 + 0.4 * phase.sin(),99 1.15,100 0.85 - 0.2 * phase.sin(),101 ))102 }103}104105// ANCHOR: deformation106fn surfel(n: Vec3, frame: RelightFrame) -> Splat {107 let a = frame.transform();108 let helper = if n.y.abs() < 0.9 { Vec3::Y } else { Vec3::X };109 let u = n.cross(helper).normalize() * 0.065;110 let v = n.cross(u).normalize() * 0.065;111 let normal_matrix = if frame.wrong_normals {112 a113 } else {114 a.inverse().transpose()115 };116 Splat {117 center: (a * n).extend(0.96).to_array(),118 u: (a * u).extend(0.0).to_array(),119 v: (a * v).extend(0.0).to_array(),120 w: (a * n * 0.008).extend(0.0).to_array(),121 normal: (normal_matrix * n).normalize().extend(0.0).to_array(),122 base: BASE.extend(0.0).to_array(),123 }124}125fn scene(frame: RelightFrame, view: Mat4) -> Vec<Splat> {126 let mut data: Vec<_> = (0..COUNT)127 .map(|i| {128 let y = 1.0 - 2.0 * (i as f32 + 0.5) / COUNT as f32;129 let phi = i as f32 * 2.3999632;130 let r = (1.0 - y * y).sqrt();131 surfel(Vec3::new(r * phi.cos(), y, r * phi.sin()), frame)132 })133 .collect();134 data.sort_by(|a, b| {135 let z = |s: &Splat| view.transform_point3(Vec3::from_slice(&s.center[..3])).z;136 z(a).total_cmp(&z(b)) // RH camera looks along -Z: far first137 });138 data139}140// END: deformation141142/// CPU explanation of the same single-scattering model as shade() in WGSL.143/// Returns [diffuse reflected radiance, specular reflected radiance].144pub fn reference_brdf(n: Vec3, v: Vec3, l: Vec3, frame: RelightFrame) -> [Vec3; 2] {145 let nl = n.dot(l).max(0.0);146 let nv = n.dot(v).max(0.0);147 if nl <= 0.0 || nv <= 0.0 {148 return [Vec3::ZERO; 2];149 }150 let h = (l + v).normalize();151 let a2 = frame.roughness.powi(4);152 let d = a2 / (std::f32::consts::PI * (n.dot(h).max(0.0).powi(2) * (a2 - 1.0) + 1.0).powi(2));153 let f0 = Vec3::splat(0.04).lerp(BASE, frame.metallic);154 let f = f0 + (Vec3::ONE - f0) * (1.0 - v.dot(h).max(0.0)).powi(5);155 let g1 = |c: f32| 2.0 * c / (c + (a2 + (1.0 - a2) * c * c).sqrt()).max(1e-6);156 let fd = (Vec3::ONE - f) * (1.0 - frame.metallic) * BASE / std::f32::consts::PI;157 let fs = d * f * g1(nl) * g1(nv) / (4.0 * nl * nv).max(1e-6);158 [fd * nl * frame.intensity, fs * nl * frame.intensity]159}160161pub struct RelightPass {162 splat_pipeline: wgpu::RenderPipeline,163 display_pipeline: wgpu::RenderPipeline,164 globals: wgpu::Buffer,165 splats: wgpu::Buffer,166 bindings: wgpu::BindGroup,167 hdr_layout: wgpu::BindGroupLayout,168 hdr: Option<(wgpu::TextureView, wgpu::BindGroup, [u32; 2])>,169 srgb: bool,170}171impl RelightPass {172 pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat) -> Self {173 let globals = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {174 label: Some("relight-globals"),175 contents: bytemuck::bytes_of(&Globals::zeroed()),176 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,177 });178 let splats = device.create_buffer(&wgpu::BufferDescriptor {179 label: Some("relight-known-surfels"),180 size: (COUNT * std::mem::size_of::<Splat>()) as u64,181 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,182 mapped_at_creation: false,183 });184 let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {185 label: Some("relight-layout"),186 entries: &[187 wgpu::BindGroupLayoutEntry {188 binding: 0,189 visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,190 ty: wgpu::BindingType::Buffer {191 ty: wgpu::BufferBindingType::Uniform,192 has_dynamic_offset: false,193 min_binding_size: None,194 },195 count: None,196 },197 wgpu::BindGroupLayoutEntry {198 binding: 1,199 visibility: wgpu::ShaderStages::VERTEX,200 ty: wgpu::BindingType::Buffer {201 ty: wgpu::BufferBindingType::Storage { read_only: true },202 has_dynamic_offset: false,203 min_binding_size: None,204 },205 count: None,206 },207 ],208 });209 let bindings = device.create_bind_group(&wgpu::BindGroupDescriptor {210 label: None,211 layout: &layout,212 entries: &[213 wgpu::BindGroupEntry {214 binding: 0,215 resource: globals.as_entire_binding(),216 },217 wgpu::BindGroupEntry {218 binding: 1,219 resource: splats.as_entire_binding(),220 },221 ],222 });223 let hdr_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {224 label: None,225 entries: &[wgpu::BindGroupLayoutEntry {226 binding: 0,227 visibility: wgpu::ShaderStages::FRAGMENT,228 ty: wgpu::BindingType::Texture {229 sample_type: wgpu::TextureSampleType::Float { filterable: false },230 view_dimension: wgpu::TextureViewDimension::D2,231 multisampled: false,232 },233 count: None,234 }],235 });236 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {237 label: Some("relight-wgsl"),238 source: wgpu::ShaderSource::Wgsl(include_str!("relight.wgsl").into()),239 });240 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {241 label: None,242 bind_group_layouts: &[Some(&layout), Some(&hdr_layout)],243 immediate_size: 0,244 });245 let splat_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {246 label: Some("relight-splat-layout"),247 bind_group_layouts: &[Some(&layout)],248 immediate_size: 0,249 });250 let pipeline = |vertex, fragment, format, blend, layout: &wgpu::PipelineLayout| {251 device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {252 label: Some(fragment),253 layout: Some(layout),254 vertex: wgpu::VertexState {255 module: &shader,256 entry_point: Some(vertex),257 compilation_options: Default::default(),258 buffers: &[],259 },260 fragment: Some(wgpu::FragmentState {261 module: &shader,262 entry_point: Some(fragment),263 compilation_options: Default::default(),264 targets: &[Some(wgpu::ColorTargetState {265 format,266 blend,267 write_mask: wgpu::ColorWrites::ALL,268 })],269 }),270 primitive: Default::default(),271 depth_stencil: None,272 multisample: Default::default(),273 multiview_mask: None,274 cache: None,275 })276 };277 let splat_pipeline = pipeline(278 "vs_main",279 "fs_main",280 wgpu::TextureFormat::Rgba16Float,281 Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),282 &splat_layout,283 );284 let display_pipeline = pipeline("vs_display", "fs_display", format, None, &pipeline_layout);285 Self {286 splat_pipeline,287 display_pipeline,288 globals,289 splats,290 bindings,291 hdr_layout,292 hdr: None,293 srgb: format.is_srgb(),294 }295 }296297 // ANCHOR: render298 pub fn render(299 &mut self,300 device: &wgpu::Device,301 queue: &wgpu::Queue,302 encoder: &mut wgpu::CommandEncoder,303 target: &wgpu::TextureView,304 size: [u32; 2],305 frame: RelightFrame,306 ) -> Result<(), String> {307 frame.validate()?;308 if size.contains(&0)309 || size310 .iter()311 .any(|&v| v > device.limits().max_texture_dimension_2d)312 {313 return Err("Invalid relight target size".into());314 }315 if self.hdr.as_ref().is_none_or(|v| v.2 != size) {316 let texture = device.create_texture(&wgpu::TextureDescriptor {317 label: Some("linear-radiance"),318 size: wgpu::Extent3d {319 width: size[0],320 height: size[1],321 depth_or_array_layers: 1,322 },323 mip_level_count: 1,324 sample_count: 1,325 dimension: wgpu::TextureDimension::D2,326 format: wgpu::TextureFormat::Rgba16Float,327 usage: wgpu::TextureUsages::RENDER_ATTACHMENT328 | wgpu::TextureUsages::TEXTURE_BINDING,329 view_formats: &[],330 });331 let view = texture.create_view(&Default::default());332 let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {333 label: None,334 layout: &self.hdr_layout,335 entries: &[wgpu::BindGroupEntry {336 binding: 0,337 resource: wgpu::BindingResource::TextureView(&view),338 }],339 });340 self.hdr = Some((view, bind, size));341 }342 let view = Mat4::look_at_rh(frame.eye(), Vec3::ZERO, Vec3::Y);343 let projection = Mat4::perspective_rh(344 45_f32.to_radians(),345 size[0] as f32 / size[1] as f32,346 0.05,347 50.0,348 );349 queue.write_buffer(350 &self.globals,351 0,352 bytemuck::bytes_of(&Globals {353 vp: (projection * view).to_cols_array_2d(),354 eye: frame.eye().extend(0.0).to_array(),355 light: frame.light().extend(frame.intensity).to_array(),356 material: [frame.roughness, frame.metallic, frame.mode as f32, 0.0],357 viewport: [358 size[0] as f32,359 size[1] as f32,360 if self.srgb { 1.0 } else { 0.0 },361 0.0,362 ],363 }),364 );365 queue.write_buffer(&self.splats, 0, bytemuck::cast_slice(&scene(frame, view)));366 let (hdr, texture_binding, _) = self.hdr.as_ref().unwrap();367 for (view, pipeline, display) in [368 (hdr, &self.splat_pipeline, false),369 (target, &self.display_pipeline, true),370 ] {371 let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {372 label: Some("relight-pass"),373 color_attachments: &[Some(wgpu::RenderPassColorAttachment {374 view,375 resolve_target: None,376 depth_slice: None,377 ops: wgpu::Operations {378 load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),379 store: wgpu::StoreOp::Store,380 },381 })],382 depth_stencil_attachment: None,383 timestamp_writes: None,384 occlusion_query_set: None,385 multiview_mask: None,386 });387 pass.set_pipeline(pipeline);388 pass.set_bind_group(0, &self.bindings, &[]);389 if display {390 pass.set_bind_group(1, texture_binding, &[]);391 pass.draw(0..3, 0..1);392 } else {393 pass.draw(0..6, 0..COUNT as u32);394 }395 }396 Ok(())397 }398 // END: render399}400401#[cfg(test)]402mod tests {403 use super::*;404 #[test]405 fn nonuniform_deformation_preserves_normal_tangent_orthogonality() {406 let frame = RelightFrame {407 time: 0.25,408 ..Default::default()409 };410 let s = surfel(Vec3::ONE.normalize(), frame);411 let n = Vec3::from_slice(&s.normal[..3]);412 assert!(n.dot(Vec3::from_slice(&s.u[..3])).abs() < 1e-6);413 assert!(n.dot(Vec3::from_slice(&s.v[..3])).abs() < 1e-6);414 let bad = surfel(415 Vec3::ONE.normalize(),416 RelightFrame {417 wrong_normals: true,418 ..frame419 },420 );421 assert!(422 Vec3::from_slice(&bad.normal[..3])423 .dot(Vec3::from_slice(&bad.u[..3]))424 .abs()425 > 0.01426 );427 }428 #[test]429 fn brdf_light_and_material_invariants() {430 let f = RelightFrame::default();431 let [d, s] = reference_brdf(Vec3::Z, Vec3::Z, Vec3::Z, f);432 assert!((d.x - 0.96 * BASE.x / std::f32::consts::PI * f.intensity).abs() < 1e-6);433 assert!(s.min_element() > 0.0);434 assert_eq!(435 reference_brdf(Vec3::Z, Vec3::Z, Vec3::NEG_Z, f),436 [Vec3::ZERO; 2]437 );438 assert_eq!(439 reference_brdf(440 Vec3::Z,441 Vec3::Z,442 Vec3::Z,443 RelightFrame { metallic: 1.0, ..f }444 )[0],445 Vec3::ZERO446 );447 let doubled = reference_brdf(448 Vec3::Z,449 Vec3::Z,450 Vec3::Z,451 RelightFrame {452 intensity: f.intensity * 2.0,453 ..f454 },455 );456 assert!((doubled[1] - s * 2.0).length() < 1e-5);457 assert_eq!(std::mem::size_of::<Globals>(), 128);458 assert_eq!(std::mem::size_of::<Splat>(), 96);459 assert!(RelightFrame {460 roughness: 0.0,461 ..f462 }463 .validate()464 .is_err());465 assert!(RelightFrame {466 intensity: f32::NAN,467 ..f468 }469 .validate()470 .is_err());471 }472}473