diff --git a/.gitignore b/.gitignore index ea8c4bf..1972091 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,15 @@ -/target +# Generated by Cargo +# will have compiled files and executables +debug/ +target/ + +# These are backup files generated by rustfmt +**/*.rs.bk + +# MSVC Windows builds of rustc generate these, which store debugging information +*.pdb + +# Manually added +/.devcontainer/ +/.vscode/ +/test_files/ \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index bb75cfb..4faf25d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -265,18 +265,95 @@ dependencies = [ "spin", ] +[[package]] +name = "futures" +version = "0.3.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38390104763dc37a5145a53c29c63c1290b5d316d6086ec32c293f6736051bb0" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ba265a92256105f45b719605a571ffe2d1f0fea3807304b522c1d778f79eed" +dependencies = [ + "futures-core", + "futures-sink", +] + [[package]] name = "futures-core" version = "0.3.25" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04909a7a7e4633ae6c4a9ab280aeb86da1236243a77b694a49eacd659a4bd3ac" +[[package]] +name = "futures-executor" +version = "0.3.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7acc85df6714c176ab5edf386123fafe217be88c0840ec11f199441134a074e2" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00f5fb52a06bdcadeb54e8d3671f8888a39697dcb0b81b23b55174030427f4eb" + +[[package]] +name = "futures-macro" +version = "0.3.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdfb8ce053d86b91919aad980c220b1fb8401a9394410e1c289ed7e66b61835d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.25" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39c15cf1a4aa79df40f1bb462fb39676d0ad9e366c2a33b590d7c66f4f81fcf9" +[[package]] +name = "futures-task" +version = "0.3.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ffb393ac5d9a6eaa9d3fdf37ae2776656b706e200c8e16b1bdb227f5198e6ea" + +[[package]] +name = "futures-util" +version = "0.3.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "197676987abd2f9cadff84926f410af1c183608d36641465df73ae8211dc65d6" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "pin-utils", + "slab", +] + [[package]] name = "getrandom" version = "0.2.8" @@ -361,6 +438,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" + [[package]] name = "lebe" version = "0.5.2" @@ -401,6 +484,12 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "memchr" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" + [[package]] name = "memoffset" version = "0.6.5" @@ -428,6 +517,18 @@ dependencies = [ "adler", ] +[[package]] +name = "mio" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57ee1c23c7c63b0c9250c339ffdc69255f110b298b901b9f6c82547b7b87caaf" +dependencies = [ + "libc", + "log", + "wasi", + "windows-sys 0.36.1", +] + [[package]] name = "nalgebra" version = "0.31.2" @@ -560,6 +661,29 @@ version = "6.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ff7415e9ae3fff1225851df9e0d9e4e5479f947619774677a63572e55e80eff" +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dc9e0dc2adc1c69d09143aff38d3d30c5c3f0df0dad82e6d25547af174ebec0" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-sys 0.42.0", +] + [[package]] name = "paste" version = "1.0.9" @@ -586,6 +710,18 @@ dependencies = [ "syn", ] +[[package]] +name = "pin-project-lite" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0a7ae3ac2f1173085d398531c705756c94a4c56843785df85a60c1a0afac116" + +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + [[package]] name = "png" version = "0.17.6" @@ -631,6 +767,15 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "progress_bar" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c8e8843722ce55deb442e7d6476a97c081b8485ea2807b031d160f592f99037" +dependencies = [ + "lazy_static", +] + [[package]] name = "quote" version = "1.0.21" @@ -670,6 +815,15 @@ dependencies = [ "num_cpus", ] +[[package]] +name = "redox_syscall" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" +dependencies = [ + "bitflags", +] + [[package]] name = "rust_hawktracer" version = "0.7.0" @@ -713,6 +867,15 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" +[[package]] +name = "signal-hook-registry" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e51e73328dc4ac0c7ccbda3a494dfa03df1de2f46018127f60c693f2648455b0" +dependencies = [ + "libc", +] + [[package]] name = "simba" version = "0.7.2" @@ -726,12 +889,31 @@ dependencies = [ "wide", ] +[[package]] +name = "slab" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4614a76b2a8be0058caa9dbbaf66d988527d86d003c11a94fbd335d7661edcef" +dependencies = [ + "autocfg", +] + [[package]] name = "smallvec" version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0" +[[package]] +name = "socket2" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02e2d2db9033d13a1567121ddd7a095ee144db4e1ca1b1bda3419bc0da294ebd" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "spin" version = "0.9.4" @@ -760,8 +942,11 @@ name = "ssimulacra2_rs" version = "0.1.0" dependencies = [ "clap", + "futures", "image", + "progress_bar", "ssimulacra2", + "tokio", "yuvxyb", ] @@ -837,6 +1022,37 @@ dependencies = [ "weezl", ] +[[package]] +name = "tokio" +version = "1.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9e03c497dc955702ba729190dc4aac6f2a0ce97f913e5b1b5912fc5039d9099" +dependencies = [ + "autocfg", + "bytes", + "libc", + "memchr", + "mio", + "num_cpus", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "winapi", +] + +[[package]] +name = "tokio-macros" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9724f9a975fb987ef7a3cd9be0350edcbe130698af5b8f7a631e23d42d052484" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "typenum" version = "1.15.0" @@ -976,6 +1192,106 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-sys" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea04155a16a59f9eab786fe12a4a450e75cdb175f9e0d80da1e17db09f55b8d2" +dependencies = [ + "windows_aarch64_msvc 0.36.1", + "windows_i686_gnu 0.36.1", + "windows_i686_msvc 0.36.1", + "windows_x86_64_gnu 0.36.1", + "windows_x86_64_msvc 0.36.1", +] + +[[package]] +name = "windows-sys" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a3e1820f08b8513f676f7ab6c1f99ff312fb97b553d30ff4dd86f9f15728aa7" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc 0.42.0", + "windows_i686_gnu 0.42.0", + "windows_i686_msvc 0.42.0", + "windows_x86_64_gnu 0.42.0", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc 0.42.0", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d2aa71f6f0cbe00ae5167d90ef3cfe66527d6f613ca78ac8024c3ccab9a19e" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bb8c3fd39ade2d67e9874ac4f3db21f0d710bee00fe7cab16949ec184eeaa47" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd0f252f5a35cac83d6311b2e795981f5ee6e67eb1f9a7f64eb4500fbc4dcdb4" + +[[package]] +name = "windows_i686_gnu" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "180e6ccf01daf4c426b846dfc66db1fc518f074baa793aa7d9b9aaeffad6a3b6" + +[[package]] +name = "windows_i686_gnu" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbeae19f6716841636c28d695375df17562ca208b2b7d0dc47635a50ae6c5de7" + +[[package]] +name = "windows_i686_msvc" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2e7917148b2812d1eeafaeb22a97e4813dfa60a3f8f78ebe204bcc88f12f024" + +[[package]] +name = "windows_i686_msvc" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84c12f65daa39dd2babe6e442988fc329d6243fdce47d7d2d155b8d874862246" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dcd171b8776c41b97521e5da127a2d86ad280114807d0b2ab1e462bc764d9e1" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf7b1b21b5362cbc318f686150e5bcea75ecedc74dd157d874d754a2ca44b0ed" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09d525d2ba30eeb3297665bd434a54297e4170c7f1a44cad4ef58095b4cd2028" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c811ca4a8c853ef420abd8592ba53ddbbac90410fab6903b3e79972a631f7680" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40009d85759725a34da6d89a94e63d7bdc50a862acf0dbc7c8e488f1edcb6f5" + [[package]] name = "yuvxyb" version = "0.1.4" diff --git a/Cargo.toml b/Cargo.toml index 2816ffb..e41de71 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,8 +11,11 @@ license = "BSD-2-Clause" [dependencies] clap = { version = "4.0.18", features = ["derive"] } +futures = "0.3.25" image = "0.24.4" +progress_bar = "1.0.3" ssimulacra2 = "0.1.0" +tokio = { version = "1.21.2", features = ["full"] } yuvxyb = "0.1.3" [profile.release] diff --git a/src/main.rs b/src/main.rs index 5f16d9f..ed7f4aa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,9 @@ use clap::Parser; +use progress_bar::*; use ssimulacra2::{compute_frame_ssimulacra2, ColorPrimaries, TransferCharacteristic, Xyb}; +use std::fs; +use std::path::Path; +use std::sync::{Arc, Mutex}; use yuvxyb::Rgb; #[derive(Parser, Debug)] @@ -12,14 +16,84 @@ struct Args { /// Distorted image #[arg(help = "Distorted image", value_hint = clap::ValueHint::FilePath)] distorted: String, + + /// Location to output a .csv file with the ssimumulacra2 values + #[arg(help = "Output location. Requires --folders", value_hint = clap::ValueHint::FilePath, requires = "folders")] + out: Option, + + // TODO: Change help text to something more useful + /// If input paths are folders, process all images in the folders + #[arg( + short, + long, + help = "If input paths are folders, process all images in the folders. This assumes the files are named the same in both folders." + )] + folders: bool, } -fn main() { +#[tokio::main] +async fn main() { let args = Args::parse(); + if !args.folders { + let result = parse(args.source, args.distorted); + println!("{result:.8}"); + } else { + // args get's moved into handle_folder, so we need to clone `out` + let out_clone = args.out.clone(); + + let mut results = handle_folder(args).await; + + // Sort by frame number + results.sort_by(|a, b| a.frame.cmp(&b.frame)); + + // println!("{:#?}", results); + + // Print Mean, min, max + println!( + "Min: {}", + results + .iter() + .map(|r| r.ssimulacra2) + .min_by(|a, b| a.partial_cmp(b).unwrap()) + .unwrap() + ); + println!( + "Max: {}", + results + .iter() + .map(|r| r.ssimulacra2) + .max_by(|a, b| a.partial_cmp(b).unwrap()) + .unwrap() + ); + println!( + "Mean: {}", + results.iter().map(|r| r.ssimulacra2).sum::() / results.len() as f64 + ); + + // Print CSV + if let Some(out) = out_clone { + let mut csv = String::new(); + csv.push_str("frame,ssimulacra2\n"); + for result in results { + csv.push_str(&format!("{},{}\n", result.frame, result.ssimulacra2)); + } + // check if `out` is a directory + if Path::new(&out).is_dir() { + let mut path = Path::new(&out).to_path_buf(); + path.push("ssimulacra2.csv"); + fs::write(path, csv).expect("Unable to write file"); + } else { + fs::write(out, csv).expect("Unable to write file"); + } + } + } +} + +fn parse(source_path: String, distorted_path: String) -> f64 { // For now just assumes the input is sRGB. Trying to keep this as simple as possible for now. - let source = image::open(args.source).expect("Failed to open source file"); - let distorted = image::open(args.distorted).expect("Failed to open distorted file"); + let source = image::open(source_path).expect("Failed to open source file"); + let distorted = image::open(distorted_path).expect("Failed to open distorted file"); let source_data = source .to_rgb32f() @@ -57,8 +131,72 @@ fn main() { ) .expect("Failed to process distorted_data into XYB"); - let result = compute_frame_ssimulacra2(source_data, distorted_data) - .expect("Failed to calculate ssimulacra2"); + // Compute and return the SSIMulacra2 value + compute_frame_ssimulacra2(source_data, distorted_data).expect("Failed to calculate ssimulacra2") +} + +async fn handle_folder(args: Args) -> Vec { + let files = fs::read_dir(args.source.clone()).unwrap(); + + let results: Arc>> = Arc::new(Mutex::new(Vec::new())); + + // let mut count = 0; + + // TODO: This is a bit ugly, but it works. Reopen the folder and iterate over it again because count consumes the iterator. + let len = fs::read_dir(args.source.clone()).unwrap().count(); + + println!("Processing {} files", len); + + let mut handles = vec![]; + + init_progress_bar(len); + set_progress_bar_action("Processing", Color::Blue, Style::Bold); + + // TODO: Figure out how to multithread this? + for (count, path) in files.enumerate() { + // count += 1; + + let arg_source = args.source.clone(); + let arg_distorted = args.distorted.clone(); + + let results_clone = Arc::clone(&results); + + handles.push(tokio::spawn(async move { + let src_path = Path::new(&arg_source); + let dst_path = Path::new(&arg_distorted); + + let file_name = path.unwrap().file_name(); + + let ssimulacra2_result = parse( + src_path + .join(file_name.clone()) + .to_str() + .unwrap() + .to_owned(), + dst_path.join(file_name).to_str().unwrap().to_owned(), + ); + + results_clone.lock().unwrap().push(FrameResult { + frame: count.try_into().unwrap(), + ssimulacra2: ssimulacra2_result, + }); + + // println!("Frame {}/{} complete!", count, len); + inc_progress_bar(); + })); + } + + futures::future::join_all(handles).await; + + finalize_progress_bar(); + + let x = results.lock().unwrap().to_vec(); + x +} - println!("{result:.8}"); +// struct to hold frame number and ssimulacra2 value +#[derive(Debug, Clone)] +struct FrameResult { + frame: u32, + ssimulacra2: f64, }