diff --git a/Cargo.toml b/Cargo.toml index 058de3f..6962336 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ name = "ssimulacra2" version = "0.4.0" edition = "2021" -rust-version = "1.61.0" +rust-version = "1.65.0" description = "Rust implementation of the SSIMULACRA2 metric" repository = "https://github.com/rust-av/ssimulacra2" homepage = "https://github.com/rust-av/ssimulacra2" @@ -15,10 +15,10 @@ default = ["rayon"] [dependencies] aligned = "0.4.1" -anyhow = "1.0.0" nalgebra = "0.32.2" num-traits = "0.2.15" rayon = { version = "1.5.3", optional = true } +thiserror = "1.0.56" wide = "0.7.5" yuvxyb = "0.3.0" diff --git a/src/blur.rs b/src/blur.rs index 5274516..bd327f4 100644 --- a/src/blur.rs +++ b/src/blur.rs @@ -3,6 +3,14 @@ use std::f64::consts::PI; use aligned::{Aligned, A16}; use nalgebra::base::{Matrix3, Matrix3x1}; +/// Structure handling image blur. +/// +/// This struct contains the necessary buffers and the kernel used for blurring +/// (currently a recursive approximation of the Gaussian filter). +/// +/// Note that the width and height of the image passed to [blur][Self::blur] needs to exactly +/// match the width and height of this instance. If you reduce the image size (e.g. via +/// downscaling), [`shrink_to`][Self::shrink_to] can be used to resize the internal buffers. pub struct Blur { kernel: RecursiveGaussian, temp: Vec, @@ -11,6 +19,8 @@ pub struct Blur { } impl Blur { + /// Create a new [Blur] for images of the given width and height. + /// This pre-allocates the necessary buffers. #[must_use] pub fn new(width: usize, height: usize) -> Self { Blur { @@ -21,12 +31,17 @@ impl Blur { } } + /// Truncates the internal buffers to fit images of the given width and height. + /// + /// This will [truncate][Vec::truncate] the internal buffers + /// without affecting the allocated memory. pub fn shrink_to(&mut self, width: usize, height: usize) { self.temp.truncate(width * height); self.width = width; self.height = height; } + /// Blur the given image. pub fn blur(&mut self, img: &[Vec; 3]) -> [Vec; 3] { [ self.blur_plane(&img[0]), diff --git a/src/lib.rs b/src/lib.rs index db11c1f..d290c3f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -41,13 +41,33 @@ mod blur; -use anyhow::{bail, Result}; pub use blur::Blur; pub use yuvxyb::{CastFromPrimitive, Frame, LinearRgb, Pixel, Plane, Rgb, Xyb, Yuv}; pub use yuvxyb::{ColorPrimaries, MatrixCoefficients, TransferCharacteristic, YuvConfig}; +// How often to downscale and score the input images. +// Each scaling step will downscale by a factor of two. const NUM_SCALES: usize = 6; +/// Errors which can occur when attempting to calculate a SSIMULACRA2 score from two input images. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum Ssimulacra2Error { + /// The conversion from input image to [LinearRgb] (via [TryFrom]) returned an [Err]. + /// Note that the conversion from LinearRgb to [Xyb] cannot fail, which means that + /// this is the only point of failure regarding image conversion. + #[error("Failed to convert input image to linear RGB")] + LinearRgbConversionFailed, + + /// The two input images do not have the same width and height. + #[error("Source and distorted image width and height must be equal")] + NonMatchingImageDimensions, + + /// One of the input images has a width and/or height of less than 8 pixels. + /// This is not currently supported by the SSIMULACRA2 metric. + #[error("Images must be at least 8x8 pixels")] + InvalidImageSize, +} + /// Computes the SSIMULACRA2 score for a given input frame and the distorted /// version of that frame. /// @@ -55,23 +75,24 @@ const NUM_SCALES: usize = 6; /// - If the source and distorted image width and height do not match /// - If the source or distorted image cannot be converted to XYB successfully /// - If the image is smaller than 8x8 pixels -pub fn compute_frame_ssimulacra2, U: TryInto>( - source: T, - distorted: U, -) -> Result { - let mut img1: LinearRgb = source - .try_into() - .map_err(|_e| anyhow::anyhow!("Failed to convert to Linear Rgb"))?; - let mut img2: LinearRgb = distorted - .try_into() - .map_err(|_e| anyhow::anyhow!("Failed to convert to Linear Rgb"))?; +pub fn compute_frame_ssimulacra2(source: T, distorted: U) -> Result +where + LinearRgb: TryFrom + TryFrom, +{ + 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() { - bail!("Source and distorted image width and height must be equal"); + return Err(Ssimulacra2Error::NonMatchingImageDimensions); } if img1.width() < 8 || img1.height() < 8 { - bail!("Images must be at least 8x8 pixels"); + return Err(Ssimulacra2Error::InvalidImageSize); } let mut width = img1.width();