channel wise enc

Author: MenxLi

src/lib.rs

@@ -23,6 +23,7 @@ macro_rules! console_log {
 pub struct ImageOptions {
     pub width: u32,
     pub height: u32,
+    pub channels: u32,
 }
 
 fn str2f(s: &str) -> f64 {
@@ -33,7 +34,7 @@ fn str2f(s: &str) -> f64 {
     x
 }
 
-fn img2vec(im: &[u8], limit_max_side: Option<u32>) -> (Vec<[u8; 3]>, ImageOptions) {
+fn img2vec(im: &[u8], limit_max_side: Option<u32>) -> (Vec<u8>, ImageOptions) {
     let img = image::load_from_memory(im).expect("Failed to load image");
     let mut rgb = img.to_rgb8();
 
@@ -50,25 +51,32 @@ fn img2vec(im: &[u8], limit_max_side: Option<u32>) -> (Vec<[u8; 3]>, ImageOption
         }
     }
 
-    let mut pixels = Vec::with_capacity(rgb.width() as usize * rgb.height() as usize);
+    let mut pixels = Vec::with_capacity(rgb.width() as usize * rgb.height() as usize * 3);
     for pixel in rgb.pixels() {
-        pixels.push([pixel[0], pixel[1], pixel[2]]);
+        pixels.push(pixel[0]);
+        pixels.push(pixel[1]);
+        pixels.push(pixel[2]);
     };
 
     (pixels, ImageOptions {
         width: rgb.width(),
         height: rgb.height(),
+        channels: 3,
     })
 }
 
-fn vec2pngblob(pixels: &Vec<[u8; 3]>, im_opt: ImageOptions, limit_max_side: Option<u32>) -> Box<[u8]> {
-    let ImageOptions { mut width, mut height } = im_opt;
+fn vec2pngblob(pixels: &Vec<u8>, im_opt: ImageOptions, limit_max_side: Option<u32>) -> Box<[u8]> {
+    let ImageOptions { mut width, mut height , channels} = im_opt;
 
     let mut img = image::RgbImage::new(width, height);
-    for (i, pixel) in pixels.iter().enumerate() {
+    for (i, pixel) in pixels.chunks(channels as usize).enumerate() {
         let x = (i % width as usize) as u32;
         let y = (i / width as usize) as u32;
-        img.put_pixel(x, y, image::Rgb(*pixel));
+        img.put_pixel(x, y, image::Rgb([
+            pixel[0],
+            pixel[1],
+            pixel[2],
+        ]));
     };
 
     if let Some(max_side) = limit_max_side {

src/logistic_map.rs

@@ -27,27 +27,22 @@ fn argsort<T: Ord>(vec: &Vec<T>) -> Vec<usize> {
     indices
 }
 
-fn confuse_xor(im: &Vec<[u8; 3]>, enc_map: &Vec<f64>) -> Vec<[u8; 3]> {
+fn confuse_xor(im: &Vec<u8>, enc_map: &Vec<f64>) -> Vec<u8> {
     let mut enc_pixels = Vec::with_capacity(im.len());
-
     for (i, pixel) in enc_map.iter().map(
         |&x| (x * 255.0) as u8
     ).enumerate() {
-        // use xor to encode the pixel
-        let r = im[i][0] ^ pixel;
-        let g = im[i][1] ^ pixel;
-        let b = im[i][2] ^ pixel;
-        enc_pixels.push([r, g, b]);
+        enc_pixels.push(im[i] ^ pixel);
     }
-
     enc_pixels
 }
 
 enum DiffuseDirection {
     Forward,
     Backward,
 }
-fn diffuse(im: &Vec<[u8; 3]>, enc_map: &Vec<f64>, direction: DiffuseDirection) -> Vec<[u8; 3]> {
+fn diffuse<T: Copy>(im: &Vec<T>, enc_map: &Vec<f64>, direction: DiffuseDirection) -> Vec<T> 
+{
     let enc_map = enc_map
         .iter()
         .map(|&x| x.to_bits())
@@ -78,7 +73,7 @@ fn diffuse(im: &Vec<[u8; 3]>, enc_map: &Vec<f64>, direction: DiffuseDirection) -
     diffuse_pixels
 }
 
-pub fn encode(im: &Vec<[u8; 3]>, x0: f64) -> Vec<[u8; 3]> {
+pub fn encode(im: &Vec<u8>, x0: f64) -> Vec<u8> {
     let enc_map = generate_map(LogisticMapOptions {
         x: x0,
         r: R,  
@@ -89,7 +84,7 @@ pub fn encode(im: &Vec<[u8; 3]>, x0: f64) -> Vec<[u8; 3]> {
     confuse_xor(&im, &enc_map)
 }
 
-pub fn decode(im: &Vec<[u8; 3]>, x0: f64) -> Vec<[u8; 3]> {
+pub fn decode(im: &Vec<u8>, x0: f64) -> Vec<u8> {
     let enc_map = generate_map(LogisticMapOptions {
         x: x0,
         r: R,