Skip to content

Commit

Permalink
Polyfill rayon in WASM
Browse files Browse the repository at this point in the history
  • Loading branch information
kornelski committed Feb 12, 2022
1 parent 5666d5d commit e9d0a53
Show file tree
Hide file tree
Showing 8 changed files with 120 additions and 13 deletions.
6 changes: 4 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,17 @@ name = "dssim"
path = "src/main.rs"

[dependencies]
dssim-core = { path = "./dssim-core", version = "3.1" }
dssim-core = { path = "./dssim-core", version = "3.2", default-features = false }
imgref = "1.9.1"
getopts = "0.2.21"
rayon = "1.5.1"
rayon = { version = "1.5.1", optional = true }
rgb = "0.8.31"
lodepng = "3.5.2"
load_image = { version = "2.16.1", features = ["lcms2-static"] }

[features]
default = ["threads", "dssim-core/default"]
threads = ["rayon"]
avif = ["load_image/avif"]
webp = ["load_image/webp"]
webp-static = ["load_image/webp-static"]
Expand Down
8 changes: 8 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -88,3 +88,11 @@ DSSIM is dual-licensed under [AGPL](LICENSE) or [commercial](https://supso.org/p
* The lightness component of SSIM is ignored when comparing color channels.
* SSIM score is pooled using mean absolute deviation. You can get per-pixel SSIM from the API to implement custom pooling.


## Compiling for WASM

For compatibility with single-threaded WASM runtimes, disable the `threads` Cargo feature. It's enabled by default, so to disable it, disable default features:

```toml
dssim-core = { version = "3.2", default-features = false }
```
6 changes: 5 additions & 1 deletion dssim-core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,17 @@ edition = "2018"
crate-type = ["lib", "staticlib"]

[dependencies]
rayon = "1.5.0"
imgref = "1.9.1"
itertools = "0.10.3"
rayon = { version = "1.5.1", optional = true }
rgb = "0.8.31"

[dev-dependencies]
lodepng = "3.6.0"

[features]
default = ["threads"]
threads = ["rayon"]

[package.metadata.docs.rs]
targets = ["x86_64-unknown-linux-gnu"]
12 changes: 7 additions & 5 deletions dssim-core/src/dssim.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ pub use crate::tolab::ToLABBitmap;
pub use crate::val::Dssim as Val;
use imgref::*;
use itertools::multizip;
#[cfg(not(feature = "threads"))]
use crate::lieon as rayon;
use rayon::prelude::*;
use rgb::{RGB, RGBA};
use std::borrow::Borrow;
Expand Down Expand Up @@ -293,8 +295,8 @@ impl Dssim {
/// `Val` is a fancy wrapper for `f64`
pub fn compare<M: Borrow<DssimImage<f32>>>(&self, original_image: &DssimImage<f32>, modified_image: M) -> (Val, Vec<SsimMap>) {
let modified_image = modified_image.borrow();
let res: Vec<_> = self.scale_weights.par_iter().cloned().zip(
modified_image.scale.par_iter().zip(original_image.scale.par_iter())
let res: Vec<_> = self.scale_weights.as_slice().par_iter().cloned().zip(
modified_image.scale.as_slice().par_iter().zip(original_image.scale.as_slice().par_iter())
).enumerate().map(|(n, (weight, (modified_image_scale, original_image_scale)))| {
let scale_width = original_image_scale.chan[0].width;
let scale_height = original_image_scale.chan[0].height;
Expand Down Expand Up @@ -395,9 +397,9 @@ impl Dssim {
debug_assert_eq!(img1_img2_blur.len(), original.mu.len());
debug_assert_eq!(img1_img2_blur.len(), original.img_sq_blur.len());

let mu_iter = original.mu.par_iter().cloned().zip_eq(modified.mu.par_iter().cloned());
let sq_iter = original.img_sq_blur.par_iter().cloned().zip_eq(modified.img_sq_blur.par_iter().cloned());
img1_img2_blur.par_iter().cloned().zip_eq(mu_iter).zip_eq(sq_iter).zip_eq(map_out.par_iter_mut())
let mu_iter = original.mu.as_slice().par_iter().cloned().zip_eq(modified.mu.as_slice().par_iter().cloned());
let sq_iter = original.img_sq_blur.as_slice().par_iter().cloned().zip_eq(modified.img_sq_blur.as_slice().par_iter().cloned());
img1_img2_blur.par_iter().cloned().zip_eq(mu_iter).zip_eq(sq_iter).zip_eq(map_out.as_mut_slice().par_iter_mut())
.for_each(|(((img1_img2_blur, (mu1, mu2)), (img1_sq_blur, img2_sq_blur)), map_out)| {
let mu1mu1 = mu1 * mu1;
let mu1mu2 = mu1 * mu2;
Expand Down
2 changes: 2 additions & 0 deletions dssim-core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ mod image;
mod linear;
mod tolab;
mod val;
#[cfg(not(feature = "threads"))]
mod lieon;

pub use crate::dssim::*;
pub use crate::image::*;
Expand Down
76 changes: 76 additions & 0 deletions dssim-core/src/lieon.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
//! Shim for single-threaded rayon replacement
pub mod prelude {
pub use super::*;
}

pub trait ParSliceLie<T> {
fn par_chunks(&self, n: usize) -> std::slice::Chunks<'_, T>;
}

pub trait ParSliceMutLie<T> {
fn par_chunks_mut(&mut self, n: usize) -> std::slice::ChunksMut<'_, T>;
}

pub trait ParIntoIterLie<T> {
type IntoIter;
fn into_par_iter(self) -> Self::IntoIter;
}

pub trait ParIterLie<T> {
type Iter;
fn par_iter(&self) -> Self::Iter;
}

pub trait ParIterMutLie<'a, T> {
type Iter;
fn par_iter_mut(&'a mut self) -> Self::Iter;
}

pub fn join<A, B>(a: impl FnOnce() -> A, b: impl FnOnce() -> B) -> (A, B) {
let a = a();
let b = b();
(a, b)
}

impl<'a, T> ParSliceLie<T> for &'a [T] {
fn par_chunks(&self, n: usize) -> std::slice::Chunks<'_, T> {
self.chunks(n)
}
}

impl<'a, T> ParSliceLie<T> for &'a mut [T] {
fn par_chunks(&self, n: usize) -> std::slice::Chunks<'_, T> {
self.chunks(n)
}
}

impl<'a, T> ParSliceMutLie<T> for &'a mut [T] {
fn par_chunks_mut(&mut self, n: usize) -> std::slice::ChunksMut<'_, T> {
self.chunks_mut(n)
}
}

impl<'a, T> ParIterLie<T> for &'a [T] {
type Iter = std::slice::Iter<'a, T>;

fn par_iter(&self) -> Self::Iter {
self.iter()
}
}

impl<'a, T> ParIterMutLie<'a, T> for &'a mut [T] {
type Iter = std::slice::IterMut<'a, T>;

fn par_iter_mut(&'a mut self) -> Self::Iter {
self.iter_mut()
}
}


impl<T> ParIntoIterLie<T> for Vec<T> {
type IntoIter = std::vec::IntoIter<T>;
fn into_par_iter(self) -> Self::IntoIter {
self.into_iter()
}
}
8 changes: 5 additions & 3 deletions dssim-core/src/tolab.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ use crate::image::ToRGB;
use crate::image::RGBAPLU;
use crate::image::RGBLU;
use imgref::*;
#[cfg(not(feature = "threads"))]
use crate::lieon as rayon;
use rayon::prelude::*;

const D65x: f64 = 0.9505;
Expand Down Expand Up @@ -69,7 +71,7 @@ impl ToLABBitmap for GBitmap {
unsafe { out.set_len(size) };

// For output width == stride
out.par_chunks_mut(width).enumerate().for_each(|(y, out_row)| {
out.as_mut_slice().par_chunks_mut(width).enumerate().for_each(|(y, out_row)| {
let start = y * self.stride();
let in_row = &self.buf()[start..start + width];
let out_row = &mut out_row[0..width];
Expand Down Expand Up @@ -102,8 +104,8 @@ fn rgb_to_lab<T: Copy + Sync + Send + 'static, F>(img: ImgRef<'_, T>, cb: F) ->
unsafe { out_b.set_len(size) };

// For output width == stride
out_l.par_chunks_mut(width).zip(
out_a.par_chunks_mut(width).zip(out_b.par_chunks_mut(width))
out_l.as_mut_slice().par_chunks_mut(width).zip(
out_a.as_mut_slice().par_chunks_mut(width).zip(out_b.as_mut_slice().par_chunks_mut(width))
).enumerate()
.for_each(|(y, (l_row, (a_row, b_row)))| {
let start = y * stride;
Expand Down
15 changes: 13 additions & 2 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
*/
#![allow(clippy::manual_range_contains)]
use getopts::Options;
#[cfg(feature = "threads")]
use rayon::prelude::*;
use std::env;
use std::path::PathBuf;
Expand Down Expand Up @@ -70,7 +71,12 @@ fn main() {

let mut attr = dssim::Dssim::new();

let files = files.par_iter().map(|file| -> Result<_, String> {
#[cfg(feature = "threads")]
let files_iter = files.par_iter();
#[cfg(not(feature = "threads"))]
let files_iter = files.iter();

let files = files_iter.map(|file| -> Result<_, String> {
let image = dssim::load_image(&attr, &file).map_err(|e| format!("Can't load {}, because: {}", file.display(), e))?;
Ok((file, image))
}).collect::<Result<Vec<_>,_>>();
Expand Down Expand Up @@ -102,7 +108,12 @@ fn main() {
println!("{:.8}\t{}", dssim, file2.display());

if map_output_file.is_some() {
ssim_maps.par_iter().enumerate().for_each(|(n, map_meta)| {
#[cfg(feature = "threads")]
let ssim_maps_iter = ssim_maps.par_iter();
#[cfg(not(feature = "threads"))]
let ssim_maps_iter = ssim_maps.iter();

ssim_maps_iter.enumerate().for_each(|(n, map_meta)| {
let avgssim = map_meta.ssim as f32;
let out: Vec<_> = map_meta.map.pixels().map(|ssim|{
let max = 1_f32 - ssim;
Expand Down

0 comments on commit e9d0a53

Please sign in to comment.