7
\$\begingroup\$

I have implemented a simple version of the median cut algorithm. It takes a vector of Color structs representing pixel in an image. I also use the ColorChanel enum representing RGBA channels.

pub struct Color {
    pub r: u8,
    pub g: u8,
    pub b: u8,
    pub a: u8,
}

#[derive(Debug)]
pub enum ColorChannel {
    R,
    G,
    B,
    A,
}

impl Color {
    pub fn channel_val(&self, channel: &ColorChannel) -> u8 {
        match channel {
            ColorChannel::R => self.r,
            ColorChannel::G => self.g,
            ColorChannel::B => self.b,
            ColorChannel::A => self.a,
        }
    }
}

impl PartialEq for Color {
    fn eq(&self, other: &Self) -> bool {
        self.r == other.r && self.g == other.g && self.b == other.b && self.a == other.a
    }
}

For each Color/Pixel vector I look for the channel with the highest range:

/// Returns the color channel with the highest range.
/// IMPORTANT: Ignores alpha channel!
///
/// # Arguments
///
/// * `colors` - Color vector from which the highest range is evaluated.
///
fn highest_range_channel(colors: &Vec<Color>) -> Option<ColorChannel> {
    if let Some(ranges) = color_ranges(colors) {
        let mut highest_range_channel = ColorChannel::R;
        let mut highest_value = ranges.r;

        if ranges.g > highest_value {
            highest_range_channel = ColorChannel::G;
            highest_value = ranges.g;
        }

        if ranges.b > highest_value {
            highest_range_channel = ColorChannel::B;
        }

        return Some(highest_range_channel);
    }

    None
}

The color ranges are calculated in the according function:

/// Returns the ranges for each color channel
///
/// # Arguments
///
/// * `colors` - Color vector from which the ranges are calculated.
///
/// # Examples
///
/// ```
/// let colors = Vec::<Color>::new();
/// let color_ranges_data = color_ranges(colors);
/// ```
///
fn color_ranges(colors: &Vec<Color>) -> Option<Color> {
    if colors.is_empty() {
        return None;
    }

    // Unwrap is ok here, because `max_by_key` only returns `None` for empty vectors
    let r_range = colors.iter().max_by_key(|c| c.r).unwrap().r;
    let g_range = colors.iter().max_by_key(|c| c.g).unwrap().g;
    let b_range = colors.iter().max_by_key(|c| c.b).unwrap().b;
    let a_range = colors.iter().max_by_key(|c| c.a).unwrap().a;

    Some(Color {
        r: r_range,
        g: g_range,
        b: b_range,
        a: a_range,
    })
}

After that I calculate the median value for the channel with the highest range. I am doing this by sorting the vector based on the desired channel and finding the value in the "middle" of the vector.

/// Sort a color vector for a specific channel.
///
/// # Arguments
///
/// * `colors` - Color data which will be sorted.
/// * `channel` - Target channel. The sorting is performed based on this value.
///
/// # Examples
///
/// ```
/// let mut colors = Vec::<Color>::new();
/// sort_colors(&mut colors, &ColorChannel::R);
/// ```
///
fn sort_colors(colors: &mut Vec<Color>, channel: &ColorChannel) {
    if colors.is_empty() {
        return;
    }

    match channel {
        ColorChannel::R => colors.sort_by(|a, b| a.r.cmp(&b.r)),
        ColorChannel::G => colors.sort_by(|a, b| a.g.cmp(&b.g)),
        ColorChannel::B => colors.sort_by(|a, b| a.b.cmp(&b.b)),
        ColorChannel::A => colors.sort_by(|a, b| a.a.cmp(&b.a)),
    }
}

/// Returns median value for a specific `ColorChannel`.
///
/// # Arguments
///
/// * `colors` - Color vector from which the median value is calculated.
/// * `channel` - Target channel for which the median is calculated.
///
/// # Examples
/// ```
/// let mut colors = Vec::<Color>::new();
/// let mut result = color_median(&mut colors, &ColorChannel::R);
/// ```
///
fn color_median(colors: &mut Vec<Color>, channel: &ColorChannel) -> Option<u8> {
    if colors.is_empty() {
        return None;
    }

    sort_colors(colors, channel);

    let mid = colors.len() / 2;
    if colors.len() % 2 == 0 {
        channel_mean(&vec![colors[mid - 1], colors[mid]], channel)
    } else {
        channel_value_by_index(colors, mid, channel)
    }
}

/// Returns a color value based on the provided channel and index parameters.
///
/// # Arguments
///
/// * `colors` - Color vector from which the value is retreived.
/// * `index` - Index of the target color in the vector.
/// * `channel` - Color channel of the searched value.
///
/// # Examples
///
/// ```
/// let mut colors: Vec<Color> = Vec::new();
/// colors.push(Color { r: 100, g: 22, b: 12, a: 0 });
/// assert_eq!(Some(100), channel_value_by_index(&colors, 0, &ColorChannel::R));
/// ```
///
fn channel_value_by_index(colors: &Vec<Color>, index: usize, channel: &ColorChannel) -> Option<u8> {
    if colors.is_empty() || index >= colors.len() {
        return None;
    }
    match channel {
        ColorChannel::R => Some(colors[index].r),
        ColorChannel::G => Some(colors[index].g),
        ColorChannel::B => Some(colors[index].b),
        ColorChannel::A => Some(colors[index].a),
    }
}

/// Calculate the mean value for a specific color channel on a vector of `Color`.
///
/// # Arguments
///
/// * `colors` - Color vector from which the mean value is calculated.
/// * `channel` - Target channel for which the mean is calculated.
///
/// # Examples
///
/// ```
/// let mut colors: Vec<Color> = Vec::new();
/// let mut result = channel_mean(&colors, &ColorChannel::R);
/// ```
///
fn channel_mean(colors: &Vec<Color>, channel: &ColorChannel) -> Option<u8> {
    let number_colors = colors.len();

    if number_colors == 0 {
        return None;
    }

    match channel {
        ColorChannel::R => Some((colors.iter().fold(0, |acc: u32, x| x.r as u32 + acc) / number_colors as u32) as u8),
        ColorChannel::G => Some((colors.iter().fold(0, |acc: u32, x| x.g as u32 + acc) / number_colors as u32) as u8),
        ColorChannel::B => Some((colors.iter().fold(0, |acc: u32, x| x.b as u32 + acc) / number_colors as u32) as u8),
        ColorChannel::A => Some((colors.iter().fold(0, |acc: u32, x| x.a as u32 + acc) / number_colors as u32) as u8),
    }
}

Now I create two new vectors/buckets with one containing all Colors above the median value and another with Color values below the median. This entire process is implemented in the median_cut function.

/// Performs the median cut on a single vector (bucket) of `Color`.
/// Returns two `color` vectors representing the colors above and colors below median value.
///
/// # Arguments
///
/// * `colors` - `Color` vector on which the median cut is performed.
///
fn median_cut(colors: &mut Vec<Color>) -> (Vec<Color>, Vec<Color>) {
    if colors.is_empty() {
        return (Vec::<Color>::new(), Vec::<Color>::new());
    }

    if let Some(highest_range_channel) = highest_range_channel(&colors) {
        if let Some(median) = color_median(colors, &highest_range_channel) {
            let mut above_median = Vec::<Color>::new();
            let mut below_median = Vec::<Color>::new();
            for color in colors {
                if color.channel_val(&highest_range_channel) > median {
                    above_median.push(*color);
                } else {
                    below_median.push(*color);
                }
            }

            return (above_median, below_median);
        }
    }

    return (Vec::<Color>::new(), Vec::<Color>::new());
}

In order to perform multiple iterations I call the recurse function. It takes a bucket, performs the median_cut on it calculates the mean color for each output bucket and performs the median_cut on the them. The function stops and returns all mean colors when the amount of iterations reaches 0.

pub fn recurse(bucket: &mut Vec<Color>, iter_count: u8) -> Option<Vec<Color>> {
    if iter_count < 1 || bucket.is_empty() {
        return None;
    }
    let mut result = Vec::<Color>::new();

    let mut new_buckets = median_cut(bucket);
    if !new_buckets.0.is_empty() {
        if let Some(c_0) = color_mean(&new_buckets.0) {
            result.push(c_0);
        }

        if let Some(new_colors) = recurse(&mut new_buckets.0, iter_count - 1) {
            result.append(&mut new_colors.clone());
        }
    }

    if !new_buckets.1.is_empty() {
        if let Some(c_1) = color_mean(&new_buckets.1) {
            result.push(c_1);
        }

        if let Some(new_colors) = recurse(&mut new_buckets.1, iter_count - 1) {
            result.append(&mut new_colors.clone());
        }
    }

    Some(result)
}

/// Returns the mean color value based on the passed colors.
///
/// # Arguments
///
/// * `colors` - Color vector from which the mean color is calculated.
///
/// # Examples
///
/// ```
/// let colors = Vec::<Color>::new();
/// let result = color_mean(&colors);
/// ```
///
fn color_mean(colors: &Vec<Color>) -> Option<Color> {
    if colors.is_empty() {
        return None;
    }

    let r_mean = (colors.iter().fold(0, |acc: u32, c| acc + c.r as u32) / colors.len() as u32) as u8;
    let g_mean = (colors.iter().fold(0, |acc: u32, c| acc + c.g as u32) / colors.len() as u32) as u8;
    let b_mean = (colors.iter().fold(0, |acc: u32, c| acc + c.b as u32) / colors.len() as u32) as u8;
    let a_mean = (colors.iter().fold(0, |acc: u32, c| acc + c.a as u32) / colors.len() as u32) as u8;

    Some(Color {
        r: r_mean,
        g: g_mean,
        b: b_mean,
        a: a_mean,
    })
}

So basically a call to the entire algorithm looks like this:

let mut pixels = Vec::new();
// fill pixels with data

if let Some(colors) = recurse(&mut pixels, 3) {
    // Do something with the output colors
}

I am very new to rust so my concerns about the code are:

  • The usage of Option. I second guess if I should have used Result with some meaningful error information. Also I am feeling like the "high level" usage of functions returning Options produces arrow-like code (like in median_cut) which I personally find hard to read.
  • Overall performance, it feels like there is a lot of iteration and cloning going on.
  • Algorithm implementation, There are scenarios which return "interesting" results. E.g. if I use an image which only contains red and blue pixels, the algorithm returns a blue, a red and a purple result. As I understand it the algorithm should perform a color reduction.

How could I optimize the code based on the above topics?

\$\endgroup\$

2 Answers 2

4
+200
\$\begingroup\$

You probably shouldn't award this a bounty; I'm just using your question as an excuse to learn Rust. Specifically, my understanding of how the [de]referencing works is clunky at best.

Does it run?

In order to get your code to compile, I had to add [derive(Debug,Clone,Copy)] to pub struct Color. But even then, we need to know if it works, which means we need something to run it on.

Adding this to what you wrote works:

use std::env;
use image::{Rgba, RgbaImage};
use image::error::ImageResult as ImageResult;
use image::io::Reader as ImageReader;


/// Convert from an image::RgbaImage to the local representation. 
fn read_pixels(image: RgbaImage) -> Vec<Color> {
    let to_color = |p: Rgba<u8>| Color {r: p[0], g: p[1], b: p[2], a: p[3], };
    let mut pixels = Vec::new();
    pixels.extend(image.pixels().map(|p| to_color(*p)));
    pixels
}

/// Read an image from disk. Returns the RgbaImage and its width and height.
fn read_image(filename: String) -> ImageResult<(RgbaImage, u32, u32)> {
    let image = ImageReader::open(&filename)?.decode()?.to_rgba8();
    let (w, h) = (image.width(), image.height());
    Ok((image, w, h))
}

/// Make a copy of colors, in which every color is replaced by the closest match in palette.
/// Uses Pythagorean distance.
fn assign_colors(colors: Vec<Color>, palette: Vec<Color>) -> Vec<Color> {
    // Skip the sqrt, we don't need it. Use u32 to prevent overflows.
    let channel_distance = |a: u8, b:u8| ((if a < b {b - a} else {a - b}) as u32).pow(2);
    let mut result = Vec::new();
    result.extend(colors.iter().map(
            |c| palette.iter().min_by_key(
                |p| channel_distance(c.r, p.r)
                  + channel_distance(c.g, p.g)
                  + channel_distance(c.b, p.b)
                  + channel_distance(c.a, p.a)
                ).unwrap()
            ));
    result
}

/// Convert from the local representation to an image::RgbaImage.
/// Could theoretically panic in a variety of ways.
fn gather_pixels(colors: Vec<Color>, width: u32, height: u32) -> RgbaImage {
    let from_color = |c: Color| Rgba::<u8>::from([c.r, c.g, c.b, c.a]);
    RgbaImage::from_fn(width, height,
                       |x, y| from_color(colors[usize::try_from(x + (width * y)).unwrap()])
                       )
}

/// The inner "main" function; wraps failure conditions.
fn handle_file(input_file: String, output_file: String) -> Result<(), String> {
    let e2str = |e| format!("{}", e);
    let (image, width, height) = read_image(input_file).map_err(e2str)?;
    let pixels = read_pixels(image);
    let palette = recurse(&mut pixels.clone(), 3)
        .ok_or(String::from("There was a problem building the palette."))?;
    let out_image = gather_pixels(assign_colors(pixels, palette), width, height);
    out_image.save(output_file).map_err(e2str)
}

fn main() {
    let input_file = env::args().nth(1).expect("Missing argument");
    let output_file = env::args().nth(2).expect("Missing argument");
    let result = handle_file(input_file, output_file);
    if let Err(error) = result {
        println!("{}", error);
    } else {
        println!("Success!");
    }
}

It takes about 11 seconds to run (on a 1.3MB 1080x1920 png).

Docstrings

You seem to be following a verbose format for your docstrings. That's not bad, but it's not what I'd do. You've done a good job of naming things in general, so adding additional comments for every argument doesn't seem to have much value.

On the other hand, why doesn't recurse have a docsting? It actually needs some explanation. While you're at it, rename it to make_palette.

Are you actually following the algorithm?

From your Wikipedia link:

Since the number of buckets doubles with each iteration, this algorithm can only generate a palette with a number of colors that is a power of two. To generate, say, a 12-color palette, one might first generate a 16-color palette and merge some of the colors in some way.

And yet during debugging of my code, I definitely noticed you were generating pallets of 14 colors.

I'm pretty sure the issue is in recurse. After calling median_cut, for each returned bucket, you first push the mean of the bucket into result, and then recurse. For a given approximate palette-size, this will waste slots on lower-saturation colors, compared to only taking the means of "leaf" buckets. This actually makes the function much simpler:

/// Perform the median cut algorithm.
/// Returns a palette with 2^iter_count colors.
/// https://en.wikipedia.org/wiki/Median_cut
pub fn make_palette(bucket: &mut Vec<Color>, iter_count: u8) -> Option<Vec<Color>> {
    if iter_count < 1 {
        return Some(Vec::from([color_mean(&bucket)?]));
    }

    let mut new_buckets = median_cut(bucket);
    let mut result = make_palette(&mut new_buckets.0, iter_count - 1)?;
    result.append(&mut make_palette(&mut new_buckets.1, iter_count - 1)?);
    Some(result)
}

Your calculation of the ranges seems to be just checking for the max value; I assume that's a mistake.

Also, why are you ignoring the alpha channel?

Can we make stuff more elegant?

If we replace Color::channel_val with ColorChannel::value, then a lot of the places where we're repeating stuff per-channel can be reduced. channel_value_by_index no longer needs to be a defined function, for example. You can also add a higher-order Color builder function.

This brings the code down to something that feels more intuitive to me:

use std::env;
use image::{Rgba, RgbaImage};
use image::error::ImageResult as ImageResult;
use image::io::Reader as ImageReader;


#[derive(Debug,Clone,Copy)]
pub struct Color {
    pub r: u8,
    pub g: u8,
    pub b: u8,
    pub a: u8,
}

#[derive(Debug,Clone,Copy)]
pub enum ColorChannel {
    R,
    G,
    B,
    A,
}

impl ColorChannel {
    pub const ALL: [ColorChannel; 4] = [Self::R, Self::G, Self::B, Self::A];

    pub fn value(&self, color: &Color) -> u8 {
        match self {
            ColorChannel::R => color.r,
            ColorChannel::G => color.g,
            ColorChannel::B => color.b,
            ColorChannel::A => color.a,
        }
    }
}

impl Color {
    pub fn from_fn<F>(f: F) -> Option<Color>
        where F: Fn(ColorChannel) -> Option<u8> {
        Some(Color {r: f(ColorChannel::R)?,
                    g: f(ColorChannel::G)?,
                    b: f(ColorChannel::B)?,
                    a: f(ColorChannel::A)?,
        })
    }
}

impl PartialEq for Color {
    fn eq(&self, other: &Self) -> bool {
        ColorChannel::ALL.iter()
            .all(|channel| channel.value(self) == channel.value(other))
    }
}

/// Returns the color channel with the highest range.
/// IMPORTANT: Ignores alpha channel!
fn highest_range_channel(colors: &Vec<Color>) -> Option<ColorChannel> {
    let ranges = color_ranges(colors)?;
    let channel: &ColorChannel = ColorChannel::ALL.iter()
        .max_by_key(|channel| channel.value(&ranges))?;
    Some(*channel)
}

/// Returns the ranges for each color channel
fn color_ranges(colors: &Vec<Color>) -> Option<Color> {
    Color::from_fn(|channel|
                   Some(colors.iter().map(|color| channel.value(color)).max()?
                        - colors.iter().map(|color| channel.value(color)).min()?))
}

/// Returns median value for a specific `ColorChannel` across `colors`.
fn channel_median(colors: &mut Vec<Color>, channel: &ColorChannel) -> Option<u8> {
    colors.sort_by_key(|a| channel.value(a));

    let mid = colors.len() / 2;
    if colors.len() % 2 == 0 {
        channel_mean(&vec![colors[mid - 1], colors[mid]], channel)
    } else {
        Some(channel.value(colors.get(mid)?))
    }
}

/// Calculate the mean value for a specific color channel on a vector of `Color`.
fn channel_mean(colors: &Vec<Color>, channel: &ColorChannel) -> Option<u8> {
    let number_colors = colors.len() as u32;

    if number_colors == 0 {
        return None;
    }

    let mean = colors.iter()
        .fold(0, |acc: u32, x| channel.value(x) as u32 + acc) / number_colors;
    Some(mean as u8)
}

/// Performs the median cut on a single vector (bucket) of `Color`.
/// Returns two vectors (buckets) with the colors above and below the chosen median.
fn median_cut(colors: &mut Vec<Color>) -> Option<(Vec<Color>, Vec<Color>)> {
    let mut above_median = Vec::<Color>::new();
    let mut below_median = Vec::<Color>::new();
    let channel = highest_range_channel(&colors)?;
    let median = channel_median(colors, &channel)?;

    for color in colors {
        if channel.value(color) > median {
            above_median.push(*color);
        } else {
            below_median.push(*color);
        }
    }

    return Some((above_median, below_median));
}

/// Perform the median cut algorithm.
/// Returns a palette with 2^iter_count colors.
/// https://en.wikipedia.org/wiki/Median_cut
pub fn make_palette(bucket: &mut Vec<Color>, iter_count: u8) -> Option<Vec<Color>> {
    if iter_count < 1 {
        return Some(vec![Color::from_fn(|channel| channel_mean(bucket, &channel))?]);
    }

    let mut new_buckets = median_cut(bucket)?;
    let mut result = make_palette(&mut new_buckets.0, iter_count - 1)?;
    result.append(&mut make_palette(&mut new_buckets.1, iter_count - 1)?);
    Some(result)
}

/// Convert from an image::RgbaImage to the local representation. 
fn read_pixels(image: RgbaImage) -> Vec<Color> {
    let to_color = |p: Rgba<u8>| Color {r: p[0], g: p[1], b: p[2], a: p[3], };
    let mut pixels = Vec::new();
    pixels.extend(image.pixels().map(|p| to_color(*p)));
    pixels
}

/// Read an image from disk. Returns the RgbaImage and its width and height.
fn read_image(filename: String) -> ImageResult<(RgbaImage, u32, u32)> {
    let image = ImageReader::open(&filename)?.decode()?.to_rgba8();
    let (w, h) = (image.width(), image.height());
    Ok((image, w, h))
}

/// Make a copy of colors, with color is replaced by the closest match in palette.
/// Uses Pythagorean distance.
fn assign_colors(colors: Vec<Color>, palette: Vec<Color>) -> Vec<Color> {
    // Skip the sqrt, we don't need it. Use u32 to prevent overflows.
    let channel_distance = |a: u8, b:u8| ((if a < b {b - a} else {a - b}) as u32).pow(2);
    let mut result = Vec::new();
    result.extend(colors.iter().map(
            |c| palette.iter().min_by_key(
                |p| ColorChannel::ALL.iter().map(
                    |channel| channel_distance(channel.value(c), channel.value(p))
                    ).sum::<u32>()
                ).unwrap()  // Maybe we should check if the palette is empty?
            ));
    result
}

/// Convert from the local representation to an image::RgbaImage.
/// Could theoretically panic in a variety of ways.
fn gather_pixels(colors: Vec<Color>, width: u32, height: u32) -> RgbaImage {
    let from_c = |c: Color| Rgba::<u8>::from([c.r, c.g, c.b, c.a]);
    RgbaImage::from_fn(width, height,
                       |x, y| from_c(colors[usize::try_from(x + (width * y)).unwrap()])
                       )
}

/// The inner "main" function; wraps failure conditions.
fn handle_file(input_file: String, output_file: String) -> Result<(), String> {
    let e2str = |e| format!("{}", e);
    let (image, width, height) = read_image(input_file).map_err(e2str)?;
    let pixels = read_pixels(image);
    let palette = make_palette(&mut pixels.clone(), 4)
        .ok_or(String::from("There was a problem building the palette."))?;
    println!("{} colors!", palette.len());
    let out_image = gather_pixels(assign_colors(pixels, palette), width, height);
    out_image.save(output_file).map_err(e2str)
}

fn main() {
    let input_file = env::args().nth(1).expect("Missing argument");
    let output_file = env::args().nth(2).expect("Missing argument");
    let result = handle_file(input_file, output_file);
    if let Err(error) = result {
        println!("{}", error);
    } else {
        println!("Success!");
    }
}

On the other hand it now takes about twice as long to run, and I don't know why.

Your questions:

  • The usage of Option.
    Think about what the situation you're trying to represent means. If the situation is "there are no elements", could you just return an empty list? If something's gone wrong, then a Result should be preferred because it will help a user figure out what went wrong. If you're still in the realm of nominal behavior, Option is probably fine.
    I think in general, it would be ideal to make your code more "arrow-like", if I understand correctly what you mean by that.
  • Overall performance.
    LOL, IDK. I have some idea's about how we might improve the performance, but it's above my pay grade :)
  • Algorithm implementation.
    I think I've covered this above.
\$\endgroup\$
4
  • \$\begingroup\$ (assign_colors(): why an abs diff before squaring?) \$\endgroup\$
    – greybeard
    Mar 9, 2022 at 5:43
  • 1
    \$\begingroup\$ additional comments for every argument doesn't seem to have much value for a human reader. There are tools that allow test automation from well-formed comments - and tools creating human readable documentation. \$\endgroup\$
    – greybeard
    Mar 9, 2022 at 5:46
  • \$\begingroup\$ @greybeard: either an absolute diff, or converting to signed integers, is necessary to prevent overflow. In retrospect, that should have been commented, and it would probably have been better to just cast everything to i32 or i64`. \$\endgroup\$ Mar 9, 2022 at 16:35
  • \$\begingroup\$ (I had my C spectacles on where "every" integral expression operand narrower than int does get widened to that, if not even wider because of other operands. Can't seem to find type handling in expression evaluation in doc.rust-lang.org/reference?!) \$\endgroup\$
    – greybeard
    Mar 9, 2022 at 17:55
4
+100
\$\begingroup\$

/*

Hello and welcome to the Rust community!

First, a couple of ideas after briefly looking through the code.

  • The algorithm can be optimized with SIMD operations. Read up on SIMD in Rust if you wish to learn more.
  • You have a struct with fields r, g, b, and a. Unfortunately you cannot do a struct of arrays instead, because you sort colors. Read up on ArrayOfStructs/StructOfArrays if you wish to learn more.
  • You guard against the emptiness of buckets in many places -- this invariant could be baked into a new type for holding buckets. This avoids almost all use of Options.
  • You have many freestanding functions. They do work on a common type. You can make a new type for holding buckets and impl these fns on that type. This could improve readability and clarity of code.
  • The code uses recursion. It is a good practice when working with algorithms to avoid recursion and work with a Vec and loop instead. Reasoning about recursion is sometimes difficult. I've even seen an amazing programmer (Niko Matsakis) make mistakes with recursion. But your recursion should be fine here, since it is a simple algorithm.
  • Use clippy! It provides some small suggestions for making this code in line with best practices.
  • There is fold, which could be refactored as map and sum instead.
  • You do impl PartialEq for Color, which could be automatically derived.
  • The clone in recurse is useless. Also, use extend instead.
  • If you have more than 2^24 pixels, for example in an image larger than 4096x4096, your color_mean may overflow. Use a u64 type there to avoid overflow.
  • recurse is not a good name for that fn. A better name is make_palette as suggested by @ShapeOfMatter.
  • sort_by can become sort_by_key.

Are sure there is any problem with your implementation? I do not see any purple color results on pixels of red and blue.

EDIT: a range is maximum minus minimum, not just maximum.

Here is how the code looks like after a refactor

I added in changes suggested by @ShapeOfMatter.

*/


// ---------
// I have implemented a simple version of the median cut algorithm.
// It takes a vector of Color structs representing pixel in an image.
// I also use the ColorChanel enum representing RGBA channels.
// ---------

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Color {
    pub r: u8,
    pub g: u8,
    pub b: u8,
    pub a: u8,
}

#[derive(Debug, Copy, Clone)]
pub enum ColorChannel {
    R,
    G,
    B,
    A,
}

/// A list of colors.
///
/// INVARIANT: is nonempty.
struct ColorBucket {
    colors: Vec<Color>,
}

impl ::std::ops::Index<ColorChannel> for Color {
    type Output = u8;
    fn index(&self, index: ColorChannel) -> &Self::Output {
        match index {
            ColorChannel::R => &self.r,
            ColorChannel::G => &self.g,
            ColorChannel::B => &self.b,
            ColorChannel::A => &self.a,
        }
    }
}

/// Helper function for calculating the mean.
fn mean(iter: impl Iterator<Item=u8> + Clone) -> u8 {
    (iter.clone().map(|x| x as u64).sum::<u64>() / iter.count() as u64) as u8
}

impl ColorBucket {
    fn from_pixels(pixels: Vec<Color>) -> Option<Self> {
        if pixels.is_empty() {
            None
        } else {
            Some(Self {
                colors: pixels,
            })
        }
    }

    // ---------
    // For each Color/Pixel vector I look for the channel with the highest range:
    // ---------

    /// Returns the color channel with the highest range.
    /// IMPORTANT: Ignores alpha channel!
    ///
    /// # Arguments
    ///
    /// * `colors` - Color vector from which the highest range is evaluated.
    ///
    fn highest_range_channel(&self) -> ColorChannel {
        let ranges = self.color_ranges();

        let mut highest_range_channel = ColorChannel::R;
        let mut highest_value = ranges.r;

        if ranges.g > highest_value {
            highest_range_channel = ColorChannel::G;
            highest_value = ranges.g;
        }

        if ranges.b > highest_value {
            highest_range_channel = ColorChannel::B;
        }

        highest_range_channel
    }


    // ---------
    // The color ranges are calculated in the according function:
    // ---------

    /// Returns the ranges for each color channel
    ///
    /// # Arguments
    ///
    /// * `colors` - Color vector from which the ranges are calculated.
    ///
    /// # Examples
    ///
    /// ```
    /// let colors = Vec::<Color>::new();
    /// let color_ranges_data = color_ranges(colors);
    /// ```
    ///
    fn color_ranges(&self) -> Color {
        // Unwrap is ok here, because `max_by_key` only returns `None` for empty vectors
        let r_range = self.colors.iter().max_by_key(|c| c.r).unwrap().r;
        let g_range = self.colors.iter().max_by_key(|c| c.g).unwrap().g;
        let b_range = self.colors.iter().max_by_key(|c| c.b).unwrap().b;
        let a_range = self.colors.iter().max_by_key(|c| c.a).unwrap().a;

        Color {
            r: r_range,
            g: g_range,
            b: b_range,
            a: a_range,
        }
    }

    // ---------
    // After that I calculate the median value for the channel with the highest range.
    // I am doing this by sorting the vector based on the desired channel and finding
    // the value in the "middle" of the vector.
    // ---------

    /// Sort a color vector for a specific channel.
    ///
    /// # Arguments
    ///
    /// * `colors` - Color data which will be sorted.
    /// * `channel` - Target channel. The sorting is performed based on this value.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut colors = Vec::<Color>::new();
    /// sort_colors(&mut colors, &ColorChannel::R);
    /// ```
    ///
    fn sort_colors(&mut self, channel: ColorChannel) {
        self.colors.sort_by_key(|x| x[channel])
    }

    /// Returns median value for a specific `ColorChannel`.
    ///
    /// # Arguments
    ///
    /// * `colors` - Color vector from which the median value is calculated.
    /// * `channel` - Target channel for which the median is calculated.
    ///
    /// # Examples
    /// ```
    /// let mut colors = Vec::<Color>::new();
    /// let mut result = color_median(&mut colors, &ColorChannel::R);
    /// ```
    ///
    fn color_median(&mut self, channel: ColorChannel) -> u8 {
        self.sort_colors(channel);

        let mid = self.colors.len() / 2;
        if self.colors.len() % 2 == 0 {
            let bucket = ColorBucket::from_pixels(vec![self.colors[mid - 1], self.colors[mid]]).unwrap();
            bucket.channel_mean(channel)
        } else {
            self.channel_value_by_index(mid, channel)
        }
    }

    /// Returns a color value based on the provided channel and index parameters.
    ///
    /// # Arguments
    ///
    /// * `colors` - Color vector from which the value is retreived.
    /// * `index` - Index of the target color in the vector.
    /// * `channel` - Color channel of the searched value.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut colors: Vec<Color> = Vec::new();
    /// colors.push(Color { r: 100, g: 22, b: 12, a: 0 });
    /// assert_eq!(Some(100), channel_value_by_index(&colors, 0, &ColorChannel::R));
    /// ```
    ///
    /// # Panics
    /// 
    /// Panics when index is out of bounds.
    /// 
    fn channel_value_by_index(&self, index: usize, channel: ColorChannel) -> u8 {
        self.colors[index][channel]
    }

    /// Calculate the mean value for a specific color channel on a vector of `Color`.
    ///
    /// # Arguments
    ///
    /// * `colors` - Color vector from which the mean value is calculated.
    /// * `channel` - Target channel for which the mean is calculated.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut colors: Vec<Color> = Vec::new();
    /// let mut result = channel_mean(&colors, &ColorChannel::R);
    /// ```
    ///
    fn channel_mean(&self, channel: ColorChannel) -> u8 {
        mean(self.colors.iter().map(|x| x[channel]))
    }

    // ---------
    // Now I create two new vectors/buckets with one containing all Colors
    // above the median value and another with Color values below the median.
    // This entire process is implemented in the median_cut function.
    // ---------

    /// Performs the median cut on a single vector (bucket) of `Color`.
    /// Returns two `color` vectors representing the colors above and colors below median value.
    ///
    /// # Arguments
    ///
    /// * `colors` - `Color` vector on which the median cut is performed.
    ///
    fn median_cut(&mut self) -> (Option<ColorBucket>, Option<ColorBucket>) {
        let highest_range_channel = self.highest_range_channel();
        let median = self.color_median(highest_range_channel);
        let mut above_median = vec![];
        let mut below_median = vec![];
        for color in &self.colors {
            if color[highest_range_channel] > median {
                above_median.push(*color);
            } else {
                below_median.push(*color);
            }
        }

        (ColorBucket::from_pixels(above_median), ColorBucket::from_pixels(below_median))
    }

    // --------
    // In order to perform multiple iterations I call the recurse function.
    // It takes a bucket, performs the median_cut on it calculates the mean
    // color for each output bucket and performs the median_cut on the them.
    // The function stops and returns all mean colors when the amount of iterations reaches 0.
    // --------

    fn recurse(&mut self, iter_count: u8, result: &mut Vec<Color>) {
        if iter_count == 0 {
            result.push(self.color_mean());
        } else {
            let new_buckets = self.median_cut();
            if let Some(mut bucket) = new_buckets.0 {
                bucket.recurse(iter_count - 1, result);
            }
            if let Some(mut bucket) = new_buckets.1 {
                bucket.recurse(iter_count - 1, result);
            }
        }
    }

    fn make_palette(&mut self, iter_count: u8) -> Vec<Color> {
        let mut result = vec![];
        self.recurse(iter_count, &mut result);
        result
    }

    /// Returns the mean color value based on the passed colors.
    ///
    /// # Arguments
    ///
    /// * `colors` - Color vector from which the mean color is calculated.
    ///
    /// # Examples
    ///
    /// ```
    /// let colors = Vec::<Color>::new();
    /// let result = color_mean(&colors);
    /// ```
    ///
    fn color_mean(&self) -> Color {
        let r = mean(self.colors.iter().map(|c| c.r));
        let g = mean(self.colors.iter().map(|c| c.g));
        let b = mean(self.colors.iter().map(|c| c.b));
        let a = mean(self.colors.iter().map(|c| c.a));

        Color { r, g, b, a }
    }
}

// --------
// So basically a call to the entire algorithm looks like this:
// --------

#[test]
fn test_colors() {
    let mut pixels = Vec::new();
    // fill pixels with data
    pixels.push(Color { r: 100, g: 120, b: 120, a: 0 });
    pixels.push(Color { r: 150, g: 150, b: 150, a: 0 });
    pixels.push(Color { r: 255, g: 255, b: 255, a: 0 });

    let mut bucket = ColorBucket::from_pixels(pixels).expect("empty list");

    let colors = bucket.make_palette(3);
    // Do something with the output colors
    // println!("{:?}", colors);
    let expected = vec![
        Color { r: 255, g: 255, b: 255, a: 0 },
        Color { r: 150, g: 150, b: 150, a: 0 },
        Color { r: 100, g: 120, b: 120, a: 0 },
    ];
    assert_eq!(colors, expected);
}

#[test]
fn test_red_blue() {
    let mut pixels = Vec::new();
    // fill pixels with data
    pixels.push(Color { r: 255, g: 0, b: 0, a: 0 });
    pixels.push(Color { r: 0, g: 255, b: 0, a: 0 });

    let mut bucket = ColorBucket::from_pixels(pixels).expect("empty list");

    let colors = bucket.make_palette(1);
    // Do something with the output colors
    // println!("{:?}", colors);
    let expected = vec![
        Color { r: 255, g: 0, b: 0, a: 0 },
        Color { r: 0, g: 255, b: 0, a: 0 },
    ];
    assert_eq!(colors, expected);
}
\$\endgroup\$
1
  • \$\begingroup\$ First of all, thank you for the detailed answer. While writing tests for the entire thing, I think I might have found a bug in your (and mine) version of color_ranges. The result should be the difference between the max and the min values. We simply return the max value here, which is only correct when min value is 0. @ShapeOfMatter got it right though. \$\endgroup\$ Mar 30, 2022 at 16:37

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.