use bytemuck::{Pod, Zeroable};
use glam::Mat4;
use crate::render::VolumeActor;
#[repr(C)]
#[derive(Debug, Clone, Copy, Pod, Zeroable)]
struct VolumeUniforms {
mvp: [[f32; 4]; 4],
camera_pos: [f32; 3],
num_steps: f32,
volume_min: [f32; 3],
opacity_scale: f32,
volume_max: [f32; 3],
_pad: f32,
}
pub struct VolumePass {
pipeline: wgpu::RenderPipeline,
bind_group_layout: wgpu::BindGroupLayout,
uniform_buffer: wgpu::Buffer,
sampler: wgpu::Sampler,
lut_sampler: wgpu::Sampler,
}
impl VolumePass {
pub fn new(device: &wgpu::Device, color_format: wgpu::TextureFormat) -> Self {
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("volume shader"),
source: wgpu::ShaderSource::Wgsl(include_str!("volume_shader.wgsl").into()),
});
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("volume bgl"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: true },
view_dimension: wgpu::TextureViewDimension::D3,
multisampled: false,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 3,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: true },
view_dimension: wgpu::TextureViewDimension::D1,
multisampled: false,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 4,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
count: None,
},
],
});
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("volume pl"),
bind_group_layouts: &[&bind_group_layout],
push_constant_ranges: &[],
});
let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
label: Some("volume pipeline"),
layout: Some(&pipeline_layout),
vertex: wgpu::VertexState {
module: &shader,
entry_point: Some("vs_volume"),
buffers: &[],
compilation_options: Default::default(),
},
fragment: Some(wgpu::FragmentState {
module: &shader,
entry_point: Some("fs_volume"),
targets: &[Some(wgpu::ColorTargetState {
format: color_format,
blend: Some(wgpu::BlendState::ALPHA_BLENDING),
write_mask: wgpu::ColorWrites::ALL,
})],
compilation_options: Default::default(),
}),
primitive: wgpu::PrimitiveState {
topology: wgpu::PrimitiveTopology::TriangleList,
cull_mode: Some(wgpu::Face::Front), ..Default::default()
},
depth_stencil: None,
multisample: Default::default(),
multiview: None,
cache: None,
});
let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("volume uniforms"),
size: std::mem::size_of::<VolumeUniforms>() as u64,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
label: Some("volume sampler"),
mag_filter: wgpu::FilterMode::Linear,
min_filter: wgpu::FilterMode::Linear,
..Default::default()
});
let lut_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
label: Some("lut sampler"),
mag_filter: wgpu::FilterMode::Linear,
min_filter: wgpu::FilterMode::Linear,
address_mode_u: wgpu::AddressMode::ClampToEdge,
..Default::default()
});
Self { pipeline, bind_group_layout, uniform_buffer, sampler, lut_sampler }
}
pub fn render(
&self,
device: &wgpu::Device,
queue: &wgpu::Queue,
encoder: &mut wgpu::CommandEncoder,
color_view: &wgpu::TextureView,
volume: &VolumeActor,
mvp: Mat4,
camera_pos: [f32; 3],
) {
let [nx, ny, nz] = volume.dimensions;
let volume_texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("volume 3d"),
size: wgpu::Extent3d { width: nx as u32, height: ny as u32, depth_or_array_layers: nz as u32 },
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D3,
format: wgpu::TextureFormat::R32Float,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
let smin = volume.scalar_range[0];
let smax = volume.scalar_range[1];
let srange = smax - smin;
let normalized: Vec<f32> = volume.scalars.iter().map(|&v| {
if srange.abs() > 1e-15 { ((v - smin) / srange) as f32 } else { 0.5 }
}).collect();
queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: &volume_texture,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
bytemuck::cast_slice(&normalized),
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(nx as u32 * 4),
rows_per_image: Some(ny as u32),
},
wgpu::Extent3d { width: nx as u32, height: ny as u32, depth_or_array_layers: nz as u32 },
);
let volume_view = volume_texture.create_view(&Default::default());
let lut_rgba = volume.transfer_function.to_lut_rgba();
let lut_f32: Vec<f32> = lut_rgba.iter().map(|&b| b as f32 / 255.0).collect();
let lut_texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("lut 1d"),
size: wgpu::Extent3d { width: 256, height: 1, depth_or_array_layers: 1 },
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D1,
format: wgpu::TextureFormat::Rgba32Float,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: &lut_texture,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
bytemuck::cast_slice(&lut_f32),
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(256 * 16),
rows_per_image: Some(1),
},
wgpu::Extent3d { width: 256, height: 1, depth_or_array_layers: 1 },
);
let lut_view = lut_texture.create_view(&Default::default());
let vol_min = [
volume.origin[0] as f32,
volume.origin[1] as f32,
volume.origin[2] as f32,
];
let vol_max = [
(volume.origin[0] + (nx - 1) as f64 * volume.spacing[0]) as f32,
(volume.origin[1] + (ny - 1) as f64 * volume.spacing[1]) as f32,
(volume.origin[2] + (nz - 1) as f64 * volume.spacing[2]) as f32,
];
let uniforms = VolumeUniforms {
mvp: mvp.to_cols_array_2d(),
camera_pos,
num_steps: volume.num_steps as f32,
volume_min: vol_min,
opacity_scale: volume.opacity_scale as f32,
volume_max: vol_max,
_pad: 0.0,
};
queue.write_buffer(&self.uniform_buffer, 0, bytemuck::bytes_of(&uniforms));
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("volume bg"),
layout: &self.bind_group_layout,
entries: &[
wgpu::BindGroupEntry { binding: 0, resource: self.uniform_buffer.as_entire_binding() },
wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(&volume_view) },
wgpu::BindGroupEntry { binding: 2, resource: wgpu::BindingResource::Sampler(&self.sampler) },
wgpu::BindGroupEntry { binding: 3, resource: wgpu::BindingResource::TextureView(&lut_view) },
wgpu::BindGroupEntry { binding: 4, resource: wgpu::BindingResource::Sampler(&self.lut_sampler) },
],
});
{
let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
label: Some("volume pass"),
color_attachments: &[Some(wgpu::RenderPassColorAttachment {
view: color_view,
resolve_target: None,
ops: wgpu::Operations {
load: wgpu::LoadOp::Load,
store: wgpu::StoreOp::Store,
},
})],
depth_stencil_attachment: None,
..Default::default()
});
pass.set_pipeline(&self.pipeline);
pass.set_bind_group(0, &bind_group, &[]);
pass.draw(0..36, 0..1); }
}
}