#![forbid(unsafe_code)]
mod blur;
mod input;
mod precompute;
#[doc(hidden)]
pub mod reference_data;
#[allow(clippy::too_many_arguments)] mod simd_ops;
mod xyb_simd;
pub use blur::Blur;
pub use input::{LinearRgbImage, ToLinearRgb};
pub use precompute::Ssimulacra2Reference;
pub use input::{srgb_to_linear, srgb_u8_to_linear, srgb_u16_to_linear};
use yuvxyb::LinearRgb;
use yuvxyb::Xyb;
pub(crate) const NUM_SCALES: usize = 6;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SimdImpl {
Scalar,
#[default]
Simd,
}
impl SimdImpl {
pub fn name(&self) -> &'static str {
match self {
SimdImpl::Scalar => "scalar",
SimdImpl::Simd => "simd (archmage)",
}
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Ssimulacra2Config {
pub impl_type: SimdImpl,
}
impl Ssimulacra2Config {
pub fn new(impl_type: SimdImpl) -> Self {
Self { impl_type }
}
pub fn simd() -> Self {
Self::new(SimdImpl::Simd)
}
pub fn scalar() -> Self {
Self::new(SimdImpl::Scalar)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Ssimulacra2Error {
#[error("Failed to convert input image to linear RGB")]
LinearRgbConversionFailed,
#[error("Source and distorted image width and height must be equal")]
NonMatchingImageDimensions,
#[error("Images must be at least 8x8 pixels")]
InvalidImageSize,
#[error("Gaussian blur operation failed")]
GaussianBlurError,
}
#[deprecated(
since = "0.8.0",
note = "use compute_ssimulacra2 with ToLinearRgb types instead"
)]
pub fn compute_frame_ssimulacra2<T, U>(source: T, distorted: U) -> Result<f64, Ssimulacra2Error>
where
LinearRgb: TryFrom<T> + TryFrom<U>,
{
compute_frame_ssimulacra2_impl(source, distorted, Ssimulacra2Config::default())
}
#[deprecated(
since = "0.8.0",
note = "use compute_ssimulacra2_with_config with ToLinearRgb types instead"
)]
pub fn compute_frame_ssimulacra2_with_config<T, U>(
source: T,
distorted: U,
config: Ssimulacra2Config,
) -> Result<f64, Ssimulacra2Error>
where
LinearRgb: TryFrom<T> + TryFrom<U>,
{
compute_frame_ssimulacra2_impl(source, distorted, config)
}
pub fn compute_ssimulacra2<S, D>(source: S, distorted: D) -> Result<f64, Ssimulacra2Error>
where
S: ToLinearRgb,
D: ToLinearRgb,
{
compute_ssimulacra2_with_config(source, distorted, Ssimulacra2Config::default())
}
pub fn compute_ssimulacra2_with_config<S, D>(
source: S,
distorted: D,
config: Ssimulacra2Config,
) -> Result<f64, Ssimulacra2Error>
where
S: ToLinearRgb,
D: ToLinearRgb,
{
let img1: LinearRgb = source.into_linear_rgb().into();
let img2: LinearRgb = distorted.into_linear_rgb().into();
compute_frame_ssimulacra2_impl(img1, img2, config)
}
fn compute_frame_ssimulacra2_impl<T, U>(
source: T,
distorted: U,
config: Ssimulacra2Config,
) -> Result<f64, Ssimulacra2Error>
where
LinearRgb: TryFrom<T> + TryFrom<U>,
{
let Ok(mut img1) = LinearRgb::try_from(source) else {
return Err(Ssimulacra2Error::LinearRgbConversionFailed);
};
let Ok(mut img2) = LinearRgb::try_from(distorted) else {
return Err(Ssimulacra2Error::LinearRgbConversionFailed);
};
if img1.width() != img2.width() || img1.height() != img2.height() {
return Err(Ssimulacra2Error::NonMatchingImageDimensions);
}
if img1.width().get() < 8 || img1.height().get() < 8 {
return Err(Ssimulacra2Error::InvalidImageSize);
}
let mut width = img1.width().get();
let mut height = img1.height().get();
let impl_type = config.impl_type;
let alloc_plane = || vec![0.0f32; width * height];
let alloc_3planes = || [alloc_plane(), alloc_plane(), alloc_plane()];
let mut mul = alloc_3planes();
let mut sigma1_sq = alloc_3planes();
let mut sigma2_sq = alloc_3planes();
let mut sigma12 = alloc_3planes();
let mut mu1 = alloc_3planes();
let mut mu2 = alloc_3planes();
let mut img1_planar = alloc_3planes();
let mut img2_planar = alloc_3planes();
let mut blur = Blur::with_simd_impl(width, height, impl_type);
let mut msssim = Msssim::default();
for scale in 0..NUM_SCALES {
if width < 8 || height < 8 {
break;
}
if scale > 0 {
img1 = downscale_by_2(&img1);
img2 = downscale_by_2(&img2);
width = img1.width().get();
height = img2.height().get();
}
let size = width * height;
for buf in [
&mut mul,
&mut sigma1_sq,
&mut sigma2_sq,
&mut sigma12,
&mut mu1,
&mut mu2,
&mut img1_planar,
&mut img2_planar,
] {
for c in buf.iter_mut() {
c.truncate(size);
}
}
blur.shrink_to(width, height);
let mut img1_xyb = linear_rgb_to_xyb(img1.clone(), impl_type);
let mut img2_xyb = linear_rgb_to_xyb(img2.clone(), impl_type);
make_positive_xyb(&mut img1_xyb);
make_positive_xyb(&mut img2_xyb);
xyb_to_planar_into(&img1_xyb, &mut img1_planar);
xyb_to_planar_into(&img2_xyb, &mut img2_planar);
image_multiply(&img1_planar, &img1_planar, &mut mul, impl_type);
blur.blur_into(&mul, &mut sigma1_sq);
image_multiply(&img2_planar, &img2_planar, &mut mul, impl_type);
blur.blur_into(&mul, &mut sigma2_sq);
image_multiply(&img1_planar, &img2_planar, &mut mul, impl_type);
blur.blur_into(&mul, &mut sigma12);
blur.blur_into(&img1_planar, &mut mu1);
blur.blur_into(&img2_planar, &mut mu2);
let avg_ssim = ssim_map(
width, height, &mu1, &mu2, &sigma1_sq, &sigma2_sq, &sigma12, impl_type,
);
let avg_edgediff = edge_diff_map(
width,
height,
&img1_planar,
&mu1,
&img2_planar,
&mu2,
impl_type,
);
msssim.scales.push(MsssimScale {
avg_ssim,
avg_edgediff,
});
}
Ok(msssim.score())
}
fn linear_rgb_to_xyb(linear_rgb: LinearRgb, impl_type: SimdImpl) -> Xyb {
match impl_type {
SimdImpl::Scalar => Xyb::from(linear_rgb),
SimdImpl::Simd => {
let width = linear_rgb.width(); let height = linear_rgb.height(); let mut data = linear_rgb.into_data();
xyb_simd::linear_rgb_to_xyb_simd(&mut data);
Xyb::new(data, width, height).expect("XYB construction should not fail")
}
}
}
pub(crate) fn linear_rgb_to_xyb_simd(linear_rgb: LinearRgb) -> Xyb {
linear_rgb_to_xyb(linear_rgb, SimdImpl::Simd)
}
pub(crate) fn make_positive_xyb(xyb: &mut Xyb) {
for pix in xyb.data_mut().iter_mut() {
pix[2] = (pix[2] - pix[1]) + 0.55;
pix[0] = (pix[0]).mul_add(14.0, 0.42);
pix[1] += 0.01;
}
}
pub(crate) fn xyb_to_planar(xyb: &Xyb) -> [Vec<f32>; 3] {
let size = xyb.width().get() * xyb.height().get();
let mut out = [vec![0.0f32; size], vec![0.0f32; size], vec![0.0f32; size]];
xyb_to_planar_into(xyb, &mut out);
out
}
pub(crate) fn xyb_to_planar_into(xyb: &Xyb, out: &mut [Vec<f32>; 3]) {
let [out0, out1, out2] = out;
for (((i, o0), o1), o2) in xyb
.data()
.iter()
.copied()
.zip(out0.iter_mut())
.zip(out1.iter_mut())
.zip(out2.iter_mut())
{
*o0 = i[0];
*o1 = i[1];
*o2 = i[2];
}
}
pub(crate) fn image_multiply(
img1: &[Vec<f32>; 3],
img2: &[Vec<f32>; 3],
out: &mut [Vec<f32>; 3],
impl_type: SimdImpl,
) {
match impl_type {
SimdImpl::Scalar => image_multiply_scalar(img1, img2, out),
SimdImpl::Simd => simd_ops::image_multiply_simd(img1, img2, out),
}
}
fn image_multiply_scalar(img1: &[Vec<f32>; 3], img2: &[Vec<f32>; 3], out: &mut [Vec<f32>; 3]) {
for ((plane1, plane2), out_plane) in img1.iter().zip(img2.iter()).zip(out.iter_mut()) {
for ((&p1, &p2), o) in plane1.iter().zip(plane2.iter()).zip(out_plane.iter_mut()) {
*o = p1 * p2;
}
}
}
pub(crate) fn downscale_by_2(in_data: &LinearRgb) -> LinearRgb {
use std::num::NonZeroUsize;
const SCALE: usize = 2;
let in_w = in_data.width().get();
let in_h = in_data.height().get();
let out_w = in_w.div_ceil(SCALE);
let out_h = in_h.div_ceil(SCALE);
let mut out_data = vec![[0.0f32; 3]; out_w * out_h];
let normalize = 1.0f32 / (SCALE * SCALE) as f32;
let in_data = &in_data.data();
for oy in 0..out_h {
for ox in 0..out_w {
for c in 0..3 {
let mut sum = 0f32;
for iy in 0..SCALE {
for ix in 0..SCALE {
let x = (ox * SCALE + ix).min(in_w - 1);
let y = (oy * SCALE + iy).min(in_h - 1);
sum += in_data[y * in_w + x][c];
}
}
out_data[oy * out_w + ox][c] = sum * normalize;
}
}
}
LinearRgb::new(
out_data,
NonZeroUsize::new(out_w).expect("out_w must be nonzero"),
NonZeroUsize::new(out_h).expect("out_h must be nonzero"),
)
.expect("Resolution and data size match")
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn ssim_map(
width: usize,
height: usize,
m1: &[Vec<f32>; 3],
m2: &[Vec<f32>; 3],
s11: &[Vec<f32>; 3],
s22: &[Vec<f32>; 3],
s12: &[Vec<f32>; 3],
impl_type: SimdImpl,
) -> [f64; 3 * 2] {
match impl_type {
SimdImpl::Scalar => ssim_map_scalar(width, height, m1, m2, s11, s22, s12),
SimdImpl::Simd => simd_ops::ssim_map_simd(width, height, m1, m2, s11, s22, s12),
}
}
fn ssim_map_scalar(
width: usize,
height: usize,
m1: &[Vec<f32>; 3],
m2: &[Vec<f32>; 3],
s11: &[Vec<f32>; 3],
s22: &[Vec<f32>; 3],
s12: &[Vec<f32>; 3],
) -> [f64; 3 * 2] {
const C2: f32 = 0.0009f32;
let one_per_pixels = 1.0f64 / (width * height) as f64;
let mut plane_averages = [0f64; 3 * 2];
for c in 0..3 {
let mut sum_d = 0.0f64;
let mut sum_d4 = 0.0f64;
for (row_m1, (row_m2, (row_s11, (row_s22, row_s12)))) in m1[c].chunks_exact(width).zip(
m2[c].chunks_exact(width).zip(
s11[c]
.chunks_exact(width)
.zip(s22[c].chunks_exact(width).zip(s12[c].chunks_exact(width))),
),
) {
for x in 0..width {
let mu1 = row_m1[x];
let mu2 = row_m2[x];
let mu11 = mu1 * mu1;
let mu22 = mu2 * mu2;
let mu12 = mu1 * mu2;
let mu_diff = mu1 - mu2;
let num_m = mu_diff.mul_add(-mu_diff, 1.0f32);
let num_s = 2.0f32.mul_add(row_s12[x] - mu12, C2);
let denom_s = (row_s11[x] - mu11) + (row_s22[x] - mu22) + C2;
let d = (1.0f32 - (num_m * num_s) / denom_s).max(0.0f32);
let d2 = d * d;
let d4 = d2 * d2;
sum_d += f64::from(d);
sum_d4 += f64::from(d4);
}
}
plane_averages[c * 2] = one_per_pixels * sum_d;
plane_averages[c * 2 + 1] = (one_per_pixels * sum_d4).sqrt().sqrt();
}
plane_averages
}
pub(crate) fn edge_diff_map(
width: usize,
height: usize,
img1: &[Vec<f32>; 3],
mu1: &[Vec<f32>; 3],
img2: &[Vec<f32>; 3],
mu2: &[Vec<f32>; 3],
impl_type: SimdImpl,
) -> [f64; 3 * 4] {
match impl_type {
SimdImpl::Scalar => edge_diff_map_scalar(width, height, img1, mu1, img2, mu2),
SimdImpl::Simd => simd_ops::edge_diff_map_simd(width, height, img1, mu1, img2, mu2),
}
}
fn edge_diff_map_scalar(
width: usize,
height: usize,
img1: &[Vec<f32>; 3],
mu1: &[Vec<f32>; 3],
img2: &[Vec<f32>; 3],
mu2: &[Vec<f32>; 3],
) -> [f64; 3 * 4] {
let one_per_pixels = 1.0f64 / (width * height) as f64;
let mut plane_averages = [0f64; 3 * 4];
for c in 0..3 {
let mut sum1 = [0.0f64; 4];
for (row1, (row2, (rowm1, rowm2))) in img1[c].chunks_exact(width).zip(
img2[c]
.chunks_exact(width)
.zip(mu1[c].chunks_exact(width).zip(mu2[c].chunks_exact(width))),
) {
for x in 0..width {
let d1: f64 = (1.0 + f64::from((row2[x] - rowm2[x]).abs()))
/ (1.0 + f64::from((row1[x] - rowm1[x]).abs()))
- 1.0;
let artifact = d1.max(0.0);
sum1[0] += artifact;
sum1[1] += artifact.powi(4);
let detail_lost = (-d1).max(0.0);
sum1[2] += detail_lost;
sum1[3] += detail_lost.powi(4);
}
}
plane_averages[c * 4] = one_per_pixels * sum1[0];
plane_averages[c * 4 + 1] = (one_per_pixels * sum1[1]).sqrt().sqrt();
plane_averages[c * 4 + 2] = one_per_pixels * sum1[2];
plane_averages[c * 4 + 3] = (one_per_pixels * sum1[3]).sqrt().sqrt();
}
plane_averages
}
#[derive(Debug, Clone, Default)]
pub(crate) struct Msssim {
pub scales: Vec<MsssimScale>,
}
#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct MsssimScale {
pub avg_ssim: [f64; 3 * 2],
pub avg_edgediff: [f64; 3 * 4],
}
impl Msssim {
#[allow(clippy::too_many_lines)]
pub fn score(&self) -> f64 {
const WEIGHT: [f64; 108] = [
0.0,
0.000_737_660_670_740_658_6,
0.0,
0.0,
0.000_779_348_168_286_730_9,
0.0,
0.0,
0.000_437_115_573_010_737_9,
0.0,
1.104_172_642_665_734_6,
0.000_662_848_341_292_71,
0.000_152_316_327_837_187_52,
0.0,
0.001_640_643_745_659_975_4,
0.0,
1.842_245_552_053_929_8,
11.441_172_603_757_666,
0.0,
0.000_798_910_943_601_516_3,
0.000_176_816_438_078_653,
0.0,
1.878_759_497_954_638_7,
10.949_069_906_051_42,
0.0,
0.000_728_934_699_150_807_2,
0.967_793_708_062_683_3,
0.0,
0.000_140_034_242_854_358_84,
0.998_176_697_785_496_7,
0.000_319_497_559_344_350_53,
0.000_455_099_211_379_206_3,
0.0,
0.0,
0.001_364_876_616_324_339_8,
0.0,
0.0,
0.0,
0.0,
0.0,
7.466_890_328_078_848,
0.0,
17.445_833_984_131_262,
0.000_623_560_163_404_146_6,
0.0,
0.0,
6.683_678_146_179_332,
0.000_377_244_079_796_112_96,
1.027_889_937_768_264,
225.205_153_008_492_74,
0.0,
0.0,
19.213_238_186_143_016,
0.001_140_152_458_661_836_1,
0.001_237_755_635_509_985,
176.393_175_984_506_94,
0.0,
0.0,
24.433_009_998_704_76,
0.285_208_026_121_177_57,
0.000_448_543_692_383_340_8,
0.0,
0.0,
0.0,
34.779_063_444_837_72,
44.835_625_328_877_896,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.000_868_055_657_329_169_8,
0.0,
0.0,
0.0,
0.0,
0.0,
0.000_531_319_187_435_874_7,
0.0,
0.000_165_338_141_613_791_12,
0.0,
0.0,
0.0,
0.0,
0.0,
0.000_417_917_180_325_133_6,
0.001_729_082_823_472_283_3,
0.0,
0.002_082_700_584_663_643_7,
0.0,
0.0,
8.826_982_764_996_862,
23.192_433_439_989_26,
0.0,
95.108_049_881_108_6,
0.986_397_803_440_068_2,
0.983_438_279_246_535_3,
0.001_228_640_504_827_849_3,
171.266_725_589_730_7,
0.980_785_887_243_537_9,
0.0,
0.0,
0.0,
0.000_513_006_458_899_067_9,
0.0,
0.000_108_540_578_584_115_37,
];
let mut ssim = 0.0f64;
let mut i = 0usize;
for c in 0..3 {
for scale in &self.scales {
for n in 0..2 {
ssim = WEIGHT[i].mul_add(scale.avg_ssim[c * 2 + n].abs(), ssim);
i += 1;
ssim = WEIGHT[i].mul_add(scale.avg_edgediff[c * 4 + n].abs(), ssim);
i += 1;
ssim = WEIGHT[i].mul_add(scale.avg_edgediff[c * 4 + n + 2].abs(), ssim);
i += 1;
}
}
}
ssim *= 0.956_238_261_683_484_4_f64;
ssim = (6.248_496_625_763_138e-5 * ssim * ssim).mul_add(
ssim,
2.326_765_642_916_932f64.mul_add(ssim, -0.020_884_521_182_843_837 * ssim * ssim),
);
if ssim > 0.0f64 {
ssim = ssim
.powf(0.627_633_646_783_138_7)
.mul_add(-10.0f64, 100.0f64);
} else {
ssim = 100.0f64;
}
ssim
}
}
#[cfg(test)]
#[allow(deprecated)]
mod tests {
use std::path::PathBuf;
use super::*;
use yuvxyb::{ColorPrimaries, Rgb, TransferCharacteristic};
#[test]
fn test_ssimulacra2() {
let source = image::open(
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("test_data")
.join("tank_source.png"),
)
.unwrap();
let distorted = image::open(
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("test_data")
.join("tank_distorted.png"),
)
.unwrap();
let source_data = source
.to_rgb32f()
.chunks_exact(3)
.map(|chunk| [chunk[0], chunk[1], chunk[2]])
.collect::<Vec<_>>();
let source_data = Xyb::try_from(
Rgb::new(
source_data,
std::num::NonZeroUsize::new(source.width() as usize).unwrap(),
std::num::NonZeroUsize::new(source.height() as usize).unwrap(),
TransferCharacteristic::SRGB,
ColorPrimaries::BT709,
)
.unwrap(),
)
.unwrap();
let distorted_data = distorted
.to_rgb32f()
.chunks_exact(3)
.map(|chunk| [chunk[0], chunk[1], chunk[2]])
.collect::<Vec<_>>();
let distorted_data = Xyb::try_from(
Rgb::new(
distorted_data,
std::num::NonZeroUsize::new(distorted.width() as usize).unwrap(),
std::num::NonZeroUsize::new(distorted.height() as usize).unwrap(),
TransferCharacteristic::SRGB,
ColorPrimaries::BT709,
)
.unwrap(),
)
.unwrap();
let result = compute_frame_ssimulacra2(source_data, distorted_data).unwrap();
let expected = 17.398_505_f64;
assert!(
(result - expected).abs() < 0.25f64,
"Result {result:.6} not equal to expected {expected:.6}",
);
}
#[test]
fn test_xyb_simd_vs_yuvxyb() {
use yuvxyb::{ColorPrimaries, TransferCharacteristic};
let source = image::open(
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("test_data")
.join("tank_source.png"),
)
.unwrap();
let source_data: Vec<[f32; 3]> = source
.to_rgb32f()
.chunks_exact(3)
.map(|chunk| [chunk[0], chunk[1], chunk[2]])
.collect();
let width = source.width() as usize;
let height = source.height() as usize;
let nz_width = std::num::NonZeroUsize::new(width).unwrap();
let nz_height = std::num::NonZeroUsize::new(height).unwrap();
let rgb_for_yuvxyb = Rgb::new(
source_data.clone(),
nz_width,
nz_height,
TransferCharacteristic::SRGB,
ColorPrimaries::BT709,
)
.unwrap();
let lrgb_for_yuvxyb = yuvxyb::LinearRgb::try_from(rgb_for_yuvxyb).unwrap();
let xyb_yuvxyb = yuvxyb::Xyb::from(lrgb_for_yuvxyb);
let rgb_for_simd = Rgb::new(
source_data,
nz_width,
nz_height,
TransferCharacteristic::SRGB,
ColorPrimaries::BT709,
)
.unwrap();
let lrgb_for_simd = LinearRgb::try_from(rgb_for_simd).unwrap();
let xyb_simd = linear_rgb_to_xyb_simd(lrgb_for_simd);
let mut max_diff = [0.0f32; 3];
for (yuvxyb_pix, simd_pix) in xyb_yuvxyb.data().iter().zip(xyb_simd.data().iter()) {
for c in 0..3 {
let diff = (yuvxyb_pix[c] - simd_pix[c]).abs();
max_diff[c] = max_diff[c].max(diff);
}
}
assert!(
max_diff[0] < 1e-5 && max_diff[1] < 1e-5 && max_diff[2] < 1e-5,
"SIMD XYB differs from yuvxyb: max_diff={:?}",
max_diff
);
}
}