lf_watermark/
lib.rs

1use std::error::Error;
2
3use image::{DynamicImage, GenericImageView, Rgb, RgbImage};
4use rustdct::DctPlanner;
5
6type Result<T> = std::result::Result<T, Box<dyn Error>>;
7
8pub fn get_watermark_from_str(words: &str) -> Result<f32> {
9    let char_map =
10        "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*(),.<>/?; ";
11
12    let mut ret = 0.0;
13
14    for i in 0..words.len() {
15        let c = words.chars().nth(i).ok_or("invalid index")?;
16        let char_idx = char_map
17            .find(c)
18            .ok_or(format!("Invalid character; {}", c))?;
19        ret += char_idx as f32;
20    }
21
22    Ok(ret
23        * option_env!("WATERMARK_STRENGTH")
24            .unwrap_or("0.01")
25            .parse::<f32>()?)
26}
27
28pub fn embed_watermark_color(image: &DynamicImage, watermark: &str) -> Result<RgbImage> {
29    let watermark = get_watermark_from_str(watermark)?;
30
31    let (width, height) = image.dimensions();
32    let len = (width * height) as usize;
33    let mut cbcr_channel = vec![(0, 0); len];
34    let mut y_channel = vec![0.0; len];
35    let idx_fn = |x: u32, y: u32| (y * width + x) as usize;
36    let normalization_factor = (2.0 / len as f32).sqrt();
37
38    let image = image.to_rgb8();
39
40    for (x, y, pixel) in image.enumerate_pixels() {
41        let idx = idx_fn(x, y);
42
43        let (y, u, v) = rgb_to_ycbcr(&pixel);
44        cbcr_channel[idx] = (u, v);
45        y_channel[idx] = y as f32 + watermark;
46    }
47
48    let mut dct_planner: DctPlanner<f32> = DctPlanner::new();
49    let dct = dct_planner.plan_dct2(len);
50    dct.process_dct2(&mut y_channel);
51
52    for y in y_channel.iter_mut() {
53        *y *= normalization_factor;
54    }
55
56    let idct = dct_planner.plan_dct3(len);
57    idct.process_dct3(&mut y_channel);
58    for y in y_channel.iter_mut() {
59        *y *= normalization_factor;
60    }
61
62    let mut img_buffer = RgbImage::new(width, height);
63    for (x, y, pixel) in img_buffer.enumerate_pixels_mut() {
64        let index = idx_fn(x, y);
65        let y_ch = y_channel[index];
66        let (cb, cr) = cbcr_channel[index];
67
68        *pixel = ycbcr_to_rgb(y_ch, cb as f32, cr as f32);
69    }
70
71    Ok(img_buffer)
72}
73
74fn rgb_to_ycbcr(pixel: &Rgb<u8>) -> (u8, u8, u8) {
75    let r = pixel[0] as f64;
76    let g = pixel[1] as f64;
77    let b = pixel[2] as f64;
78
79    let y = (0.299 * r + 0.587 * g + 0.114 * b).round() as u8;
80    let cb = (-0.169 * r - 0.331 * g + 0.5 * b + 128.0).round() as u8;
81    let cr = (0.5 * r - 0.419 * g - 0.081 * b + 128.0).round() as u8;
82
83    (y, cb, cr)
84}
85
86fn ycbcr_to_rgb(y: f32, cb: f32, cr: f32) -> Rgb<u8> {
87    let r = (y + 1.402 * (cr as f32 - 128.0)).round() as u8;
88    let g = (y - 0.34414 * (cb as f32 - 128.0) - 0.71414 * (cr as f32 - 128.0)).round() as u8;
89    let b = (y + 1.772 * (cb as f32 - 128.0)).round() as u8;
90
91    Rgb([r, g, b])
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn test_get_watermark_from_str() {
100        let words = "Hello, World!";
101        let bytes = get_watermark_from_str(words).unwrap();
102        assert_eq!(bytes, 5.35);
103    }
104
105    #[test]
106    fn test_rgb_to_ycbcr() {
107        // NOTE: this ycbcr conversion make a little changes to the original rgb value
108        for (r, g, b) in vec![
109            (255, 255, 255),
110            (254, 0, 0),
111            (0, 255, 1),
112            (0, 0, 254),
113            (0, 0, 0),
114            (128, 128, 128),
115        ] {
116            let (y, cb, cr) = rgb_to_ycbcr(&Rgb([r, g, b]));
117            let color = ycbcr_to_rgb(y as f32, cb as f32, cr as f32);
118
119            assert_eq!(Rgb([r, g, b]), color, "rgb: {:?} {:?}", (r, g, b), color);
120        }
121    }
122
123    #[test]
124    fn test_watermark() {
125        let img = image::open("image.png").unwrap();
126        let watermark = "Hello, World!";
127        let watermarked_img = embed_watermark_color(&img, watermark);
128        assert!(watermarked_img.is_ok(), "Failed to embed watermark");
129        assert!(
130            watermarked_img.unwrap().save("output.png").is_ok(),
131            "Failed to save image"
132        )
133    }
134}
lf_watermark/lib.rs

lf_watermark/
lib.rs
