Implement camera colors

Still needs some unit tests, and to fix the gasket example

enkou-shaders/examples/gasket.rs

@@ -1,12 +1,12 @@
 use anyhow::{Context, Result};
 use enkou_shaders::Coefficients2;
 use enkou_shaders::camera::Camera;
-use enkou_shaders::camera::entry::main_camera;
+use enkou_shaders::camera::entry::main_image_render;
 use enkou_shaders::chaos_game::entry::main_chaos_game;
 use enkou_shaders::transform::Transform;
 use enkou_shaders::variation::Variation;
-use glam::{Affine2, IVec2, UVec2, Vec2, Vec4, uvec2, vec2};
-use image::{GrayImage, Luma};
+use glam::{Affine2, UVec2, Vec2, Vec4, uvec2, vec2};
+use image::{Rgba, Rgba32FImage};
 use std::mem;
 use std::process::Command;
 use tempfile::NamedTempFile;
@@ -62,20 +62,26 @@ pub fn main() -> Result<()> {
         IMAGE_DIMENSION.as_vec2(),
     );
 
+    let palette = &[Vec4::ONE; 2];
+
     let mut output_points_pixel = Vec::new();
-    output_points_pixel.resize(ITERATIONS as usize, IVec2::ZERO);
+    output_points_pixel.resize(ITERATIONS as usize, Vec4::ZERO);
 
-    main_camera(&camera, &output_points_ifs, &mut output_points_pixel);
+    main_image_render(
+        &camera,
+        palette,
+        &output_points_ifs,
+        &mut output_points_pixel,
+    );
 
-    let mut image = GrayImage::new(IMAGE_DIMENSION.x, IMAGE_DIMENSION.y);
-    let dimensions = image.dimensions();
-    output_points_pixel
-        .iter()
-        .skip_while(|p| {
-            p.x < 0 || (p.x as u32) > dimensions.0 || p.y < 0 || (p.y as u32) > dimensions.1
-        })
-        .map(|p| (p.x as u32, p.y as u32))
-        .for_each(|(x, y)| image.put_pixel(x, y, Luma([255u8])));
+    let mut image = Rgba32FImage::new(IMAGE_DIMENSION.x, IMAGE_DIMENSION.y);
+    for x in 0..image.dimensions().0 {
+        for y in 0..image.dimensions().1 {
+            let pixel_index = y * IMAGE_DIMENSION.x + x;
+            let pixel = output_points_pixel[pixel_index as usize];
+            image.put_pixel(x, y, Rgba(pixel.into()));
+        }
+    }
 
     let temp = NamedTempFile::with_suffix(".png").context("Unable to create file for image")?;
     image.save(temp.path()).context("Unable to save image")?;

enkou-shaders/src/camera.rs

@@ -2,8 +2,47 @@
 //!
 //! Map points from the IFS coordinate system to pixel coordinates. This is a lossy transformation.
 use bytemuck::{Pod, Zeroable};
-use glam::{Affine2, IVec2, UVec2, Vec2, vec2};
-use libm::powf;
+use glam::{Affine2, IVec2, UVec2, Vec2, Vec4, Vec4Swizzles, vec2};
+use libm::{ceilf, floorf, powf};
+
+/// Blending modes for mapping IFS color values (which are on a scale `[0, 1]`)
+/// to RGBA colors.
+#[derive(Copy, Clone)]
+#[repr(u32)]
+pub enum BlendMode {
+    /// Map IFS color values to a linear blend of the nearest two palette colors
+    Linear = 0,
+
+    /// Map IFS color values to the nearest single palette color
+    Step = 1,
+}
+
+impl Default for BlendMode {
+    fn default() -> Self {
+        BlendMode::Linear
+    }
+}
+
+impl BlendMode {
+    /// Map an IFS color value to RGBA color from the provided palette.
+    pub fn ifs_to_rgb(&self, color: f32, palette: &[Vec4]) -> Vec4 {
+        let palette_index = color * palette.len() as f32;
+        let palette_index_lower = floorf(palette_index) as usize;
+        let palette_index_upper = ceilf(palette_index) as usize;
+
+        match self {
+            BlendMode::Linear => {
+                (palette[palette_index_lower] + palette[palette_index_upper]) / 2.0
+            }
+            BlendMode::Step => palette[palette_index_lower],
+        }
+    }
+}
+
+// UNSAFE: Sound because enum has guaranteed layout (u32) and defined zero-value
+unsafe impl bytemuck::Zeroable for BlendMode {}
+// UNSAFE: Sound because enum has guaranteed layout (u32) and defined zero-value
+unsafe impl bytemuck::Pod for BlendMode {}
 
 /// Settings used to map IFS coordinates to pixel coordinates.
 ///
@@ -14,6 +53,8 @@ use libm::powf;
 pub struct Camera {
     dimensions: UVec2,
     transform: Affine2,
+    blend_mode: BlendMode,
+    image_gamma: f32,
 }
 
 impl Camera {
@@ -50,62 +91,70 @@ impl Camera {
         Camera {
             dimensions,
             transform,
+            blend_mode: BlendMode::Linear,
+            image_gamma: 1.5,
         }
     }
 
-    /// Map a point from IFS coordinates to pixel coordinates.
-    ///
-    /// ```
-    /// # use glam::{vec2, ivec2, uvec2, Vec2};
-    /// # use crate::enkou_shaders::camera::Camera;
-    /// // Output image is 600x600 pixels, centered at the origin, no rotation, no zoom,
-    /// // and scaled such that it covers the range [-2, 2].
-    /// // Use the origin as the IFS coordinate, so the pixel coordinate is the center of the image
-    /// let camera = Camera::new(
-    ///     uvec2(600, 600),
-    ///     Vec2::ZERO,
-    ///     0.0,
-    ///     Vec2::ZERO,
-    ///     vec2(150.0, 150.0)
-    /// );
-    /// assert_eq!(camera.transform_point(vec2(0.0, 0.0)), ivec2(300, 300));
-    /// ```
-    pub fn transform_point(&self, point: Vec2) -> IVec2 {
+    fn transform_point(&self, point: Vec2) -> IVec2 {
         self.transform.transform_point2(point).as_ivec2()
     }
 
-    /// Map a point from IFS coordinates to pixel coordinates (like [`transform_point`](Camera::transform_point)),
-    /// and check that the result is within the provided image dimensions.
-    pub fn transform_point_to_image(&self, point: Vec2) -> Option<UVec2> {
-        let pixel_coordinates = self.transform_point(point);
+    /// Map a point from IFS coordinates to a pixel index and RGBA value; if the IFS coordinate
+    /// is outside the viewable range, return [`None`].
+    pub fn transform_point_to_image(&self, point: Vec4, palette: &[Vec4]) -> Option<(usize, Vec4)> {
+        let pixel_coordinates = self.transform_point(point.xy());
         if pixel_coordinates.x < 0
             || pixel_coordinates.y < 0
             || (pixel_coordinates.x as u32) >= self.dimensions.x
             || (pixel_coordinates.y as u32) >= self.dimensions.y
         {
-            None
-        } else {
-            Some(pixel_coordinates.as_uvec2())
+            return None;
         }
+
+        let to_pixel_index = self.dimensions.with_y(0);
+        let pixel_index = pixel_coordinates.as_uvec2().dot(to_pixel_index) as usize;
+
+        let rgba = self.blend_mode.ifs_to_rgb(point.w, palette);
+
+        Some((pixel_index, rgba))
     }
 }
 
-/// Shader entry point for running the camera transformation over a list of IFS coordinates
+#[allow(missing_docs)]
 pub mod entry {
     use crate::camera::Camera;
-    use glam::{Vec4, Vec4Swizzles};
-    use spirv_std::glam::IVec2;
+    use glam::Vec4;
+    use libm::log10f;
     use spirv_std::spirv;
 
-    /// Transform IFS coordinates to pixel coordinates
-    #[spirv(compute(entry_point_name = "main_camera", threads(1)))]
-    pub fn main_camera(
+    /// Render an output image from a list of IFS coordinates.
+    ///
+    /// Arguments:
+    /// * `camera` - Camera settings for mapping IFS coordinates to pixel coordinates
+    /// * `palette` - Color palette to use when mapping IFS color to RGB colors. Individual elements
+    ///     are assumed to be RGBA values on the scale of `[0, 1]`
+    /// * `coordinates_ifs` - IFS coordinates to use for the output image
+    /// * `image` - Buffer for the output image
+    #[spirv(compute(entry_point_name = "main_image_render", threads(1)))]
+    pub fn main_image_render(
         #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] camera: &Camera,
+        #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] palette: &[Vec4],
         #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] coordinates_ifs: &[Vec4],
-        #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] coordinates_pixel: &mut [IVec2],
+        #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] image: &mut [Vec4],
     ) {
-        for i in 0..coordinates_ifs.len() {
-            coordinates_pixel[i] = camera.transform_point(coordinates_ifs[i].xy())
+        for coordinate_index in 0..coordinates_ifs.len() {
+            camera
+                .transform_point_to_image(coordinates_ifs[coordinate_index], palette)
+                .map(|(pixel_index, rgba)| image[pixel_index] += rgba);
+        }
+
+        for pixel_index in 0..image.len() {
+            // TODO: Fix the bootleg gamma adjustment
+            let pixel_unscaled = image[pixel_index];
+            let pixel =
+                pixel_unscaled * log10f(pixel_unscaled.w) / (pixel_unscaled.w * camera.image_gamma);
+            image[pixel_index] = pixel;
         }
     }
 }
@@ -117,7 +166,7 @@ mod test {
     use libm::powf;
 
     #[test]
-    pub fn manual_camera() {
+    fn manual_camera() {
         let starting_point = vec2(1.0, 1.0);
 
         // Move the origin; points move right and up by one unit, giving us (2.0, 2.0)
@@ -155,7 +204,7 @@ mod test {
     }
 
     #[test]
-    pub fn point_outside_camera() {
+    fn point_outside_camera() {
         // Scale 250 for an image 1000 x 1000 gives an effective range of [-2, 2]
         let camera = Camera::new(
             uvec2(1000, 1000),
@@ -170,7 +219,7 @@ mod test {
     }
 
     #[test]
-    pub fn point_outside_camera_negative() {
+    fn point_outside_camera_negative() {
         // Scale 250 for an image 1000 x 1000 gives an effective range of [-2, 2]
         let camera = Camera::new(
             uvec2(1000, 1000),
@@ -185,7 +234,7 @@ mod test {
     }
 
     #[test]
-    pub fn aspect_ratio() {
+    fn aspect_ratio() {
         // Scale 100 for an image 1600 x 900 gives an effective X range of [-8, 8],
         // and effective Y range of [-4.5, 4.5]
         let camera = Camera::new(