Added function mse

src/lib.rs

@@ -1,13 +1,37 @@
 extern crate image;
 
-use image::{Pixel, GenericImage, ImageResult};
+use image::{Pixel, GenericImage, ImageResult, Rgba};
 
-/// Compute a PSNR between two images.
+/// Compute a MSE between an image and its color.
 ///
 /// Will panic an error if the two images don't have the same size.
-pub fn psnr<P: Pixel<Subpixel=u8>, T: GenericImage<Pixel=P>>(img1: &T, img2: &T) -> f64 {
+pub fn mse_color<P: Pixel<Subpixel=u8>, T: GenericImage<Pixel=P>>(img: &T, color: &Rgba<u8>) -> f64 {
 
-    let twenty_log10_max: f64 = 20.0 * (std::u8::MAX as f64).log10();
+    let mut mse = 0.0;
+
+    for x in 0 .. img.dimensions().0 {
+        for y in 0 .. img.dimensions().1 {
+
+            let p1 = img.get_pixel(x, y).to_rgba();
+
+            let dx = p1.data[0] as f64 - color.data[0] as f64;
+            let dy = p1.data[1] as f64 - color.data[1] as f64;
+            let dz = p1.data[2] as f64 - color.data[2] as f64;
+
+            let dx = dx * dx;
+            let dy = dy * dy;
+            let dz = dz * dz;
+
+            mse += dx + dy + dz;
+        }
+    }
+
+    mse / (img.dimensions().0 * img.dimensions().1 * 3) as f64
+}
+/// Compute a MSE between two images.
+///
+/// Will panic an error if the two images don't have the same size.
+pub fn mse<P: Pixel<Subpixel=u8>, T: GenericImage<Pixel=P>>(img1: &T, img2: &T) -> f64 {
 
     assert_eq!(img1.dimensions(), img2.dimensions());
 
@@ -31,8 +55,12 @@ pub fn psnr<P: Pixel<Subpixel=u8>, T: GenericImage<Pixel=P>>(img1: &T, img2: &T)
         }
     }
 
-    mse /= (img1.dimensions().0 * img1.dimensions().1 * 3) as f64;
+    mse / (img1.dimensions().0 * img1.dimensions().1 * 3) as f64
-    twenty_log10_max - 10.0 * mse.log10()
+}
+
+pub fn psnr<P: Pixel<Subpixel=u8>, T: GenericImage<Pixel=P>>(img1: &T, img2: &T) -> f64 {
+    let twenty_log10_max: f64 = 20.0 * (std::u8::MAX as f64).log10();
+    twenty_log10_max - 10.0 * mse(img1, img2).log10()
 }
 
 pub fn psnr_files(path1: &str, path2: &str) -> ImageResult<f64> {

src/main.rs

@@ -11,24 +11,32 @@ fn psnr_dir(path1: &str, path2: &str) -> Vec<(String, ImageResult<f64>)> {
     let mut dirs1: Vec<_> = fs::read_dir(path1).unwrap()
         .map(|x|x.unwrap())
         .filter(|x| x.file_name().into_string().unwrap().ends_with(".png"))
+        .map(|x| x.file_name().to_owned().into_string().unwrap())
         .collect();
 
-    let mut dirs2: Vec<_> = fs::read_dir(path2).unwrap()
+    let metadata = fs::metadata(path2).unwrap();
+
+    let mut dirs2: Vec<_> = if metadata.is_dir() {
+        fs::read_dir(path2).unwrap()
         .map(|x|x.unwrap())
         .filter(|x| x.file_name().into_string().unwrap().ends_with(".png"))
-        .collect();
+        .map(|x| x.file_name().to_owned().into_string().unwrap())
+        .collect()
+    } else {
+        vec![path2.to_owned()]
+    };
 
-    dirs1.sort_by_key(|x| x.file_name());
+    dirs1.sort();
-    dirs2.sort_by_key(|x| x.file_name());
+    dirs2.sort();
 
     let mut res = vec!();
 
     for file in dirs1 {
-        let p1 = path1.to_owned() + &file.file_name().to_owned().into_string().unwrap();
+        let p1 = path1.to_owned() + &file;
-        let p2 = path2.to_owned() + &file.file_name().to_owned().into_string().unwrap();
+        let p2 = if dirs2.len() != 1 { path2.to_owned() + &file } else { path2.to_owned() };
         let psnr = psnr_files(&p1, &p2);
 
-        let elt = (file.file_name().to_owned().into_string().unwrap(), psnr.unwrap());
+        let elt = (file, psnr.unwrap());
         println!("{} {}", elt.0, elt.1);
         res.push((elt.0, Ok(elt.1)));
     }