imgfx/
filter.rs

1use std::str::FromStr;
2
3use crate::{calc_luminance, rgb_to_hsv, utils::get_channel_by_name_rgba_u8};
4use image::{DynamicImage, GenericImageView, ImageBuffer, Rgba, RgbaImage};
5use rayon::prelude::*;
6
7/// Specify whether the filter should replace colors that are INCLUDED in the range or EXCLUDED
8/// from the range.
9#[derive(Copy, Clone)]
10pub enum FilterType {
11    Include,
12    Exclude,
13}
14
15/// Clap FromStr
16impl FromStr for FilterType {
17    type Err = String;
18
19    fn from_str(s: &str) -> Result<Self, Self::Err> {
20        match s.to_lowercase().as_str() {
21            "include" => Ok(FilterType::Include),
22            "exclude" => Ok(FilterType::Exclude),
23
24            _ => Err(format!("Invalid FilterType name: {}", s)),
25        }
26    }
27}
28
29/// What property to filter by? Minimum and maximum values vary by property. For example, hue is
30/// 0-360, while red is 0-255.
31#[derive(Copy, Clone)]
32pub enum FilterParam {
33    Luminance,
34    Red,
35    Green,
36    Blue,
37    Hue,
38    Saturation,
39    Value,
40}
41
42/// Clap FromStr
43impl FromStr for FilterParam {
44    type Err = String;
45
46    fn from_str(s: &str) -> Result<Self, Self::Err> {
47        match s.to_lowercase().as_str() {
48            "luminance" => Ok(FilterParam::Luminance),
49            "red" => Ok(FilterParam::Red),
50            "green" => Ok(FilterParam::Green),
51            "blue" => Ok(FilterParam::Blue),
52            "hue" => Ok(FilterParam::Hue),
53            "saturation" => Ok(FilterParam::Saturation),
54            "value" => Ok(FilterParam::Value),
55            "l" => Ok(FilterParam::Luminance),
56            "r" => Ok(FilterParam::Red),
57            "g" => Ok(FilterParam::Green),
58            "b" => Ok(FilterParam::Blue),
59            "h" => Ok(FilterParam::Hue),
60            "s" => Ok(FilterParam::Saturation),
61            "v" => Ok(FilterParam::Value),
62
63            _ => Err(format!("Invalid FilterParam name: {}", s)),
64        }
65    }
66}
67
68/// A threshold range that the filter will check in between. This is a dedicated struct because
69/// for a CLI frontend, I want to minimize String usage after the initial arg parsing.
70#[derive(Clone, Copy, Debug)]
71pub struct ThresholdRange {
72    min: f64,
73    max: f64,
74}
75
76/// The filter to perform on the image.
77pub struct Filter {
78    pub filter_type: FilterType,
79    pub filter_param: FilterParam,
80    pub threshold_ranges: Vec<ThresholdRange>,
81}
82
83/// Parse a vector of strings which contain a number into a vector of ThresholdRanges.
84pub fn parse_filter_vec(thresholds_str_vec: Vec<String>) -> Vec<ThresholdRange> {
85    let mut thresholds: Vec<ThresholdRange> = vec![];
86
87    let mut iter = thresholds_str_vec.iter();
88    while let Some(min_str) = iter.next() {
89        if let Some(max_str) = iter.next() {
90            let min = min_str
91                .parse::<f64>()
92                .expect("Failed to parse min threshold as f64");
93            let max = max_str
94                .parse::<f64>()
95                .expect("Failed to parse max threshold as f64");
96
97            thresholds.push(ThresholdRange { min, max });
98        } else {
99            eprintln!(
100                "Warning: Threshold range input has an unmatched min value: {}",
101                min_str
102            );
103        }
104    }
105
106    thresholds
107}
108
109/// Generate the closure which returns a boolean whether the Rgba<u8> satisfies the filter.
110fn generate_filter(filter: Filter) -> impl Fn(&Rgba<u8>) -> bool {
111    move |pixel| match filter.filter_param {
112        FilterParam::Luminance => {
113            let luminance = calc_luminance(*pixel);
114            filter
115                .threshold_ranges
116                .iter()
117                .any(|range| luminance > range.min && luminance < range.max)
118        }
119        FilterParam::Red => filter.threshold_ranges.iter().any(|range| {
120            let red = pixel.0[0] as f64;
121            red > range.min && red < range.max
122        }),
123        FilterParam::Green => filter.threshold_ranges.iter().any(|range| {
124            let green = pixel.0[1] as f64;
125            green > range.min && green < range.max
126        }),
127        FilterParam::Blue => filter.threshold_ranges.iter().any(|range| {
128            let blue = pixel.0[2] as f64;
129            blue > range.min && blue < range.max
130        }),
131        FilterParam::Hue => {
132            let (h, _, _) = rgb_to_hsv(*pixel);
133            filter
134                .threshold_ranges
135                .iter()
136                .any(|range| h > range.min && h < range.max)
137        }
138        FilterParam::Saturation => {
139            let (_, s, _) = rgb_to_hsv(*pixel);
140            filter
141                .threshold_ranges
142                .iter()
143                .any(|range| s > range.min && s < range.max)
144        }
145        FilterParam::Value => {
146            let (_, _, v) = rgb_to_hsv(*pixel);
147            filter
148                .threshold_ranges
149                .iter()
150                .any(|range| v > range.min && v < range.max)
151        }
152    }
153}
154
155/// Perform the filter operation on the image. lhs will remap the colors before filtering.
156pub fn filter(
157    img: DynamicImage,
158    lhs: Option<Vec<String>>,
159    filter: Filter,
160    replace_with: Rgba<u8>,
161) -> RgbaImage {
162    let (width, height) = img.dimensions();
163
164    let mut output: RgbaImage = ImageBuffer::new(width, height);
165
166    let filter_sorter = generate_filter(Filter {
167        filter_type: filter.filter_type,
168        filter_param: filter.filter_param,
169        threshold_ranges: filter.threshold_ranges,
170    });
171
172    output.par_enumerate_pixels_mut().for_each(|(x, y, pixel)| {
173        let in_pixel = img.get_pixel(x, y);
174
175        // Parse lhs
176        let lhs = match lhs {
177            Some(ref lhs) => (
178                get_channel_by_name_rgba_u8(&lhs[0], &in_pixel),
179                get_channel_by_name_rgba_u8(&lhs[1], &in_pixel),
180                get_channel_by_name_rgba_u8(&lhs[2], &in_pixel),
181            ),
182            None => (in_pixel[0], in_pixel[1], in_pixel[2]),
183        };
184
185        // Include or exclude
186        match filter.filter_type {
187            FilterType::Include => {
188                // Remaps the lhs here
189                if !filter_sorter(&Rgba([lhs.0, lhs.1, lhs.2, 255u8])) {
190                    *pixel = replace_with;
191                } else {
192                    *pixel = in_pixel;
193                }
194            }
195
196            // Or here
197            FilterType::Exclude => {
198                if filter_sorter(&Rgba([lhs.0, lhs.1, lhs.2, 255u8])) {
199                    *pixel = replace_with;
200                } else {
201                    *pixel = in_pixel;
202                }
203            }
204        }
205    });
206
207    output
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213    use image::{Pixel, Rgb};
214    use std::env;
215    use std::path::PathBuf;
216
217    fn get_file_path(file_name: String) -> PathBuf {
218        let mut path = env::current_dir().expect("Failed to get current directory");
219        path.push("assets/control-images/");
220        path.push(file_name);
221        path
222    }
223
224    fn load_image(file_name: String) -> DynamicImage {
225        let path = get_file_path(file_name);
226        let img = image::open(path).expect("Failed to open image.");
227
228        img
229    }
230
231    fn get_color_from_control(img: DynamicImage) -> Rgb<u8> {
232        let pixel = img.get_pixel(0, 0);
233        return pixel.to_rgb();
234    }
235}
imgfx/filter.rs

imgfx/
filter.rs
