bs_trace/image/
image.rs

1use super::prelude::*;
2use crate::linalg::vector::Vector;
3use const_default::ConstDefault;
4use std::io::Write;
5use std::ops::{Mul, MulAssign};
6
7/// A colour, with [r], [g] and [b] values.
8#[derive(Debug, Clone, Copy, PartialEq)]
9pub struct Colour(pub Vector<f64, 3>);
10
11/// An image with compile-time-known width and height, written `Image<W, H>`.
12///
13/// The bitmap data is represented as a grid of [Colour]s.
14pub struct Image<const W: usize, const H: usize> {
15    pub rows: Vec<Vec<Colour>>,
16}
17
18impl Colour {
19    pub const WHITE: Colour = Colour::new(1.0, 1.0, 1.0);
20    pub const BLACK: Colour = Colour::new(0.0, 0.0, 0.0);
21
22    pub const fn new(r: f64, g: f64, b: f64) -> Self {
23        debug_assert!(0.0 <= r && r <= 1.0);
24        debug_assert!(0.0 <= g && g <= 1.0);
25        debug_assert!(0.0 <= b && b <= 1.0);
26        Self(Vector::new([r, g, b]))
27    }
28
29    /// Returns the red value for this colour.
30    pub const fn r(&self) -> f64 {
31        *self.0.x()
32    }
33
34    /// Returns the green value for this colour.
35    pub const fn g(&self) -> f64 {
36        *self.0.y()
37    }
38
39    /// Returns the blue value for this colour.
40    pub const fn b(&self) -> f64 {
41        *self.0.z()
42    }
43
44    pub fn map<F>(self, f: F) -> Self
45    where
46        F: Fn(f64) -> f64,
47    {
48        Colour(self.0.map(f))
49    }
50
51    pub fn to_pixel(&self) -> Pixel {
52        // TODO these clamps may be unnecessary if I lock down new further
53        let r = self.r().clamp(0.0, 1.0);
54        let g = self.g().clamp(0.0, 1.0);
55        let b = self.b().clamp(0.0, 1.0);
56
57        let r = (r * 255.0) as u8;
58        let g = (g * 255.0) as u8;
59        let b = (b * 255.0) as u8;
60        Pixel::new(r, g, b)
61    }
62}
63
64impl ConstDefault for Colour {
65    const DEFAULT: Self = Colour::new(0.0, 0.0, 0.0);
66}
67
68impl<const W: usize, const H: usize> Image<W, H> {
69    /// Constructs an image from the given row data.
70    ///
71    /// This method will panic if there are not precisely `H` rows, or any of the rows do not contain precisely `W` colours.
72    pub fn new(rows: Vec<Vec<Colour>>) -> Self {
73        assert_eq!(rows.len(), H);
74        for row in rows.iter() {
75            assert_eq!(row.len(), W);
76        }
77
78        unsafe { Self::new_unchecked(rows) }
79    }
80
81    /// Constructs an image from the given row data.
82    ///
83    /// # Safety
84    ///
85    /// It is the responsibility of the caller to ensure that `rows` contains precisely `H`
86    /// rows, and that each row contains precisely `W` colours.
87    pub unsafe fn new_unchecked(rows: Vec<Vec<Colour>>) -> Self {
88        Self { rows }
89    }
90
91    /// Generates an image using the given function in each pixel's (x,y) coordinates -- the
92    /// top-left corner of the image is (0,0), x grows to the right in the positive direction,
93    /// and y grows downwards in the positive direction.
94    pub fn gen<F>(mut f: F) -> Self
95    where
96        F: FnMut(usize, usize) -> Colour,
97    {
98        let mut rows = Vec::<Vec<Colour>>::with_capacity(H);
99        for y in 0..H {
100            let mut row_data = Vec::<Colour>::with_capacity(W);
101            for x in 0..W {
102                row_data.push(f(x, y));
103            }
104            rows.push(row_data);
105        }
106
107        unsafe { Self::new_unchecked(rows) }
108    }
109
110    /// Converts this image to PPM format.
111    pub fn to_ppm(&self) -> Ppm<W, H> {
112        let mut pixel_rows = Vec::<Vec<Pixel>>::with_capacity(H);
113        for colour_row in self.rows.iter() {
114            let mut pixel_row = Vec::<Pixel>::with_capacity(W);
115            for colour in colour_row.iter() {
116                pixel_row.push(colour.to_pixel());
117            }
118            pixel_rows.push(pixel_row);
119        }
120        Ppm::new(pixel_rows)
121    }
122
123    pub fn write_as_png<OUT: Write>(&self, writer: OUT) {
124        let mut encoder = png::Encoder::new(writer, W as u32, H as u32);
125        encoder.set_color(png::ColorType::Rgb);
126        encoder.set_depth(png::BitDepth::Eight);
127        let mut png_writer = encoder.write_header().unwrap();
128        let image_data: Vec<u8> = self
129            .rows
130            .iter()
131            .flat_map(|row| row.iter().flat_map(|colour| colour.to_pixel().into_iter()))
132            .collect();
133        png_writer.write_image_data(&image_data).unwrap();
134    }
135}
136
137impl From<Vector<f64, 3>> for Colour {
138    fn from(v: Vector<f64, 3>) -> Self {
139        Self::new(*v.x(), *v.y(), *v.z()) // TODO is this optimised to be equivalent to Self(v)?
140    }
141}
142
143impl Mul<Colour> for Colour {
144    type Output = Colour;
145
146    fn mul(self, rhs: Colour) -> Self::Output {
147        Colour(self.0 * rhs.0)
148    }
149}
150
151impl MulAssign<Colour> for Colour {
152    fn mul_assign(&mut self, rhs: Colour) {
153        self.0 *= rhs.0;
154    }
155}
bs_trace/image/image.rs

bs_trace/image/
image.rs
