raytracers/rtiow/src/renderer.rs

use std;
use std::fmt;
use std::path::Path;
use std::path::PathBuf;
use std::str;
use std::sync;
use std::thread;

use crossbeam_channel as channel;
use image;
use image::RgbImage;
use num_cpus;
use rand;
use rand::Rng;

use camera::Camera;
use hitable::Hit;
use ray::Ray;
use scenes;
use vec3::Vec3;

#[derive(Debug)]
pub enum Model {
    Bench,
    Book,
    Tutorial,
    Cube,
    BVH,
    Test,
}

impl Model {
    pub fn scene(&self, opt: &Opt) -> Scene {
        match self {
            Model::Book => scenes::book::new(&opt),
            Model::Bench => scenes::bench::new(&opt),
            Model::Cube => scenes::cube::new(&opt),
            Model::Tutorial => scenes::tutorial::new(&opt),
            Model::BVH => scenes::bvh::new(&opt),
            Model::Test => scenes::test::new(&opt),
        }
    }
}

#[derive(Debug)]
pub struct ModelParseError(String);

impl fmt::Display for ModelParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "unknown model enum type '{}'", self.0)
    }
}

impl str::FromStr for Model {
    type Err = ModelParseError;
    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s {
            "bench" => Ok(Model::Bench),
            "book" => Ok(Model::Book),
            "tutorial" => Ok(Model::Tutorial),
            "cube" => Ok(Model::Cube),
            "bvh" => Ok(Model::BVH),
            "test" => Ok(Model::Test),
            _ => Err(ModelParseError(s.to_owned())),
        }
    }
}

#[derive(Debug, StructOpt)]
#[structopt(name = "tracer", about = "An experimental ray tracer.")]
pub struct Opt {
    /// Image width
    #[structopt(short = "w", long = "width", default_value = "1024")]
    pub width: usize,
    /// Image height
    #[structopt(short = "h", long = "height", default_value = "1024")]
    pub height: usize,
    /// Sub-samples per pixel
    #[structopt(short = "s", long = "subsample", default_value = "1")]
    pub subsamples: usize,
    #[structopt(long = "model", default_value = "test")]
    pub model: Model,
    #[structopt(long = "pprof")]
    pub pprof: bool,
    /// Use acceleration data structure, may be BVH or kd-tree depending on scene.
    #[structopt(long = "use_accel")]
    pub use_accel: bool,

    /// Output directory
    #[structopt(parse(from_os_str), default_value = "/tmp/tracer")]
    pub output: PathBuf,
}

pub struct Scene {
    pub world: Box<Hit>,
    pub camera: Camera,
    pub subsamples: usize,
    pub width: usize,
    pub height: usize,
}

fn color(r: Ray, world: &Hit, depth: usize) -> Vec3 {
    if let Some(rec) = world.hit(r, 0.001, std::f32::MAX) {
        let scatter_response = rec.material.scatter(&r, &rec);
        if depth < 50 && scatter_response.reflected {
            return scatter_response.attenutation
                * color(scatter_response.scattered, world, depth + 1);
        }
        return Default::default();
    }
    // No hit, choose color from background.
    let unit_direction = r.direction.unit_vector();
    let t = 0.5 * (unit_direction.y + 1.);
    Vec3::new(1., 1., 1.) * (1. - t) + Vec3::new(0.5, 0.7, 1.) * t
}

fn trace_pixel(x: usize, y: usize, scene: &Scene) -> Vec3 {
    let mut rng = rand::thread_rng();
    let u = (rng.gen_range::<f32>(0., 1.) + x as f32) / scene.width as f32;
    let v = (rng.gen_range::<f32>(0., 1.) + y as f32) / scene.height as f32;
    let ray = scene.camera.get_ray(u, v);
    color(ray, scene.world.as_ref(), 0)
}

fn render_worker(
    tid: usize,
    scene: &Scene,
    input_chan: channel::Receiver<usize>,
    output_chan: channel::Sender<(usize, Vec<Vec3>)>,
) {
    for subsample in input_chan {
        let mut pixel_data: Vec<Vec3> = Vec::with_capacity(scene.width * scene.height);
        for y in 0..scene.height {
            for x in 0..scene.width {
                let p = trace_pixel(x, y, scene);
                pixel_data.push(p);
            }
        }
        output_chan.send((subsample, pixel_data));
    }
    info!("Shutting down worker {}", tid);
}

pub fn render(scene: Scene, output_dir: &Path) -> std::result::Result<(), std::io::Error> {
    let (seq_tx, seq_rx) = channel::unbounded();
    let (pixel_data_tx, pixel_data_rx) = channel::unbounded();

    let scene = sync::Arc::new(scene);
    for i in 0..num_cpus::get() {
        let s = sync::Arc::clone(&scene);
        let seq_rx = seq_rx.clone();
        let pixel_data_tx = pixel_data_tx.clone();
        thread::spawn(move || {
            render_worker(i, &s, seq_rx, pixel_data_tx);
        });
    }
    drop(seq_rx);
    drop(pixel_data_tx);

    (1..=scene.subsamples).for_each(|idx| seq_tx.send(idx));
    drop(seq_tx);

    let mut acc_count = 0;
    let mut acc: Vec<Vec3> = Vec::with_capacity(scene.width * scene.height);
    for _ in 0..(scene.width * scene.height) {
        acc.push(Default::default());
    }

    let mut img = RgbImage::new(scene.width as u32, scene.height as u32);
    for (_subsample, pixel_data) in pixel_data_rx {
        acc_count += 1;
        pixel_data.iter().enumerate().for_each(|(idx, p)| {
            let x = idx % scene.width;
            let y = idx / scene.width;
            let y_inv = scene.height - y - 1;
            acc[idx] = acc[idx] + *p;

            // Gamma correct, use gamma 2 correction, which is 1/gamma where gamma=2 which is 1/2 or
            // sqrt.
            let col = acc[idx] / acc_count as f32;
            let col = Vec3::new(col[0].sqrt(), col[1].sqrt(), col[2].sqrt());
            img.put_pixel(
                x as u32,
                y_inv as u32,
                image::Rgb([
                    (col[0] * 255.).min(255.) as u8,
                    (col[1] * 255.).min(255.) as u8,
                    (col[2] * 255.).min(255.) as u8,
                ]),
            );
        });
        let path = output_dir.join(format!("iteration{:05}.png", acc_count));
        trace!(target: "renderer", "Saving {}", path.to_string_lossy());
        img.save(&path)
            .unwrap_or_else(|_| panic!("Failed save {}", path.to_string_lossy()));
    }
    let path = output_dir.join("final.png");
    // Write the contents of this image to the Writer in PNG format.
    trace!(target: "renderer", "Saving {}", path.to_string_lossy());
    img.save(path)
}