More PBR

kernels/src/bsdf.rs

@@ -1,4 +1,5 @@
-use spirv_std::glam::Vec3;
+use shared_structs::MaterialData;
+use spirv_std::{glam::{Vec3, Vec2, Vec4Swizzles}, Sampler, Image};
 #[allow(unused_imports)]
 use spirv_std::num_traits::Float;
 
@@ -269,3 +270,33 @@ impl BSDF for PBR {
         }
     }
 }
+
+pub fn get_pbr_bsdf(material: &MaterialData, uv: Vec2, atlas: &Image!(2D, type=f32, sampled), sampler: &Sampler) -> PBR {
+    let albedo = if material.has_albedo_texture() {
+        let scaled_uv = material.albedo.xy() + uv * material.albedo.zw();
+        let albedo = atlas.sample_by_lod(*sampler, scaled_uv, 0.0);
+        albedo.xyz()
+    } else {
+        material.albedo.xyz()
+    };
+    let roughness = if material.has_roughness_texture() {
+        let scaled_uv = material.roughness.xy() + uv * material.roughness.zw();
+        let roughness = atlas.sample_by_lod(*sampler, scaled_uv, 0.0);
+        roughness.x
+    } else {
+        material.roughness.x
+    };
+    let metallic = if material.has_metallic_texture() {
+        let scaled_uv = material.metallic.xy() + uv * material.metallic.zw();
+        let metallic = atlas.sample_by_lod(*sampler, scaled_uv, 0.0);
+        metallic.x
+    } else {
+        material.metallic.x
+    };
+
+    PBR {
+        albedo,
+        roughness,
+        metallic,
+    }
+}

kernels/src/lib.rs

@@ -6,7 +6,7 @@ use intersection::BVHReference;
 use shared_structs::{TracingConfig, BVHNode, MaterialData, PerVertexData};
 #[allow(unused_imports)]
 use spirv_std::num_traits::Float;
-use spirv_std::{glam, spirv};
+use spirv_std::{glam, spirv, Sampler, Image};
 
 mod bsdf;
 mod rng;
@@ -28,8 +28,8 @@ pub fn main_material(
     #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] nodes_buffer: &[BVHNode],
     #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] indirect_indices_buffer: &[u32],
     #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] material_data_buffer: &[MaterialData],
-    #[spirv(descriptor_set = 0, binding = 8)] sampler: &spirv_std::Sampler,
-    #[spirv(descriptor_set = 0, binding = 9)] albedo_atlas: &spirv_std::Image!(2D, type=f32, sampled),
+    #[spirv(descriptor_set = 0, binding = 8)] sampler: &Sampler,
+    #[spirv(descriptor_set = 0, binding = 9)] atlas: &Image!(2D, type=f32, sampled),
 ) {
     let index = (id.y * config.width + id.x) as usize;
 
@@ -82,24 +82,18 @@ pub fn main_material(
 
             let material_index = trace_result.triangle.w;
             let material = material_data_buffer[material_index as usize];
-            let albedo = if material.has_albedo_texture() {
-                let uv_a = per_vertex_buffer[trace_result.triangle.x as usize].uv0;
-                let uv_b = per_vertex_buffer[trace_result.triangle.y as usize].uv0;
-                let uv_c = per_vertex_buffer[trace_result.triangle.z as usize].uv0;
-                let uv = bary.x * uv_a + bary.y * uv_b + bary.z * uv_c;
-                let uv = Vec2::new(uv.x, uv.y);
-                let scaled_uv = material.albedo.xy() + uv * material.albedo.zw();
-                let albedo = albedo_atlas.sample_by_lod(*sampler, scaled_uv, 0.0);
-                albedo.xyz()
-            } else {
-                material.albedo.xyz()
-            };
-
-            let bsdf = bsdf::PBR {
-                albedo: albedo,
-                roughness: 0.5,
-                metallic: 0.3,
-            };
+            if material.emissive.xyz() != Vec3::ZERO {
+                // Emissives don't bounce light
+                output[index] += (throughput * material.emissive.xyz()).extend(1.0);
+                return;
+            }
+
+            let uv_a = per_vertex_buffer[trace_result.triangle.x as usize].uv0;
+            let uv_b = per_vertex_buffer[trace_result.triangle.y as usize].uv0;
+            let uv_c = per_vertex_buffer[trace_result.triangle.z as usize].uv0;
+            let uv = bary.x * uv_a + bary.y * uv_b + bary.z * uv_c;
+            let bsdf = bsdf::get_pbr_bsdf(&material, uv, atlas, sampler);
+
             let bsdf_sample = bsdf.sample(-ray_direction, norm, &mut rng_state);
             throughput *= bsdf_sample.spectrum / bsdf_sample.pdf;
 

shared_structs/src/lib.rs

@@ -29,10 +29,16 @@ impl Default for TracingConfig {
 
 #[repr(C)]
 #[derive(Copy, Clone, Pod, Zeroable, Default)]
-pub struct MaterialData {
-    pub albedo: Vec4, // either albedo color or atlas location
+pub struct MaterialData { // each Vec4 is either a color or an atlas location
+    pub emissive: Vec4,
+    pub albedo: Vec4,
+    pub roughness: Vec4,
+    pub metallic: Vec4,
+    pub normals: Vec4,
     has_albedo_texture: u32,
-    _padding: [u32; 3],
+    has_metallic_texture: u32,
+    has_roughness_texture: u32,
+    has_normal_texture: u32,
 }
 
 impl MaterialData {
@@ -43,6 +49,30 @@ impl MaterialData {
     pub fn set_has_albedo_texture(&mut self, has_albedo_texture: bool) {
         self.has_albedo_texture = if has_albedo_texture { 1 } else { 0 };
     }
+
+    pub fn has_metallic_texture(&self) -> bool {
+        self.has_metallic_texture != 0
+    }
+
+    pub fn set_has_metallic_texture(&mut self, has_metallic_texture: bool) {
+        self.has_metallic_texture = if has_metallic_texture { 1 } else { 0 };
+    }
+
+    pub fn has_roughness_texture(&self) -> bool {
+        self.has_roughness_texture != 0
+    }
+
+    pub fn set_has_roughness_texture(&mut self, has_roughness_texture: bool) {
+        self.has_roughness_texture = if has_roughness_texture { 1 } else { 0 };
+    }
+
+    pub fn has_normal_texture(&self) -> bool {
+        self.has_normal_texture != 0
+    }
+
+    pub fn set_has_normal_texture(&mut self, has_normal_texture: bool) {
+        self.has_normal_texture = if has_normal_texture { 1 } else { 0 };
+    }
 }
 
 #[repr(C)]

src/asset.rs

@@ -1,7 +1,7 @@
 use glam::{UVec4, Vec4, Mat4, Vec2};
 use gpgpu::{GpuBuffer, BufOps, GpuConstImage, primitives::pixels::{Rgba8UintNorm}, ImgOps};
 use image::DynamicImage;
-use russimp::{scene::{Scene, PostProcess::*}, node::Node, material::{DataContent, TextureType, Texture, Material}};
+use russimp::{scene::{Scene, PostProcess::*}, node::Node, material::{DataContent, TextureType, Texture, Material, PropertyTypeInfo}};
 use shared_structs::{MaterialData, PerVertexData};
 
 use crate::{bvh::{BVH, BVHBuilder}, trace::FW};
@@ -11,7 +11,7 @@ pub struct World<'fw> {
     pub index_buffer: GpuBuffer<'fw, UVec4>,
     pub bvh: BVH<'fw>,
 
-    pub albedo_atlas: GpuConstImage<'fw, Rgba8UintNorm>,
+    pub atlas: GpuConstImage<'fw, Rgba8UintNorm>,
     pub material_data_buffer: GpuBuffer<'fw, MaterialData>,    
 }
 
@@ -34,6 +34,14 @@ fn load_texture(material: &Material, texture_type: TextureType) -> Option<Dynami
     material.textures.get(&texture_type).and_then(|texture| convert_texture(&texture.borrow()))
 }
 
+fn load_float_array(material: &Material, name: &str) -> Option<Vec<f32>> {
+    let prop = material.properties.iter().find(|p| p.key == name)?;
+    match &prop.data {
+        PropertyTypeInfo::FloatArray(col) => Some(col.clone()),
+        _ => None
+    }
+}
+
 impl<'fw> World<'fw> {
     pub fn from_path(path: &str) -> Self {
         let blend = Scene::from_file(
@@ -101,21 +109,55 @@ impl<'fw> World<'fw> {
         // Gather material data
         let mut material_datas = vec![MaterialData::default(); blend.materials.len()];
 
-        let mut albedo_textures = Vec::new();
+        let mut textures = Vec::new();
         for (material_index, material) in blend.materials.iter().enumerate() {
             let current_material_data = &mut material_datas[material_index];
             if let Some(texture) = load_texture(&material, TextureType::Diffuse) {
-                albedo_textures.push(texture);
+                textures.push(texture);
                 current_material_data.set_has_albedo_texture(true);
             }
+            if let Some(texture) = load_texture(&material, TextureType::Metalness) {
+                textures.push(texture);
+                current_material_data.set_has_metallic_texture(true);
+            }
+            if let Some(texture) = load_texture(&material, TextureType::Roughness) {
+                textures.push(texture);
+                current_material_data.set_has_roughness_texture(true);
+            }
+            if let Some(texture) = load_texture(&material, TextureType::Normals) {
+                textures.push(texture);
+                current_material_data.set_has_normal_texture(true);
+            }
+            if let Some(col) = load_float_array(&material, "$clr.diffuse") {
+                current_material_data.albedo = Vec4::new(col[0], col[1], col[2], col[3]);
+            }
+            if let Some(col) = load_float_array(&material, "$clr.emissive") {
+                // HACK: Multiply by 20 since assimp 5.2.5 doesn't support emissive strength :(
+                current_material_data.emissive = Vec4::new(col[0], col[1], col[2], col[3]) * 5.0;
+            }
+            if let Some(col) = load_float_array(&material, "$mat.metallicFactor") {
+                current_material_data.metallic = Vec4::splat(col[0]);
+            }
+            if let Some(col) = load_float_array(&material, "$mat.roughnessFactor") {
+                current_material_data.roughness = Vec4::splat(col[0]);
+            }
         }
 
-        let (albedo_atlas_raw, mut albedo_sts) = crate::atlas::pack_textures(&albedo_textures, 1024, 1024);
-        let albedo_atlas = GpuConstImage::from_bytes(&FW, &albedo_atlas_raw.to_rgba8(), 1024, 1024);
+        let (atlas_raw, mut sts) = crate::atlas::pack_textures(&textures, 4096, 4096);
+        let atlas = GpuConstImage::from_bytes(&FW, &atlas_raw.to_rgba8(), 4096, 4096);
 
         for material_data in material_datas.iter_mut() {
             if material_data.has_albedo_texture() {
-                material_data.albedo = albedo_sts.remove(0); // TODO: Optimize this
+                material_data.albedo = sts.remove(0); // TODO: Optimize this
+            }
+            if material_data.has_metallic_texture() {
+                material_data.metallic = sts.remove(0);
+            }
+            if material_data.has_roughness_texture() {
+                material_data.roughness = sts.remove(0);
+            }
+            if material_data.has_normal_texture() {
+                material_data.normals = sts.remove(0);
             }
         }
 
@@ -144,7 +186,7 @@ impl<'fw> World<'fw> {
             per_vertex_buffer,
             index_buffer,
             bvh,
-            albedo_atlas,
+            atlas,
             material_data_buffer,
         }
     }

src/trace.rs

@@ -38,7 +38,7 @@ impl<'fw> PathTracingKernel<'fw> {
             .bind_buffer(&world.bvh.indirect_indices_buffer, GpuBufferUsage::ReadOnly)
             .bind_buffer(&world.material_data_buffer, GpuBufferUsage::ReadOnly)
             .bind_sampler(&sampler)
-            .bind_const_image(&world.albedo_atlas);
+            .bind_const_image(&world.atlas);
         let program = Program::new(&shader, "main_material").add_descriptor_set(bindings);
         let kernel = Kernel::new(&FW, program);