use crate::error::{Error, Result};
use crate::metrics::{
self, MetricConfig, MetricResult, PerceptionLevel, butteraugli, dssim, ssimulacra2,
};
use crate::viewing::ViewingCondition;
use imgref::ImgVec;
use rgb::{RGB8, RGBA};
fn rgb8_to_rgba_f32(img: &ImgVec<RGB8>) -> ImgVec<RGBA<f32>> {
let pixels: Vec<RGBA<f32>> = img
.pixels()
.map(|p| {
let r = srgb_to_linear(p.r);
let g = srgb_to_linear(p.g);
let b = srgb_to_linear(p.b);
RGBA::new(r, g, b, 1.0)
})
.collect();
ImgVec::new(pixels, img.width(), img.height())
}
#[inline]
fn srgb_to_linear(srgb: u8) -> f32 {
let s = f32::from(srgb) / 255.0;
if s <= 0.04045 {
s / 12.92
} else {
((s + 0.055) / 1.055).powf(2.4)
}
}
pub fn evaluate_single(
reference: &ImgVec<RGB8>,
encoded: &ImgVec<RGB8>,
config: &MetricConfig,
) -> Result<MetricResult> {
if reference.width() != encoded.width() || reference.height() != encoded.height() {
return Err(Error::DimensionMismatch {
expected: (reference.width(), reference.height()),
actual: (encoded.width(), encoded.height()),
});
}
let width = reference.width();
let height = reference.height();
let reference_img: ImgVec<RGB8>;
let reference_final = if config.xyb_roundtrip {
let ref_bytes: Vec<u8> = reference.pixels().flat_map(|p| [p.r, p.g, p.b]).collect();
let roundtripped = metrics::xyb_roundtrip(&ref_bytes, width, height);
let pixels: Vec<RGB8> = roundtripped
.chunks_exact(3)
.map(|chunk| RGB8::new(chunk[0], chunk[1], chunk[2]))
.collect();
reference_img = ImgVec::new(pixels, width, height);
&reference_img
} else {
reference
};
let mut result = MetricResult::default();
if config.dssim {
let ref_rgba = rgb8_to_rgba_f32(reference_final);
let enc_rgba = rgb8_to_rgba_f32(encoded);
let viewing = ViewingCondition::desktop();
result.dssim = Some(dssim::calculate_dssim(&ref_rgba, &enc_rgba, &viewing)?);
}
if config.ssimulacra2 || config.butteraugli || config.psnr {
let ref_buf: Vec<u8> = reference_final
.pixels()
.flat_map(|p| [p.r, p.g, p.b])
.collect();
let enc_buf: Vec<u8> = encoded.pixels().flat_map(|p| [p.r, p.g, p.b]).collect();
if config.ssimulacra2 {
result.ssimulacra2 = Some(ssimulacra2::calculate_ssimulacra2(
&ref_buf, &enc_buf, width, height,
)?);
}
if config.butteraugli {
result.butteraugli = Some(butteraugli::calculate_butteraugli(
&ref_buf, &enc_buf, width, height,
)?);
}
if config.psnr {
result.psnr = Some(metrics::calculate_psnr(&ref_buf, &enc_buf, width, height));
}
}
Ok(result)
}
pub fn assert_quality(
reference: &ImgVec<RGB8>,
encoded: &ImgVec<RGB8>,
min_ssimulacra2: Option<f64>,
max_dssim: Option<f64>,
) -> Result<()> {
let config = MetricConfig {
dssim: max_dssim.is_some(),
ssimulacra2: min_ssimulacra2.is_some(),
butteraugli: false,
psnr: false,
xyb_roundtrip: false,
};
let result = evaluate_single(reference, encoded, &config)?;
if let Some(threshold) = min_ssimulacra2 {
if let Some(score) = result.ssimulacra2 {
if score < threshold {
return Err(Error::QualityBelowThreshold {
metric: "SSIMULACRA2".to_string(),
value: score,
threshold,
});
}
}
}
if let Some(threshold) = max_dssim {
if let Some(score) = result.dssim {
if score > threshold {
return Err(Error::QualityBelowThreshold {
metric: "DSSIM".to_string(),
value: score,
threshold,
});
}
}
}
Ok(())
}
pub fn assert_perception_level(
reference: &ImgVec<RGB8>,
encoded: &ImgVec<RGB8>,
min_level: PerceptionLevel,
) -> Result<()> {
let config = MetricConfig {
dssim: true,
ssimulacra2: false,
butteraugli: false,
psnr: false,
xyb_roundtrip: false,
};
let result = evaluate_single(reference, encoded, &config)?;
if let Some(dssim) = result.dssim {
let actual_level = PerceptionLevel::from_dssim(dssim);
let min_level_value = min_level as u8;
let actual_level_value = actual_level as u8;
if actual_level_value > min_level_value {
return Err(Error::QualityBelowThreshold {
metric: format!("PerceptionLevel (DSSIM {dssim:.6})"),
value: actual_level_value.into(),
threshold: min_level_value.into(),
});
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_image(width: usize, height: usize, pattern: u8) -> ImgVec<RGB8> {
let pixels: Vec<RGB8> = (0..width * height)
.map(|i| {
let base = (i + usize::from(pattern)) % 256;
RGB8::new(base as u8, (base + 50) as u8, (base + 100) as u8)
})
.collect();
ImgVec::new(pixels, width, height)
}
#[test]
fn test_evaluate_single_identical() {
let img = create_test_image(64, 64, 0);
let config = MetricConfig::perceptual();
let result = evaluate_single(&img, &img, &config).unwrap();
assert!(result.dssim.unwrap() < 0.0001);
assert!(result.ssimulacra2.unwrap() > 99.0);
assert!(result.butteraugli.unwrap() < 0.1);
}
#[test]
fn test_evaluate_single_dimension_mismatch() {
let img1 = create_test_image(64, 64, 0);
let img2 = create_test_image(32, 32, 0);
let config = MetricConfig::perceptual();
let result = evaluate_single(&img1, &img2, &config);
assert!(result.is_err());
}
#[test]
fn test_assert_quality_pass() {
let img = create_test_image(64, 64, 0);
assert!(assert_quality(&img, &img, Some(90.0), Some(0.001)).is_ok());
}
#[test]
fn test_assert_quality_fail_ssimulacra2() {
let img1 = create_test_image(64, 64, 0);
let img2 = create_test_image(64, 64, 50);
assert!(assert_quality(&img1, &img2, Some(99.0), None).is_err());
}
#[test]
fn test_assert_perception_level() {
let img = create_test_image(64, 64, 0);
assert!(assert_perception_level(&img, &img, PerceptionLevel::Imperceptible).is_ok());
}
}