raytracers/rtiow/renderer/src/output.rs

use std::collections::HashMap;
use std::fs::File;
use std::io::BufWriter;
use std::path::Path;
use std::sync::Arc;
use std::sync::Mutex;
use std::time;

use chrono::Local;
use image;
use lazy_static::lazy_static;
use log::info;
use serde_derive::Serialize;

use crate::renderer::Scene;
use crate::vec3::Vec3;

// Main RGB image output from rendering the scene.
pub const MAIN_IMAGE: &str = "@final";
// Debug image for adaptive pixels subsampling.
// Red indicates recursion hit maximum depth splitting the pixel.
// Green indicates no subdivision necessary
pub const ADAPTIVE_DEPTH: &str = "adaptive_depth";
// Grey scale showing rays cast per pixel.
pub const RAYS_PER_PIXEL: &str = "rays_per_pixel";

lazy_static! {
    static ref DEBUGGER: Arc<Mutex<Debugger>> = Arc::new(Mutex::new(Debugger::new()));
}

#[derive(Serialize)]
struct ImageMetadata {
    name: String,
    image: String,
    binary: String,
    ratio: f32,
    size: (usize, usize),
    format: ImageType,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct Data<'s> {
    timestamp: i64,
    render_time_seconds: f32,
    scene: &'s Scene,
    image_metadata: Vec<ImageMetadata>,
}

#[derive(Clone, Copy, Serialize)]
pub enum ImageType {
    RGB01,
    Grey01,
    GreyNormalized,
}

struct Image {
    w: usize,
    h: usize,
    pix: Vec<Vec3>,
}

impl Image {
    fn new(w: usize, h: usize) -> Image {
        Image {
            w,
            h,
            pix: (0..w * h).map(|_| [0., 0., 0.].into()).collect(),
        }
    }

    fn put_pixel(&mut self, x: usize, y: usize, p: Vec3) {
        let offset = x + y * self.w;
        self.pix[offset] = p;
    }
}

struct Debugger {
    images: HashMap<String, (ImageType, Image)>,
}

impl Debugger {
    fn new() -> Debugger {
        Debugger {
            images: HashMap::new(),
        }
    }
}

pub fn register_image(name: String, dimensions: (usize, usize), it: ImageType) {
    let mut debugger = DEBUGGER.lock().unwrap();
    debugger
        .images
        .insert(name, (it, Image::new(dimensions.0, dimensions.1)));
}

pub fn set_pixel(name: &str, x: usize, y: usize, pixel: Vec3) {
    let mut debugger = DEBUGGER.lock().unwrap();
    let (_it, img) = debugger
        .images
        .get_mut(name)
        .unwrap_or_else(|| panic!("couldn't find image named '{}'", name));
    let y_inv = img.h - y - 1;
    img.put_pixel(x, y_inv, pixel);
}

pub fn set_pixel_grey(name: &str, x: usize, y: usize, grey: f32) {
    let mut debugger = DEBUGGER.lock().unwrap();
    let (_it, img) = debugger
        .images
        .get_mut(name)
        .unwrap_or_else(|| panic!("couldn't find image named '{}'", name));
    let y_inv = img.h - y - 1;
    img.put_pixel(x, y_inv, [grey, grey, grey].into());
}

trait ImageSaver {
    fn save<Q>(&self, path: Q) -> std::io::Result<()>
    where
        Q: AsRef<Path> + Sized;
}

pub fn write_images<P: AsRef<Path>>(
    scene: &Scene,
    render_time: time::Duration,
    output_dir: P,
) -> std::io::Result<()> {
    let output_dir: &Path = output_dir.as_ref();
    let debugger = DEBUGGER.lock().unwrap();
    let now = Local::now();
    // Write out images in consistent order.
    let mut names = debugger.images.keys().collect::<Vec<_>>();
    names.sort();
    let mut image_metadata = Vec::new();
    for name in &names {
        let (it, img) = debugger.images.get(*name).unwrap();
        let image = format!("{}.png", name);
        let binary = format!("{}.json", name);
        let ratio = img.w as f32 / img.h as f32;
        let size = (img.w, img.h);
        let image_path = output_dir.join(&image);
        let binary_path = output_dir.join(&binary);
        image_metadata.push(ImageMetadata {
            name: name.to_string(),
            image,
            binary,
            ratio,
            size,
            format: *it,
        });
        info!("Saving {}", image_path.to_string_lossy());
        match it {
            ImageType::RGB01 => {
                let mut out_img = image::RgbImage::new(img.w as u32, img.h as u32);
                out_img
                    .enumerate_pixels_mut()
                    .enumerate()
                    .for_each(|(i, (_x, _y, p))| {
                        let pixel = img.pix[i];
                        *p = image::Rgb([
                            (pixel[0] * 255.).min(255.) as u8,
                            (pixel[1] * 255.).min(255.) as u8,
                            (pixel[2] * 255.).min(255.) as u8,
                        ])
                    });
                out_img.save(image_path)?;
            }
            ImageType::Grey01 => {
                let mut out_img = image::GrayImage::new(img.w as u32, img.h as u32);
                out_img
                    .enumerate_pixels_mut()
                    .enumerate()
                    .for_each(|(i, (_x, _y, p))| {
                        let pixel = img.pix[i];
                        *p = image::Luma([(pixel[0] * 255.).min(255.) as u8])
                    });
                out_img.save(image_path)?;
            }
            ImageType::GreyNormalized => {
                let mut out_img = image::GrayImage::new(img.w as u32, img.h as u32);

                let max_val = img.pix.iter().map(|v| v.x).fold(0., f32::max);
                out_img
                    .enumerate_pixels_mut()
                    .enumerate()
                    .for_each(|(i, (_x, _y, p))| {
                        let pixel = img.pix[i];
                        *p = image::Luma([(pixel[0] / max_val * 255.).min(255.) as u8])
                    });
                out_img.save(image_path)?;
            }
        };
        info!("Saving {}", binary_path.to_string_lossy());
        let f = File::create(output_dir.join(binary_path))?;
        let f = BufWriter::new(f);
        match it {
            ImageType::RGB01 => {
                serde_json::ser::to_writer(
                    f,
                    &img.pix
                        .iter()
                        .map(|v| [v.x, v.y, v.z])
                        .collect::<Vec<[f32; 3]>>(),
                )?;
            }
            ImageType::Grey01 | ImageType::GreyNormalized => {
                serde_json::ser::to_writer(f, &img.pix.iter().map(|v| v.x).collect::<Vec<f32>>())?;
            }
        };
    }
    let f = File::create(output_dir.join("data.json"))?;
    let f = BufWriter::new(f);
    serde_json::ser::to_writer(
        f,
        &Data {
            timestamp: now.timestamp(),
            render_time_seconds: render_time.as_secs_f32(),
            scene,
            image_metadata,
        },
    )?;
    Ok(())
}