Merge pull request #1 from studio-YOLO/wasm

merge: wasm branch into main

.github/workflows/wasm-pack.yml

@@ -0,0 +1,47 @@
+name: wasm-pack workflow
+on:
+  push:
+    paths:
+      - 'lib/src/lib.rs'
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Check out repository
+      uses: actions/checkout@v2
+
+    - name: Compile Rust project with wasm-pack
+      run: |
+        curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
+        wasm-pack build lib/ --release --target web
+
+    - name: Copy necessary files
+      run: |
+        cp lib/pkg/editpix_wasm_bg.wasm src/core/editpix_wasm_bg.wasm
+        cp lib/pkg/editpix_wasm.js src/core/editpix_wasm.js
+        rm -r lib/pkg
+
+    - name: Check for changes
+      id: check_changes
+      run: |
+        if git diff --quiet --exit-code; then
+          echo "No changes detected. Exiting workflow."
+          echo "::set-output name=changes::false"
+          exit 0; # Exit with success code
+        else
+          echo "::set-output name=changes::true"
+        fi
+
+    - name: Commit build results
+      if: steps.check_changes.outputs.changes == 'true'
+      run: |
+        git config --local user.email "github-actions[bot]@users.noreply.github.com"
+        git config --local user.name "github-actions[bot]"
+        git add *
+        git commit -m "CI push"
+
+    - name: Push changes
+      if: steps.check_changes.outputs.changes == 'true'
+      uses: ad-m/github-push-action@master

.gitignore

@@ -1,2 +1,6 @@
 node_modules
-coverage
+coverage
+.DS_Store
+lib/target
+lib/pkg
+lib/Cargo.lock

lib/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "editpix-wasm"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+wasm-bindgen = "0.2"
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[dependencies.web-sys]
+version = "0.3.68"
+features = [
+  'Document',
+  'Element',
+  'HtmlElement',
+  'Node',
+  'Window',
+]
+

lib/src/.gitignore

@@ -0,0 +1,3 @@
+target/
+pkg/
+.DS_Store

lib/src/lib.rs

@@ -0,0 +1,108 @@
+extern crate wasm_bindgen;
+use wasm_bindgen::prelude::*;
+extern crate web_sys;
+
+#[inline]
+fn euclidean_distance(color1: &[u8; 3], color2: &[u8; 3]) -> f64 {
+    f64::sqrt(
+        ((color1[0] as i16 - color2[0] as i16) as f64).powi(2) +
+        ((color1[1] as i16 - color2[1] as i16) as f64).powi(2) +
+        ((color1[2] as i16 - color2[2] as i16) as f64).powi(2)
+    )
+}
+
+#[inline(always)]
+fn initialize_centroids(colors: &[[u8; 3]], color_number: usize) -> Vec<[u8; 3]> {
+    let mut centroids: Vec<[u8; 3]> = Vec::new();
+
+    for i in 0..color_number {
+        centroids.push(colors[i * colors.len() / color_number])
+    }
+
+    centroids
+}
+
+fn assign_to_centroids(colors: &[[u8; 3]], centroids: Vec<[u8; 3]>) -> Vec<usize> {
+    let mut assignments: Vec<usize> = Vec::new();
+    for i in 0..colors.len() {
+        let mut min_distance = f64::INFINITY;
+        let mut closest_centroid = 3;
+        for j in 0..centroids.len() {
+            let distance = euclidean_distance(colors.get(i).unwrap(), centroids.get(j).unwrap());
+            if distance < min_distance {
+                min_distance = distance;
+                closest_centroid = j;
+            }
+        }
+        assignments.push(closest_centroid);
+    }
+
+    assignments
+}
+
+
+fn calculate_new_centroids(colors: &[[u8; 3]], assignments: &[usize], color_number: usize) -> Vec<[u8; 3]> {
+    let mut new_centroids: Vec<[u32; 3]> = vec![[0,0,0]; color_number];   
+    let mut counts = vec![0; color_number];
+
+
+    for i in 0..colors.len() {
+        let assignment = assignments[i];
+        new_centroids[assignment][0] += colors[i][0] as u32;
+        new_centroids[assignment][1] += colors[i][1] as u32;
+        new_centroids[assignment][2] += colors[i][2] as u32;
+        counts[assignment]+=1;
+    }
+
+    for  j in  0..color_number  {
+        if counts[j] > 0 {
+          new_centroids[j][0] /= counts[j];
+          new_centroids[j][1] /= counts[j];
+          new_centroids[j][2] /= counts[j];
+        }
+    }
+
+    let vec_u8: Vec<[u8; 3]> = new_centroids
+        .iter()
+        .map(|&arr| [arr[0] as u8, arr[1] as u8, arr[2] as u8])
+        .collect();
+
+    vec_u8
+
+}
+
+#[wasm_bindgen]
+pub fn k_means(colors_r: Vec<u8>, color_number: usize, max_iterations: usize) -> Vec<u8> {
+    let colors: Vec<[u8; 3]> = colors_r
+        .chunks_exact(3) // Get chunks of 3 elements
+        .map(|chunk| {
+            let mut array: [u8; 3] = [0; 3];
+            array.copy_from_slice(chunk);
+            array
+        })
+        .collect();
+
+
+    let mut centroids = initialize_centroids(&colors, color_number);
+
+    let mut iterations: usize = 0;
+    let mut previous_assignments;
+    let mut assignments: Vec<usize> = Vec::new();
+
+    loop {
+        previous_assignments = assignments.clone();
+        assignments = assign_to_centroids(&colors, centroids);
+        centroids = calculate_new_centroids(&colors, &assignments, color_number);
+        iterations += 1;
+        if !(iterations < max_iterations && assignments != previous_assignments) {
+            break;
+        }
+    }
+
+    let serialized_vector: Vec<u8> = centroids
+        .into_iter()
+        .flat_map(|array| array.into_iter())
+        .collect();
+
+    serialized_vector
+}

src/editpix.js

@@ -5,6 +5,7 @@ import convertFromGrayToRgb from "./core/gray_to_rgb.js";
 import kMeans from "./core/kmean.js";
 import imageManager from "./image_manager.js";
 import higherColorContrast from "./core/higher_contrast.js";
+import init, { k_means } from "./core/editpix_wasm.js"
 
 var EditPix = function () { };
 
@@ -20,6 +21,14 @@ EditPix.prototype.getColorPalette = (image, colorNumber = 5, quality = 1) => {
     })
 }
 
+EditPix.prototype.getColorPaletteWasm = async (image, colorNumber = 5, quality = 1) => {
+    utils.validate(quality, colorNumber);
+    const pixelArray = utils.removeAlphaSerialized(imageManager.getPixelArray(image));
+    await init();
+    let a = k_means(pixelArray, colorNumber, quality * 10);
+    return utils.deserializeArray(a);
+}
+
 EditPix.prototype.getDominantColor = function(image, quality = 1) {
     return this.getColorPalette(image, 1, quality);
 }

src/utils.js

@@ -6,6 +6,14 @@ function removeAlpha(pixelArray) {
     return result
 }
 
+function removeAlphaSerialized(pixelArray) {
+    let result = [];
+    for (let i = 0; i < pixelArray.length / 4; i++) {
+        result.push(pixelArray[i * 4], pixelArray[i * 4 + 1], pixelArray[i * 4 + 2]);
+    }
+    return result
+}
+
 function rgbToHex(rgbColors) {
     let hexColors = [];
     rgbColors.forEach(color => {
@@ -40,9 +48,19 @@ function validate(quality, colorNumber) {
     }
 }
 
+function deserializeArray(serializedArray, chunkSize) {
+    const result = [];
+    for (let i = 0; i < serializedArray.length; i += chunkSize) {
+        result.push(serializedArray.slice(i, i + chunkSize));
+    }
+    return result;
+}
+
 export default {
     rgbToHex,
     hexToRgb,
     validate,
-    removeAlpha
+    removeAlpha,
+    removeAlphaSerialized,
+    deserializeArray
 };