texture_synthesis/
utils.rs

1use crate::{Dims, Error};
2use std::path::Path;
3
4/// Helper type used to define the source of `ImageSource`'s data
5#[derive(Clone)]
6pub enum ImageSource<'a> {
7    /// A raw buffer of image data, see `image::load_from_memory` for details
8    /// on what is supported
9    Memory(&'a [u8]),
10    /// The path to an image to load from disk. The image format is inferred
11    /// from the file extension, see `image::open` for details
12    Path(&'a Path),
13    /// An already loaded image that is passed directly to the generator
14    Image(image::DynamicImage),
15}
16
17impl<'a> ImageSource<'a> {
18    pub fn from_path(path: &'a Path) -> Self {
19        Self::Path(path)
20    }
21}
22
23impl<'a> From<image::DynamicImage> for ImageSource<'a> {
24    fn from(img: image::DynamicImage) -> Self {
25        Self::Image(img)
26    }
27}
28
29impl<'a, S> From<&'a S> for ImageSource<'a>
30where
31    S: AsRef<Path> + 'a,
32{
33    fn from(path: &'a S) -> Self {
34        Self::Path(path.as_ref())
35    }
36}
37
38pub fn load_dynamic_image(src: ImageSource<'_>) -> Result<image::DynamicImage, image::ImageError> {
39    match src {
40        ImageSource::Memory(data) => image::load_from_memory(data),
41        ImageSource::Path(path) => image::open(path),
42        ImageSource::Image(img) => Ok(img),
43    }
44}
45
46/// Helper type used to mask `ImageSource`'s channels
47#[derive(Clone, Copy)]
48pub enum ChannelMask {
49    R,
50    G,
51    B,
52    A,
53}
54
55pub(crate) fn load_image(
56    src: ImageSource<'_>,
57    resize: Option<Dims>,
58) -> Result<image::RgbaImage, Error> {
59    let img = load_dynamic_image(src)?;
60
61    let img = match resize {
62        None => img.to_rgba8(),
63        Some(ref size) => {
64            use image::GenericImageView;
65
66            if img.width() != size.width || img.height() != size.height {
67                image::imageops::resize(
68                    &img.to_rgba8(),
69                    size.width,
70                    size.height,
71                    image::imageops::CatmullRom,
72                )
73            } else {
74                img.to_rgba8()
75            }
76        }
77    };
78
79    Ok(img)
80}
81
82pub(crate) fn apply_mask(mut image: image::RgbaImage, mask: ChannelMask) -> image::RgbaImage {
83    let channel = match mask {
84        ChannelMask::R => 0,
85        ChannelMask::G => 1,
86        ChannelMask::B => 2,
87        ChannelMask::A => 3,
88    };
89
90    for pixel_iter in image.enumerate_pixels_mut() {
91        let pixel = pixel_iter.2;
92        pixel[0] = pixel[channel];
93        pixel[1] = pixel[channel];
94        pixel[2] = pixel[channel];
95        pixel[3] = 255;
96    }
97
98    image
99}
100
101pub(crate) fn transform_to_guide_map(
102    image: image::RgbaImage,
103    size: Option<Dims>,
104    blur_sigma: f32,
105) -> image::RgbaImage {
106    use image::GenericImageView;
107    let dyn_img = image::DynamicImage::ImageRgba8(image);
108
109    if let Some(s) = size {
110        if dyn_img.width() != s.width || dyn_img.height() != s.height {
111            dyn_img.resize(s.width, s.height, image::imageops::Triangle);
112        }
113    }
114
115    dyn_img.blur(blur_sigma).grayscale().to_rgba8()
116}
117
118pub(crate) fn get_histogram(img: &image::RgbaImage) -> Vec<u32> {
119    let mut hist = vec![0; 256]; //0-255 incl
120
121    let pixels = &img;
122
123    //populate the hist
124    for pixel_value in pixels
125        .iter()
126        .step_by(/*since RGBA image, we only care for 1st channel*/ 4)
127    {
128        hist[*pixel_value as usize] += 1; //increment histogram by 1
129    }
130
131    hist
132}
133
134//source will be modified to fit the target
135pub(crate) fn match_histograms(source: &mut image::RgbaImage, target: &image::RgbaImage) {
136    let target_hist = get_histogram(target);
137    let source_hist = get_histogram(source);
138
139    //get commutative distrib
140    let target_cdf = get_cdf(&target_hist);
141    let source_cdf = get_cdf(&source_hist);
142
143    //clone the source image, modify and return
144    let (dx, dy) = source.dimensions();
145
146    for x in 0..dx {
147        for y in 0..dy {
148            let pixel_value = source.get_pixel(x, y)[0]; //we only care about the first channel
149            let pixel_source_cdf = source_cdf[pixel_value as usize];
150
151            //now need to find by value similar cdf in the target
152            let new_pixel_val = target_cdf
153                .iter()
154                .position(|cdf| *cdf > pixel_source_cdf)
155                .unwrap_or((pixel_value + 1) as usize) as u8
156                - 1;
157
158            let new_color: image::Rgba<u8> =
159                image::Rgba([new_pixel_val, new_pixel_val, new_pixel_val, 255]);
160            source.put_pixel(x, y, new_color);
161        }
162    }
163}
164
165pub(crate) fn get_cdf(a: &[u32]) -> Vec<f32> {
166    let mut cumm = vec![0.0; 256];
167
168    for i in 0..a.len() {
169        if i != 0 {
170            cumm[i] = cumm[i - 1] + (a[i] as f32);
171        } else {
172            cumm[i] = a[i] as f32;
173        }
174    }
175
176    //normalize
177    let max = cumm[255];
178    for i in cumm.iter_mut() {
179        *i /= max;
180    }
181
182    cumm
183}