pbrt/src/integrators/pipeline.rs
2026-01-18 16:29:27 +00:00

257 lines
8.1 KiB
Rust

use crate::core::image::Image;
use crate::core::{options::PBRTOptions, sampler::get_camera_sample};
use crate::spectra::get_spectra_context;
use indicatif::{ProgressBar, ProgressStyle};
use std::io::Write;
use std::path::Path;
use super::*;
struct PbrtProgress {
bar: ProgressBar,
}
impl PbrtProgress {
fn new(total_work: u64, description: &str, quiet: bool) -> Self {
if quiet {
return Self {
bar: ProgressBar::hidden(),
};
}
let bar = ProgressBar::new(total_work);
bar.set_style(
ProgressStyle::default_bar()
.template("[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}")
.unwrap()
.progress_chars("=>-"),
);
bar.set_message(description.to_string());
Self { bar }
}
fn update(&self, amount: u64) {
self.bar.inc(amount);
}
fn done(&self) {
self.bar.finish_with_message("Done");
}
fn elapsed_seconds(&self) -> f32 {
self.bar.elapsed().as_secs_f32()
}
}
fn generate_tiles(bounds: Bounds2i) -> Vec<Bounds2i> {
let mut tiles = Vec::new();
const TILE_SIZE: i32 = 16;
for y in (bounds.p_min.y()..bounds.p_max.y()).step_by(TILE_SIZE as usize) {
for x in (bounds.p_min.x()..bounds.p_max.x()).step_by(TILE_SIZE as usize) {
let p_min = Point2i::new(x, y);
let p_max = Point2i::new(
(x + TILE_SIZE).min(bounds.p_max.x()),
(y + TILE_SIZE).min(bounds.p_max.y()),
);
tiles.push(Bounds2i::from_points(p_min, p_max));
}
}
tiles
}
pub fn render<T>(
integrator: &T,
_base: &IntegratorBase,
camera: &Camera,
sampler_prototype: &Sampler,
arena: &mut Arena,
) where
T: RayIntegratorTrait,
{
let options = get_options();
if let Some((p_pixel, sample_index)) = options.debug_start {
let s_index = sample_index as usize;
let mut tile_sampler = sampler_prototype.clone();
tile_sampler.start_pixel_sample(p_pixel, s_index, None);
evaluate_pixel_sample(
integrator,
camera,
&mut tile_sampler,
p_pixel,
s_index,
arena,
);
return;
}
let pixel_bounds = camera.get_film().pixel_bounds();
let spp = sampler_prototype.samples_per_pixel();
let total_work = (pixel_bounds.area() as u64) * (spp as u64);
let progress = PbrtProgress::new(total_work, "Rendering", options.quiet);
let mut wave_start = 0;
let mut wave_end = 1;
let mut next_wave_size = 1;
let mut reference_image: Option<Image> = None;
let mut mse_out_file: Option<std::fs::File> = None;
if let Some(ref_path) = &options.mse_reference_image {
let image_and_metadata =
Image::read(Path::new(&ref_path), None).expect("Could not load image");
let image = image_and_metadata.image;
let metadata = image_and_metadata.metadata;
let resolution = image.resolution();
// reference_image = Some(image);
let mse_pixel_bounds = metadata
.pixel_bounds
.unwrap_or_else(|| Bounds2i::from_points(Point2i::new(0, 0), resolution));
if !mse_pixel_bounds.overlaps(&pixel_bounds) {
panic!("Output pixel bounds dont fit inside the reference image");
}
let crop_p_min = Point2i::from(pixel_bounds.p_min - mse_pixel_bounds.p_min);
let crop_p_max = Point2i::from(pixel_bounds.p_max - mse_pixel_bounds.p_min);
let crop_bounds = Bounds2i::from_points(crop_p_min, crop_p_max);
let cropped_image = image.crop(crop_bounds).clone();
let cropped_resolution = cropped_image.resolution();
let expected_res = Point2i::new(
pixel_bounds.p_max.x() - pixel_bounds.p_min.x(),
pixel_bounds.p_max.y() - pixel_bounds.p_min.y(),
);
reference_image = Some(cropped_image);
assert_eq!(
cropped_resolution, expected_res,
"Cropped reference image resolution mismatch"
);
if let Some(out_path) = &options.mse_reference_output {
mse_out_file = Some(
std::fs::File::create(out_path)
.expect(&format!("Failed to create MSE output file: {}", out_path)),
);
}
}
let tiles = generate_tiles(pixel_bounds);
while wave_start < spp {
tiles.par_iter().for_each(|tile_bounds| {
let mut sampler = sampler_prototype.clone();
for p_pixel in tile_bounds {
for sample_index in wave_start..wave_end {
sampler.start_pixel_sample(*p_pixel, sample_index, None);
evaluate_pixel_sample(
integrator,
camera,
&mut sampler,
*p_pixel,
sample_index,
arena,
);
}
}
let work_done = (tile_bounds.area() as u64) * ((wave_end - wave_start) as u64);
progress.update(work_done);
});
wave_start = wave_end;
wave_end = (wave_end + next_wave_size).min(spp);
if reference_image.is_none() {
next_wave_size = (2 * next_wave_size).min(64);
}
if wave_start == spp {
progress.done();
}
if wave_start == spp || options.write_partial_images || reference_image.is_some() {
let mut metadata = ImageMetadata {
render_time_seconds: Some(progress.elapsed_seconds()),
samples_per_pixel: Some(wave_start as i32),
..Default::default()
};
if wave_start == spp || options.write_partial_images {
camera.init_metadata(&mut metadata);
camera
.get_film()
.write_image(&metadata, 1.0 / wave_start as Float);
}
if let Some(ref_img) = &reference_image {
let splat_scale = 1.0 / (wave_start as Float);
let film_metadata = ImageMetadata::default();
let film_image = camera.get_film().get_image(&film_metadata, splat_scale);
let (mse_values, _mse_debug_img) =
film_image.mse(film_image.all_channels_desc(), ref_img, false);
let mse_avg = mse_values.average();
if let Some(file) = &mut mse_out_file {
writeln!(file, "{}, {:.9}", wave_start, mse_avg).ok();
file.flush().ok();
}
metadata.mse = Some(mse_avg);
}
}
}
}
pub fn evaluate_pixel_sample<T: RayIntegratorTrait>(
integrator: &T,
camera: &Camera,
sampler: &mut Sampler,
pixel: Point2i,
_sample_index: usize,
arena: &mut Arena,
) {
let mut lu = sampler.get1d();
if get_options().disable_wavelength_jitter {
lu = 0.5;
}
let lambda = camera.get_film().sample_wavelengths(lu);
let film = camera.get_film();
let filter = film.get_filter();
let camera_sample = get_camera_sample(sampler, pixel, filter);
if let Some(mut camera_ray) = camera.generate_ray_differential(camera_sample, &lambda) {
debug_assert!(camera_ray.ray.d.norm() > 0.999);
debug_assert!(camera_ray.ray.d.norm() < 1.001);
let ray_diff_scale = (sampler.samples_per_pixel() as Float).sqrt().max(0.125);
if get_options().disable_pixel_jitter {
camera_ray.ray.scale_differentials(ray_diff_scale);
}
let initialize_visible_surface = film.uses_visible_surface();
let (mut l, visible_surface) =
integrator.li(camera_ray.ray, &lambda, sampler, initialize_visible_surface, arena);
l *= camera_ray.weight;
let std_spectra = get_spectra_context();
if l.has_nans() || l.y(&lambda, std_spectra).is_infinite() {
l = SampledSpectrum::new(0.);
}
film.add_sample(
pixel,
l,
&lambda,
visible_surface.as_ref(),
camera_sample.filter_weight,
);
}
}