From 278da46531d3d03869636ba8edb4da2a6bef2b25 Mon Sep 17 00:00:00 2001 From: Daniel Oom Date: Thu, 4 Jul 2024 13:57:27 +0200 Subject: [PATCH] Optimize buffer to egui::Image conversion --- Cargo.lock | 2 + pathtracer-cli/Cargo.toml | 1 + pathtracer-cli/src/main.rs | 7 +- pathtracer-gui/src/gui.rs | 54 +++++++------ pathtracer-gui/src/workers.rs | 2 +- raytracer-cli/Cargo.toml | 1 + raytracer-cli/src/main.rs | 6 +- tracing/Cargo.toml | 2 +- tracing/src/image_buffer.rs | 147 ++++++++++++++++++---------------- tracing/src/pathtracer.rs | 14 ++-- tracing/src/raytracer.rs | 11 +-- 11 files changed, 134 insertions(+), 113 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 89c0c8fb..f5531b64 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1700,6 +1700,7 @@ name = "pathtracer-cli" version = "1.0.0" dependencies = [ "clap", + "image 0.25.1", "kdtree", "rand", "scene", @@ -1896,6 +1897,7 @@ name = "raytracer-cli" version = "1.0.0" dependencies = [ "clap", + "image 0.25.1", "kdtree", "scene", "tracing 1.0.0", diff --git a/pathtracer-cli/Cargo.toml b/pathtracer-cli/Cargo.toml index cb9a39ae..34cc7f14 100644 --- a/pathtracer-cli/Cargo.toml +++ b/pathtracer-cli/Cargo.toml @@ -9,6 +9,7 @@ ray_logging = ["tracing/ray_logging"] [dependencies] clap = { version = "4.5.4", features = ["derive"] } +image = { version = "0.25.1", default-features = false, features = ["png"] } kdtree = { version = "1.0.0", path = "../kdtree" } rand = { version = "0.8.5", default-features = false, features = ["std", "small_rng"] } scene = { version = "1.0.0", path = "../scene" } diff --git a/pathtracer-cli/src/main.rs b/pathtracer-cli/src/main.rs index 194fe339..f4cef190 100644 --- a/pathtracer-cli/src/main.rs +++ b/pathtracer-cli/src/main.rs @@ -1,4 +1,5 @@ use clap::Parser; +use image::ImageFormat; use kdtree::{ build::build_kdtree, build_sah::{self, SahKdTreeBuilder}, @@ -218,9 +219,9 @@ fn main() { ); println!("Writing {}...", args.output.display()); - (buffer / total_iterations as f32) - .gamma_correct() - .save_png(&args.output) + buffer + .as_rgb_image(total_iterations as u16) + .save_with_format(&args.output, ImageFormat::Png) .unwrap(); }); } diff --git a/pathtracer-gui/src/gui.rs b/pathtracer-gui/src/gui.rs index 45a259ee..af5a0bd3 100644 --- a/pathtracer-gui/src/gui.rs +++ b/pathtracer-gui/src/gui.rs @@ -10,7 +10,7 @@ use std::{ use egui::Vec2; use nalgebra::Vector3; use scene::camera::{Camera, Pinhole}; -use tracing::pathtracer::Pathtracer; +use tracing::{image_buffer::ImageBuffer, pathtracer::Pathtracer}; use crate::workers::{spawn_worker, RenderResult}; @@ -19,7 +19,33 @@ struct RenderState { duration: time::Duration, } -fn spawn_receiver_thread( +fn convert(iterations: u16, buffer: ImageBuffer) -> egui::ColorImage { + let size = [buffer.width as usize, buffer.height as usize]; + let pixels = buffer + .into_rgba_iter(iterations) + .map(|p| egui::Color32::from_rgba_premultiplied(p[0], p[1], p[2], p[3])) + .collect(); + egui::ColorImage { size, pixels } +} + +fn receiver_thread( + rx: Receiver, + render_ptr: Arc>>, + image_ptr: Arc>>, + ctx: egui::Context, +) { + while let Ok(result) = rx.recv() { + let image = convert(result.iteration, result.image); + render_ptr.lock().unwrap().replace(RenderState { + iteration: result.iteration, + duration: result.duration, + }); + image_ptr.lock().unwrap().replace(image); + ctx.request_repaint(); + } +} + +fn spawn_receiver( rx: Receiver, render_ptr: Arc>>, image_ptr: Arc>>, @@ -27,27 +53,7 @@ fn spawn_receiver_thread( ) -> Result, std::io::Error> { std::thread::Builder::new() .name("GUI receiver".to_string()) - .spawn(move || { - while let Ok(result) = rx.recv() { - render_ptr.lock().unwrap().replace(RenderState { - iteration: result.iteration, - duration: result.duration, - }); - image_ptr - .lock() - .unwrap() - .replace(egui::ColorImage::from_rgba_unmultiplied( - [ - result.image.width() as usize, - result.image.height() as usize, - ], - &(result.image / result.iteration as f32) - .gamma_correct() - .to_rgba8(), - )); - ctx.request_repaint(); - } - }) + .spawn(move || receiver_thread(rx, render_ptr, image_ptr, ctx)) } pub(crate) struct PathtracerGui { @@ -95,7 +101,7 @@ impl PathtracerGui { self.worker_tx = Some(tx); self.worker_rx = Some(rx); self.receiver_thread = Some( - spawn_receiver_thread( + spawn_receiver( self.worker_rx.take().unwrap(), Arc::clone(&self.render), Arc::clone(&self.image), diff --git a/pathtracer-gui/src/workers.rs b/pathtracer-gui/src/workers.rs index a75dfd63..c414f608 100644 --- a/pathtracer-gui/src/workers.rs +++ b/pathtracer-gui/src/workers.rs @@ -114,7 +114,7 @@ fn worker_loop( let (duration, buffer) = measure(|| render_subdivided(&pathtracer, &pinhole, pinhole.size() / 4)); - combined_buffer.add_mut(&buffer); + combined_buffer += buffer; iteration += 1; let _ = tx.send(RenderResult { iteration, diff --git a/raytracer-cli/Cargo.toml b/raytracer-cli/Cargo.toml index 455c2679..1d31d7ce 100644 --- a/raytracer-cli/Cargo.toml +++ b/raytracer-cli/Cargo.toml @@ -5,6 +5,7 @@ edition = "2021" [dependencies] clap = { version = "4.5.4", features = ["derive"] } +image = { version = "0.25.1", default-features = false, features = ["png"] } kdtree = { version = "1.0.0", path = "../kdtree" } scene = { version = "1.0.0", path = "../scene" } tracing = { version = "1.0.0", path = "../tracing" } diff --git a/raytracer-cli/src/main.rs b/raytracer-cli/src/main.rs index f90f65d5..d2678fd6 100644 --- a/raytracer-cli/src/main.rs +++ b/raytracer-cli/src/main.rs @@ -1,4 +1,5 @@ use clap::Parser; +use image::ImageFormat; use kdtree::{build::build_kdtree, build_sah, build_sah::SahKdTreeBuilder}; use scene::{camera::Pinhole, Scene}; use std::{fmt::Display, str::FromStr}; @@ -91,5 +92,8 @@ fn main() { raytracer.render(&mut buffer); println!("Writing {}...", args.output.display()); - buffer.gamma_correct().save_png(&args.output).unwrap(); + buffer + .as_rgb_image(1) + .save_with_format(&args.output, ImageFormat::Png) + .unwrap(); } diff --git a/tracing/Cargo.toml b/tracing/Cargo.toml index d838ac7b..c1c8b0e8 100644 --- a/tracing/Cargo.toml +++ b/tracing/Cargo.toml @@ -9,7 +9,7 @@ ray_logging = [] [dependencies] assert_approx_eq = "1.1.0" geometry = { version = "1.0.0", path = "../geometry" } -image = { version = "0.25.1", default-features = false, features = ["png"] } +image = { version = "0.25.1", default-features = false } kdtree = { version = "1.0.0", path = "../kdtree" } nalgebra = { version = "0.32.5", default-features = false, features = ["alloc"] } rand = { version = "0.8.5", default-features = false, features = ["small_rng", "std"] } diff --git a/tracing/src/image_buffer.rs b/tracing/src/image_buffer.rs index 43b843bc..a7601133 100644 --- a/tracing/src/image_buffer.rs +++ b/tracing/src/image_buffer.rs @@ -1,107 +1,116 @@ -use std::{ - ops::{Add, Div}, - path::Path, -}; +use std::ops::{Add, AddAssign, Index, IndexMut}; -use nalgebra::{DMatrix, Vector2, Vector3}; +use nalgebra::{Vector2, Vector3}; #[derive(Clone)] -pub struct ImageBuffer(DMatrix>); +pub struct ImageBuffer { + pub width: u32, + pub height: u32, + pub pixels: Vec>, +} -fn gamma_correct(x: f32) -> f32 { +fn gamma_correct(x: Vector3) -> Vector3 { const GAMMA_POWER: f32 = 1.0 / 2.2; - 1.0f32.min(x.powf(GAMMA_POWER)) + x.zip_map(&Vector3::from_element(GAMMA_POWER), |b, e| b.powf(e)) + .inf(&Vector3::from_element(1.0)) } impl ImageBuffer { + #[inline] pub fn new(width: u32, height: u32) -> Self { - ImageBuffer(DMatrix::from_element( - height as usize, - width as usize, - Vector3::zeros(), - )) - } - - pub fn width(&self) -> u32 { - self.0.ncols() as u32 - } - - pub fn height(&self) -> u32 { - self.0.nrows() as u32 - } - - pub fn add_mut(&mut self, rhs: &Self) { - self.0 - .iter_mut() - .zip(rhs.0.iter()) - .for_each(|(a, b)| *a += *b); + Self { + width, + height, + pixels: [Vector3::zeros()].repeat((width * height) as usize), + } } - pub fn add_pixel_mut(&mut self, x: u32, y: u32, p: Vector3) { - self.0[(y as usize, x as usize)] += p; + #[inline] + fn index(&self, idx: Vector2) -> usize { + (self.width * idx.y + idx.x) as usize } + #[inline] pub fn add_at(mut self, at: Vector2, rhs: &Self) -> Self { - for y in 0..rhs.0.nrows() { - for x in 0..rhs.0.ncols() { - self.0[(at.y as usize + y, at.x as usize + x)] += rhs.0[(y, x)]; + debug_assert!(at.x + rhs.width <= self.width && at.y + rhs.height <= self.height); + for y in 0..rhs.height { + for x in 0..rhs.width { + self[Vector2::new(at.x + x, at.y + y)] += rhs[Vector2::new(x, y)]; } } self } - pub fn gamma_correct(self) -> Self { - ImageBuffer(self.0.map(|p| p.map(gamma_correct))) - } - - pub fn to_rgb8(self) -> Vec { - self.0 - .transpose() + pub fn to_rgb8(self, iterations: u16) -> Vec { + let iterations_inv = 1.0 / iterations as f32; + self.pixels .into_iter() - .flat_map(|p| { - [ - (p.x * 255.0).round() as u8, - (p.y * 255.0).round() as u8, - (p.z * 255.0).round() as u8, - ] + .flat_map(|p| -> [u8; 3] { + (gamma_correct(p * iterations_inv) * 255.0) + .map(|c| c.round() as u8) + .into() }) .collect() } - pub fn to_rgba8(self) -> Vec { - self.0 - .transpose() - .into_iter() - .flat_map(|p| { - [ - (p.x * 255.0).round() as u8, - (p.y * 255.0).round() as u8, - (p.z * 255.0).round() as u8, - u8::MAX, - ] - }) - .collect() + #[inline] + pub fn into_rgba_iter(self, iterations: u16) -> impl Iterator { + let iterations_inv = 1.0 / iterations as f32; + self.pixels.into_iter().map(move |p| -> [u8; 4] { + let p = (gamma_correct(p * iterations_inv) * 255.0).map(|c| c.round() as u8); + [p.x, p.y, p.z, u8::MAX] + }) } - pub fn save_png(self, path: &Path) -> Result<(), image::ImageError> { - image::RgbImage::from_raw(self.width(), self.height(), self.to_rgb8()) - .unwrap() - .save_with_format(path, image::ImageFormat::Png) + #[inline] + pub fn as_rgb_image(self, iterations: u16) -> image::RgbImage { + image::RgbImage::from_raw(self.width, self.height, self.to_rgb8(iterations)).unwrap() + } +} + +impl Index> for ImageBuffer { + type Output = Vector3; + + #[inline] + fn index(&self, index: Vector2) -> &Self::Output { + &self.pixels[self.index(index)] + } +} + +impl IndexMut> for ImageBuffer { + #[inline] + fn index_mut(&mut self, index: Vector2) -> &mut Self::Output { + let index = self.index(index); + &mut self.pixels[index] } } impl Add for ImageBuffer { type Output = Self; + #[inline] fn add(self, rhs: Self) -> Self::Output { - Self(self.0 + &rhs.0) + debug_assert!(self.width == rhs.width && self.height == rhs.height); + Self { + width: self.width, + height: self.height, + pixels: self + .pixels + .into_iter() + .zip(rhs.pixels.iter()) + .map(|(a, b)| a + b) + .collect::>(), + } } } -impl Div for ImageBuffer { - type Output = Self; - - fn div(self, rhs: f32) -> Self::Output { - Self(self.0.map(|p| p / rhs)) +impl AddAssign for ImageBuffer { + #[inline] + fn add_assign(&mut self, rhs: Self) { + debug_assert!(self.width == rhs.width && self.height == rhs.height); + self.pixels + .iter_mut() + .zip(rhs.pixels.iter()) + .for_each(|(a, b)| *a += *b); } } diff --git a/tracing/src/pathtracer.rs b/tracing/src/pathtracer.rs index db583d81..4e901989 100644 --- a/tracing/src/pathtracer.rs +++ b/tracing/src/pathtracer.rs @@ -166,13 +166,10 @@ impl Pathtracer { rng: &mut SmallRng, buffer: &mut ImageBuffer, ) { - for y in 0..buffer.height() { - for x in 0..buffer.width() { - buffer.add_pixel_mut( - x, - y, - self.render_pixel(pinhole, Vector2::new(x, y), ray_logger, rng), - ); + for y in 0..buffer.height { + for x in 0..buffer.width { + let pixel = Vector2::new(x, y); + buffer[pixel] += self.render_pixel(pinhole, pixel, ray_logger, rng); } } } @@ -189,8 +186,7 @@ impl Pathtracer { for sub_y in 0..sub_size.y { for sub_x in 0..sub_size.x { let pixel = sub_start + Vector2::new(sub_x, sub_y); - let color = self.render_pixel(pinhole, pixel, ray_logger, rng); - buffer.add_pixel_mut(pixel.x, pixel.y, color); + buffer[pixel] += self.render_pixel(pinhole, pixel, ray_logger, rng); } } } diff --git a/tracing/src/raytracer.rs b/tracing/src/raytracer.rs index 98fe6e16..c6bea1de 100644 --- a/tracing/src/raytracer.rs +++ b/tracing/src/raytracer.rs @@ -65,13 +65,14 @@ impl Raytracer { } pub fn render(&self, buffer: &mut ImageBuffer) { - let buffer_size = Vector2::new(buffer.width() as f32, buffer.height() as f32); - for y in 0..buffer.height() { - for x in 0..buffer.width() { - let pixel_center = Vector2::new(x as f32, y as f32) + Vector2::new(0.5, 0.5); + let buffer_size = Vector2::new(buffer.width as f32, buffer.height as f32); + for y in 0..buffer.height { + for x in 0..buffer.width { + let pixel = Vector2::new(x, y); + let pixel_center = Vector2::new(pixel.x as f32 + 0.5, pixel.y as f32 + 0.5); let scene_direction = pixel_center.component_div(&buffer_size); let ray = self.camera.ray(scene_direction.x, scene_direction.y); - buffer.add_pixel_mut(x, y, self.trace_ray(&ray)); + buffer[pixel] += self.trace_ray(&ray); } } }