raytracers/rtiow/src/renderer.rs

use std;
use std::fmt;
use std::io;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::str;
use std::sync;
use std::sync::mpsc::sync_channel;
use std::sync::mpsc::Receiver;
use std::sync::mpsc::SyncSender;
use std::sync::Arc;
use std::sync::Mutex;
use std::thread;

use image;
use image::RgbImage;
use num_cpus;
use rand;
use rand::Rng;

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

lazy_static! {
    static ref RAY_COUNTER: prometheus::CounterVec =
        register_counter_vec!("rays", "Number of rays fired", &["level"]).unwrap();
}

#[derive(Debug)]
pub enum Model {
    Bench,
    Book,
    Tutorial,
    BVH,
    Test,
    CornellBox,
    CornellSmoke,
    PerlinDebug,
    Final,
    Mandelbrot,
}

impl Model {
    pub fn scene(&self, opt: &Opt) -> Scene {
        match self {
            Model::Book => scenes::book::new(&opt),
            Model::Bench => scenes::bench::new(&opt),
            Model::Tutorial => scenes::tutorial::new(&opt),
            Model::BVH => scenes::bvh::new(&opt),
            Model::Test => scenes::test::new(&opt),
            Model::CornellBox => scenes::cornell_box::new(&opt),
            Model::CornellSmoke => scenes::cornell_smoke::new(&opt),
            Model::PerlinDebug => scenes::perlin_debug::new(&opt),
            Model::Final => scenes::final_scene::new(&opt),
            Model::Mandelbrot => scenes::mandelbrot::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),
            "bvh" => Ok(Model::BVH),
            "test" => Ok(Model::Test),
            "cornell_box" => Ok(Model::CornellBox),
            "cornell_smoke" => Ok(Model::CornellSmoke),
            "perlin_debug" => Ok(Model::PerlinDebug),
            "final" => Ok(Model::Final),
            "mandelbrot" => Ok(Model::Mandelbrot),
            _ => Err(ModelParseError(s.to_owned())),
        }
    }
}

impl std::string::ToString for Model {
    fn to_string(&self) -> String {
        match self {
            Model::Bench => "bench".to_string(),
            Model::Book => "book".to_string(),
            Model::Tutorial => "tutorial".to_string(),
            Model::BVH => "bvh".to_string(),
            Model::Test => "test".to_string(),
            Model::CornellBox => "cornell_box".to_string(),
            Model::CornellSmoke => "cornell_smoke".to_string(),
            Model::PerlinDebug => "perlin_debug".to_string(),
            Model::Final => "final".to_string(),
            Model::Mandelbrot => "mandelbrot".to_string(),
        }
    }
}

#[derive(Debug, StructOpt)]
#[structopt(name = "tracer", about = "An experimental ray tracer.")]
pub struct Opt {
    /// Prometheus push gateway address, use "" to disable
    #[structopt(
        short = "a",
        long = "pushgateway",
        default_value = "pushgateway.z.xinu.tv:80"
    )]
    pub push_gateway: String,
    /// 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,
    /// Number of threads
    #[structopt(short = "t", long = "num_threads")]
    pub num_threads: Option<usize>,
    /// Sub-samples per pixel
    #[structopt(short = "s", long = "subsample", default_value = "8")]
    pub subsamples: usize,
    /// Select scene to render, one of: "bench", "book", "tutorial", "bvh", "test", "cornell_box",
    /// "cornell_smoke", "perlin_debug", "final"
    #[structopt(long = "model", default_value = "perlin_debug")]
    pub model: Model,
    /// Path to store pprof profile data, i.e. /tmp/cpuprofile.pprof
    #[structopt(long = "pprof", parse(from_os_str))]
    pub pprof: Option<PathBuf>,
    /// 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 fn opt_hash(opt: &Opt) -> String {
    // TODO(wathiede): add threads.
    format!(
        "w:{}-h:{}-s:{}-pprof:{}-model:{}-use_accel:{}-{}",
        opt.width,
        opt.height,
        opt.subsamples,
        opt.pprof.is_some(),
        opt.model.to_string(),
        opt.use_accel,
        opt.output.display().to_string().replace("/", "_")
    )
}

pub struct Scene {
    pub world: Box<dyn Hit>,
    pub camera: Camera,
    pub subsamples: usize,
    pub num_threads: Option<usize>,
    pub width: usize,
    pub height: usize,
    pub global_illumination: bool,
    pub env_map: Option<EnvMap>,
}

// color will trace ray up to 50 bounces deep accumulating color as it goes.  If
// global_illumination is true, a default light background color is assumed and will light the
// world.  If false, it is expected the scene has emissive light sources.
fn color(
    r: Ray,
    world: &dyn Hit,
    depth: usize,
    global_illumination: bool,
    env_map: &Option<EnvMap>,
) -> Vec3 {
    RAY_COUNTER.with_label_values(&[&depth.to_string()]).inc();
    if let Some(rec) = world.hit(r, 0.001, std::f32::MAX) {
        let (u, v) = rec.uv;
        let emitted = rec.material.emitted(u, v, rec.p);
        let scatter_response = rec.material.scatter(&r, &rec);
        if depth < 50 && scatter_response.reflected {
            return emitted
                + scatter_response.attenutation
                    * color(
                        scatter_response.scattered,
                        world,
                        depth + 1,
                        global_illumination,
                        env_map,
                    );
        } else {
            return emitted;
        }
    }
    if global_illumination {
        return match env_map {
            Some(env_map) => env_map.color(r.direction.unit_vector()),
            None => {
                let unit_direction = r.direction.unit_vector();
                // No hit, choose color from background.
                let t = 0.5 * (unit_direction.y + 1.);
                Vec3::new(1., 1., 1.) * (1. - t) + Vec3::new(0.5, 0.7, 1.) * t
            }
        };
    }
    // No global illumination, so background is black.
    Vec3::new(0., 0., 0.)
}

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,
        scene.global_illumination,
        &scene.env_map,
    )
}

enum Request {
    Pixel { x: usize, y: usize },
    Line { width: usize, y: usize },
}

enum Response {
    Pixel {
        x: usize,
        y: usize,
        pixel: Vec3,
    },
    Line {
        width: usize,
        y: usize,
        pixels: Vec<Vec3>,
    },
}
fn render_pixel(scene: &Scene, x: usize, y: usize) -> Vec3 {
    let mut pixel: Vec3 = Default::default();
    for _ in 0..scene.subsamples {
        pixel = pixel + trace_pixel(x, y, scene);
    }
    pixel = pixel / scene.subsamples as f32;
    // Gamma correct, use gamma 2 correction, which is 1/gamma where gamma=2 which is 1/2 or
    // sqrt.
    Vec3::new(pixel[0].sqrt(), pixel[1].sqrt(), pixel[2].sqrt())
}

fn render_worker(
    tid: usize,
    scene: &Scene,
    input_chan: Arc<Mutex<Receiver<Request>>>,
    output_chan: &SyncSender<Response>,
) {
    loop {
        let job = { input_chan.lock().unwrap().recv() };
        match job {
            Err(err) => {
                info!("Shutting down render_worker {}: {}", tid, err);
                return;
            }
            Ok(req) => match req {
                Request::Line { width, y } => {
                    trace!("tid {} width {} y {}", tid, width, y);
                    let pixels = (0..width).map(|x| render_pixel(scene, x, y)).collect();
                    output_chan.send(Response::Line { width, y, pixels });
                }
                Request::Pixel { x, y } => {
                    trace!("tid {} x {} y {}", tid, x, y);
                    let pixel = render_pixel(scene, x, y);
                    output_chan.send(Response::Pixel { x, y, pixel });
                }
            },
        }
    }
    trace!(target: "renderer", "Shutting down worker {}", tid);
}

pub fn render(scene: Scene, output_dir: &Path) -> std::result::Result<(), std::io::Error> {
    let num_threads = scene.num_threads.unwrap_or_else(num_cpus::get);
    let (pixel_req_tx, pixel_req_rx) = sync_channel(2 * num_threads);
    let (pixel_resp_tx, pixel_resp_rx) = sync_channel(2 * num_threads);

    let scene = Arc::new(scene);
    let pixel_req_rx = Arc::new(Mutex::new(pixel_req_rx));
    println!("Creating {} render threads", num_threads);
    for i in 0..num_threads {
        let s = sync::Arc::clone(&scene);
        let pixel_req_rx = pixel_req_rx.clone();
        let pixel_resp_tx = pixel_resp_tx.clone();
        thread::spawn(move || {
            render_worker(i, &s, pixel_req_rx, &pixel_resp_tx);
        });
    }
    drop(pixel_req_rx);
    drop(pixel_resp_tx);

    let (w, h) = (scene.width, scene.height);
    thread::spawn(move || {
        let batch_line_requests = true;
        if batch_line_requests {
            for y in 0..h {
                pixel_req_tx.send(Request::Line { width: w, y });
            }
        } else {
            for y in 0..h {
                for x in 0..w {
                    pixel_req_tx.send(Request::Pixel { x, y });
                }
            }
        }
        drop(pixel_req_tx);
    });

    println!("Rendering with {} subsamples", scene.subsamples);
    let mut img = RgbImage::new(scene.width as u32, scene.height as u32);
    let total = scene.width * scene.height;
    let mut cur_pixel = 0;
    let mut last_progress = 1000;
    for resp in pixel_resp_rx {
        match resp {
            Response::Pixel { x, y, pixel } => {
                let y_inv = scene.height - y - 1;
                img.put_pixel(
                    x as u32,
                    y_inv as u32,
                    image::Rgb([
                        (pixel[0] * 255.).min(255.) as u8,
                        (pixel[1] * 255.).min(255.) as u8,
                        (pixel[2] * 255.).min(255.) as u8,
                    ]),
                );
                let progress = 100 * cur_pixel / total;
                if progress != last_progress {
                    last_progress = progress;
                    if progress % 10 == 0 {
                        print!("{}%", progress);
                    } else {
                        print!(".");
                    }
                    io::stdout().flush().unwrap();
                }
                cur_pixel += 1;
            }
            Response::Line {
                width: _,
                y,
                pixels,
            } => {
                for (x, pixel) in pixels.iter().enumerate() {
                    let y_inv = scene.height - y - 1;
                    img.put_pixel(
                        x as u32,
                        y_inv as u32,
                        image::Rgb([
                            (pixel[0] * 255.).min(255.) as u8,
                            (pixel[1] * 255.).min(255.) as u8,
                            (pixel[2] * 255.).min(255.) as u8,
                        ]),
                    );
                    let progress = 100 * cur_pixel / total;
                    if progress != last_progress {
                        last_progress = progress;
                        if progress % 10 == 0 {
                            print!("{}%", progress);
                        } else {
                            print!(".");
                        }
                        io::stdout().flush().unwrap();
                    }
                    cur_pixel += 1;
                }
            }
        }
    }
    println!();
    io::stdout().flush().unwrap();
    let path = output_dir.join("final.png");
    // Write the contents of this image to the Writer in PNG format.
    trace!(target: "renderer", "Saving {}", path.display());
    img.save(path)
}