From 62bfa0a63ba5c394a42a1e5fcce51a02cff12fef Mon Sep 17 00:00:00 2001 From: Ferrah Aiko Date: Fri, 20 Sep 2024 19:09:21 -0300 Subject: [PATCH] begin good stuff --- Cargo.toml | 24 ++-- src/cli/commands/avif.rs | 45 ++++--- src/cli/commands/mod.rs | 14 ++ src/cli/commands/png.rs | 259 ++++++++++++++++++++++++++++++++++++ src/decoders/mod.rs | 0 src/encoders/avif/encode.rs | 8 +- src/encoders/avif/error.rs | 2 +- src/image_file.rs | 35 ++++- src/main.rs | 9 +- src/ssim/mod.rs | 49 +++++++ src/utils.rs | 61 ++++++--- 11 files changed, 452 insertions(+), 54 deletions(-) create mode 100644 src/cli/commands/png.rs create mode 100644 src/decoders/mod.rs create mode 100644 src/ssim/mod.rs diff --git a/Cargo.toml b/Cargo.toml index ac6eabd..eabad9f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ name = "avif-converter" description = "Simple tool to batch convert multiple images to AVIF" authors = ["Ferrah Aiko "] license = "MIT" -version = "1.10.0" +version = "1.11.0" edition = "2021" readme = "" @@ -13,7 +13,7 @@ readme = "" bytesize = "1.2.0" clap = { version = "4", features = ["derive", "cargo"] } color-eyre = { version = "0.6.2", default_features = false } -env_logger = { version = "0.10.0", default_features = false, features = [ +env_logger = { version = "0.11.3", default_features = false, features = [ "auto-color", ] } hex = "0.4.3" @@ -22,30 +22,30 @@ indicatif = { version = "0.17" } log = "0.4" md5 = { version = "0.7.0", default_features = false } num_cpus = "1.16.0" -owo-colors = "3.5.0" +owo-colors = "4.0.0" rand = "0.8.5" rgb = "0.8.36" sha2 = { version = "0.10.6", features = ["asm"] } spinoff = "0.8.0" once_cell = "1.17.1" threadpool = "1.8.1" -image = { version = "0.24.7", default-features = false, features = [ - "rgb", - "png", - "jpeg", - "webp", - "bmp", -] } -rav1e = { version = "0.6.6", default_features = false, features = ["threading", "asm"] } +image = { version = "0.25.1", default-features = false, features = ["png", "jpeg", "webp", "bmp", "avif-native", "nasm", "rayon", "avif"] } +rav1e = { version = "0.7.1", default_features = false, features = ["threading", "asm"] } thiserror = "1.0" loop9 = "0.1.3" avif-serialize = "0.8.1" notify-rust = { version = "4.8.0", features = ["images"] } -thread-priority = "0.13.1" +thread-priority = "1.1.0" notify = "6.0.1" blake2 = { version = "0.10.6" } +opencv = { version = "0.93.0", default-features = false, features = ["imgproc", "imgcodecs", "rgb"], optional = true} [profile.release] lto = false opt-level = 3 panic = "abort" + +[features] +default = ["ssim"] +ssim = ["opencv"] +opencv = ["dep:opencv"] diff --git a/src/cli/commands/avif.rs b/src/cli/commands/avif.rs index e078539..8e0cc43 100644 --- a/src/cli/commands/avif.rs +++ b/src/cli/commands/avif.rs @@ -1,4 +1,4 @@ -use std::{path::PathBuf, sync::atomic::Ordering, time::Instant}; +use std::{path::PathBuf, process::exit, sync::atomic::Ordering, time::Instant}; use bytesize::ByteSize; use clap::Args; @@ -11,16 +11,18 @@ use crate::{ cli::{Args as Globals, FINAL_STATS, ITEMS_PROCESSED, SUCCESS_COUNT}, console::ConsoleMsg, image_file::ImageFile, - utils::{calculate_tread_count, search_dir, sys_threads, PROGRESS_BAR}, + utils::{calculate_tread_count, parse_files, sys_threads, PROGRESS_BAR}, }; use color_eyre::Result; +use super::EncodeFuncs; + #[derive(Args, Debug, Clone)] #[clap(author, about, long_about = None)] pub struct Avif { /// File or directory containing images to convert - #[clap(value_name = "PATH")] - pub path: PathBuf, + #[clap(value_name = "PATH", required = true)] + pub path: Vec, /// Enable benchmark mode #[clap( @@ -39,31 +41,40 @@ pub struct Avif { /// Send a notification to the desktop when all jobs are finished #[clap(short = 'N', long, default_value_t = false)] pub notify: bool, + + /// Measure SSIM of encoded vs original image/s. + #[cfg(feature = "ssim")] + #[clap(long = "ssim", default_value_t = false)] + pub ssim: bool, + + /// Save SSIM difference as an image along with the encoded file. + #[cfg(feature = "ssim")] + #[clap(long = "ssim_save", default_value_t = false, requires = "ssim")] + pub ssim_save: bool, } -impl Avif { - pub fn run_conv(self, globals: &Globals) -> Result<()> { +impl EncodeFuncs for Avif { + fn run_conv(self, globals: &Globals) -> Result<()> { let console = ConsoleMsg::new(globals.quiet, self.notify); let error_con = ConsoleMsg::new(globals.quiet, self.notify); - let u = { - if self.path.is_dir() { - self.dir_conv(console, globals) - } else if self.path.is_file() { - self.single_file_conv(console, globals) - } else { - bail!("Unsupported operation") - } + let l_size = self.path.len(); + + let u = if l_size > 1 { + self.batch_conv(console, globals) + } else { + self.single_file_conv(console, globals) }; if let Err(error) = u { error_con.notify_error(&error.to_string())?; + exit(1); } Ok(()) } - fn dir_conv(self, console: ConsoleMsg, globals: &Globals) -> Result<()> { + fn batch_conv(self, console: ConsoleMsg, globals: &Globals) -> Result<()> { if self.output_file.is_some() { bail!("Cannot assign an output file to a directory") } @@ -71,7 +82,7 @@ impl Avif { let mut console = console; console.set_spinner("Searching for files..."); - let mut paths = search_dir(&self.path); + let mut paths = parse_files(&self.path); let psize = paths.len(); paths.sort_by(|a, b| a.metadata.name.cmp(&b.metadata.name)); @@ -203,7 +214,7 @@ impl Avif { fn single_file_conv(self, console: ConsoleMsg, globals: &Globals) -> Result<()> { let mut console = console; - let mut image = ImageFile::new_from_path(&self.path)?; + let mut image = ImageFile::new_from_path(&self.path[0])?; let image_size = image.metadata.size; console.print_message(format!( diff --git a/src/cli/commands/mod.rs b/src/cli/commands/mod.rs index ab7a291..df711b0 100644 --- a/src/cli/commands/mod.rs +++ b/src/cli/commands/mod.rs @@ -1,8 +1,14 @@ use clap::Subcommand; +use crate::console::ConsoleMsg; + use self::{avif::Avif, watch::Watch}; +use super::Args as Globals; +use color_eyre::Result; + pub mod avif; +//pub mod png; pub mod watch; #[derive(Debug, Subcommand, Clone)] @@ -12,3 +18,11 @@ pub enum Commands { /// Watch directory for new image files and convert them Watch(Watch), } + +pub trait EncodeFuncs { + fn run_conv(self, globals: &Globals) -> Result<()>; + + fn batch_conv(self, console: ConsoleMsg, globals: &Globals) -> Result<()>; + + fn single_file_conv(self, console: ConsoleMsg, globals: &Globals) -> Result<()>; +} diff --git a/src/cli/commands/png.rs b/src/cli/commands/png.rs new file mode 100644 index 0000000..a3eb90a --- /dev/null +++ b/src/cli/commands/png.rs @@ -0,0 +1,259 @@ +use std::{path::PathBuf, sync::atomic::Ordering, time::Instant}; + +use bytesize::ByteSize; +use clap::Args; +use color_eyre::eyre::bail; +use log::{debug, trace}; +use owo_colors::OwoColorize; +use threadpool::ThreadPool; + +use crate::{ + cli::{Args as Globals, FINAL_STATS, ITEMS_PROCESSED, SUCCESS_COUNT}, + console::ConsoleMsg, + image_file::ImageFile, + utils::{calculate_tread_count, search_dir, sys_threads, PROGRESS_BAR}, +}; +use color_eyre::Result; + +use super::EncodeFuncs; + +#[derive(Args, Debug, Clone)] +#[clap(author, about, long_about = None)] +pub struct Png { + /// File or directory containing images to convert + #[clap(value_name = "PATH")] + pub path: PathBuf, + + /// Enable benchmark mode + #[clap( + long, + default_value_t = false, + conflicts_with = "name_type", + conflicts_with = "keep", + conflicts_with = "output_file", + global = true + )] + pub benchmark: bool, + + #[clap(short, long, conflicts_with = "name_type", value_name = "OUTPUT")] + pub output_file: Option, + + /// Send a notification to the desktop when all jobs are finished + #[clap(short = 'N', long, default_value_t = false)] + pub notify: bool, +} + +impl EncodeFuncs for Png { + fn run_conv(self, globals: &Globals) -> Result<()> { + let console = ConsoleMsg::new(globals.quiet, self.notify); + let error_con = ConsoleMsg::new(globals.quiet, self.notify); + + let u = { + if self.path.is_dir() { + self.dir_conv(console, globals) + } else if self.path.is_file() { + self.single_file_conv(console, globals) + } else { + bail!("Unsupported operation") + } + }; + + if let Err(error) = u { + error_con.notify_error(&error.to_string())?; + } + + Ok(()) + } + + fn dir_conv(self, console: ConsoleMsg, globals: &Globals) -> Result<()> { + if self.output_file.is_some() { + bail!("Cannot assign an output file to a directory") + } + + let mut console = console; + console.set_spinner("Searching for files..."); + + let mut paths = search_dir(&self.path); + let psize = paths.len(); + + paths.sort_by(|a, b| a.metadata.name.cmp(&b.metadata.name)); + + let con = console.finish_spinner(&format!("Found {psize} files.")); + + let job_num = calculate_tread_count(globals.threads, psize); + + let pool = ThreadPool::with_name("Encoder Thread".to_string(), job_num.spawn_threads); + + let initial_size: u64 = paths.iter().map(|item| item.metadata.size).sum(); + + con.setup_bar(psize as u64); + + let start = Instant::now(); + + for mut item in paths.drain(..) { + let globals = globals.clone(); + pool.execute(move || { + Globals::set_encoder_priority(globals.priority); + let enc_start = Instant::now(); + + let bar = if globals.quiet { + None + } else { + Some(PROGRESS_BAR.clone()) + }; + + if let Ok(r_size) = item.convert_to_avif_stored( + globals.quality, + globals.speed, + job_num.task_threads, + globals.bit_depth, + globals.remove_alpha, + bar, + ) { + SUCCESS_COUNT.fetch_add(1, Ordering::SeqCst); + FINAL_STATS.fetch_add(r_size, Ordering::SeqCst); + } + + if !self.benchmark { + item.save_avif(None, globals.name_type, globals.keep) + .unwrap(); + } + + trace!( + "Finished encoding: {} | {:?} | {:?}", + item.original_name(), + enc_start.elapsed().bold().cyan(), + start.elapsed().bold().green() + ); + + drop(item); + + ITEMS_PROCESSED.fetch_add(1, Ordering::SeqCst); + + if globals.quiet { + debug!( + "Items Processed: {}", + ITEMS_PROCESSED.load(Ordering::Relaxed) + ); + } + }); + } + + pool.join(); + + let elapsed = start.elapsed(); + + con.finish_bar(); + + let texts = [ + *"Original folder size".bold().0, + *"New folder size".bold().0, + ]; + + debug!("Final stats: {}", FINAL_STATS.load(Ordering::Relaxed)); + debug!("Initial size: {}", initial_size); + + let initial_delta = FINAL_STATS.load(Ordering::Relaxed) as f32 / initial_size as f32; + + let delta = (initial_delta * 100.) - 100.; + + debug!("Delta: {}", delta); + + let percentage = if delta < 0. { + let st1 = format!("{delta:.2}%"); + format!("{}", st1.green()) + } else { + let st1 = format!("+{delta:.2}%"); + format!("{}", st1.red()) + }; + + let times = { + let ratio = 1. / initial_delta; + debug!("Ratio: {}", ratio); + if ratio > 0. { + let st1 = format!("~{:.1}X smaller", ratio); + format!("{}", st1.green()) + } else { + let st1 = format!("~{:.1}X bigger", ratio); + format!("{}", st1.red()) + } + }; + + con.print_message(format!( + "Encoded {} files in {elapsed:.2?}.\n{} {} | {} {} ({} or {})", + SUCCESS_COUNT.load(Ordering::SeqCst), + texts[0], + ByteSize::b(initial_size).to_string_as(true).blue().bold(), + texts[1], + ByteSize::b(FINAL_STATS.load(Ordering::SeqCst)) + .to_string_as(true) + .green() + .bold(), + percentage, + times + )); + + con.notify_text(&format!( + "Encoded {} files in {elapsed:.2?}\n{} → {}", + SUCCESS_COUNT.load(Ordering::SeqCst), + ByteSize::b(initial_size).to_string_as(true), + ByteSize::b(FINAL_STATS.load(Ordering::SeqCst)).to_string_as(true) + ))?; + + Ok(()) + } + + fn single_file_conv(self, console: ConsoleMsg, globals: &Globals) -> Result<()> { + let mut console = console; + let mut image = ImageFile::new_from_path(&self.path)?; + let image_size = image.metadata.size; + + console.print_message(format!( + "Encoding single file {} ({})", + image.metadata.name.bold(), + ByteSize::b(image.metadata.size) + .to_string_as(true) + .bold() + .blue() + )); + + console.set_spinner("Processing..."); + + let start = Instant::now(); + + let fsz = image.convert_to_avif_stored( + globals.quality, + globals.speed, + sys_threads(globals.threads), + globals.bit_depth, + globals.remove_alpha, + None, + )?; + + if !self.benchmark { + image.save_avif(self.output_file, globals.name_type, globals.keep)?; + } + + let bmp = image.bitmap.clone(); + + drop(image); + + console.notify_image( + &format!( + "Finished in {:.2?} \n {} → {}", + start.elapsed(), + ByteSize::b(image_size).to_string_as(true), + ByteSize::b(fsz).to_string_as(true) + ), + bmp, + )?; + + console.finish_spinner(&format!( + "Encoding finished in {:?} ({})", + start.elapsed(), + ByteSize::b(fsz).to_string_as(true).bold().green() + )); + + Ok(()) + } +} diff --git a/src/decoders/mod.rs b/src/decoders/mod.rs new file mode 100644 index 0000000..e69de29 diff --git a/src/encoders/avif/encode.rs b/src/encoders/avif/encode.rs index 3a539ea..1ea74c9 100644 --- a/src/encoders/avif/encode.rs +++ b/src/encoders/avif/encode.rs @@ -18,9 +18,9 @@ pub struct EncodedImage { /// AVIF (HEIF+AV1) encoded image data pub avif_file: Vec, /// FYI: number of bytes of AV1 payload used for the color - pub color_byte_size: usize, + pub _color_byte_size: usize, /// FYI: number of bytes of AV1 payload used for the alpha channel - pub alpha_byte_size: usize, + pub _alpha_byte_size: usize, } /// Encoder config builder @@ -287,8 +287,8 @@ impl Encoder { Ok(EncodedImage { avif_file, - color_byte_size, - alpha_byte_size, + _color_byte_size: color_byte_size, + _alpha_byte_size: alpha_byte_size, }) } } diff --git a/src/encoders/avif/error.rs b/src/encoders/avif/error.rs index 95f3562..b6ec59f 100644 --- a/src/encoders/avif/error.rs +++ b/src/encoders/avif/error.rs @@ -2,7 +2,7 @@ use thiserror::Error; #[derive(Debug)] #[doc(hidden)] -pub struct EncodingErrorDetail; // maybe later +pub struct _EncodingErrorDetail; // maybe later /// Failures enum #[derive(Debug, Error)] diff --git a/src/image_file.rs b/src/image_file.rs index 03c2926..a926efa 100644 --- a/src/image_file.rs +++ b/src/image_file.rs @@ -39,7 +39,8 @@ impl ImageFile { || ext == "jpeg" || ext == "jfif" || ext == "webp" - || ext == "bmp") + || ext == "bmp" + || ext == "avif") { bail!("Unsupported image format"); } @@ -138,24 +139,50 @@ impl ImageFile { let avif_name = fpath.join(format!("{fname}.avif")); + // If `path` is Some, save to the provided path if let Some(new_path) = path { - fs::write(new_path, &self.encoded_data)?; + let target_avif_name = new_path.join(format!("{fname}.avif")); + + if !keep { + // If `keep` is false and we have a target path, we want to replace the original file + let mut orig_file = OpenOptions::new().write(true).open(&binding)?; + orig_file.set_len(self.encoded_data.len() as u64)?; + orig_file.seek(std::io::SeekFrom::Start(0))?; + orig_file.write_all(&self.encoded_data)?; + + // Attempt to rename (move) to the new path + match fs::rename(&binding, &target_avif_name) { + Ok(_) => return Ok(()), // Success, file moved + Err(_) => { + // Rename failed (likely due to different filesystems), fallback to copy+delete + fs::copy(&binding, &target_avif_name)?; + fs::remove_file(&binding)?; // Remove original after successful copy + } + } + + return Ok(()); + } + + // If `keep` is true, just save the AVIF to the target location + fs::write(&target_avif_name, &self.encoded_data)?; + return Ok(()); } + // If no `path` is provided, proceed with in-place modifications if !keep { let mut orig_file = OpenOptions::new().write(true).open(&binding)?; orig_file.set_len(self.encoded_data.len() as u64)?; - orig_file.seek(std::io::SeekFrom::Start(0))?; - orig_file.write_all(&self.encoded_data)?; + // Rename (move) the file to the new AVIF name fs::rename(&binding, &avif_name)?; return Ok(()); } + // If `keep` is true, save AVIF to the same directory fs::write(&avif_name, &self.encoded_data)?; Ok(()) diff --git a/src/main.rs b/src/main.rs index 390e80b..8c224d0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,20 @@ -use cli::{commands::Commands, Args}; +use cli::{ + commands::{Commands, EncodeFuncs}, + Args, +}; use color_eyre::eyre::Result; mod cli; mod console; +mod decoders; mod encoders; mod image_file; mod name_fun; mod utils; +#[cfg(feature = "ssim")] +mod ssim; + fn main() -> Result<()> { color_eyre::install()?; env_logger::builder().format_timestamp(None).init(); diff --git a/src/ssim/mod.rs b/src/ssim/mod.rs new file mode 100644 index 0000000..100e967 --- /dev/null +++ b/src/ssim/mod.rs @@ -0,0 +1,49 @@ +use image::{GrayImage, Luma}; +use rayon::prelude::*; + +pub fn calculate_ssim_and_diff(img1: &GrayImage, img2: &GrayImage) -> (f64, GrayImage) { + assert_eq!(img1.dimensions(), img2.dimensions()); + + let (width, height) = img1.dimensions(); + let mut diff_image = GrayImage::new(width, height); // To store the difference image + + let total_ssim: f64 = (0..height) + .into_par_iter() + .map(|y| { + let mut ssim_row_total = 0.0; + + for x in 0..width { + let p1 = img1.get_pixel(x, y)[0] as f64; + let p2 = img2.get_pixel(x, y)[0] as f64; + + let mean_p1 = p1; + let mean_p2 = p2; + + let variance_p1 = p1 * p1; + let variance_p2 = p2 * p2; + + let covariance = p1 * p2; + + let c1 = 0.01 * 255.0; // Constants to stabilize division + let c2 = 0.03 * 255.0; + + let ssim = ((2.0 * mean_p1 * mean_p2 + c1) * (2.0 * covariance + c2)) + / ((mean_p1.powi(2) + mean_p2.powi(2) + c1) * (variance_p1 + variance_p2 + c2)); + + // Add SSIM value to the total + ssim_row_total += ssim; + + // Generate difference image by scaling the absolute difference + let diff_value = ((p1 - p2).abs() * 255.0 / 255.0) as u8; // Scale the difference to fit 0-255 range + diff_image.put_pixel(x, y, Luma([diff_value])); // Store difference in diff image + } + + ssim_row_total + }) + .sum(); // Sum all the rows' SSIM totals in parallel + + // Compute the final SSIM score (average over all pixels) + let avg_ssim = total_ssim / (width * height) as f64; + + (avg_ssim, diff_image) +} diff --git a/src/utils.rs b/src/utils.rs index d716db8..cff4054 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,4 +1,4 @@ -use std::{fmt::Write, fs, path::Path}; +use std::{fmt::Write, fs, path::PathBuf}; use indicatif::{ProgressBar, ProgressState, ProgressStyle}; use once_cell::sync::Lazy; @@ -8,17 +8,31 @@ use crate::image_file::ImageFile; pub static PROGRESS_BAR: Lazy = Lazy::new(|| ProgressBar::new(0).with_style(bar_style())); -pub fn search_dir(dir: &Path) -> Vec { - let paths = fs::read_dir(dir).unwrap(); - - Vec::from_iter(paths.filter_map(|entry| { - let entry = entry.unwrap(); - let path = entry.path(); - if let Ok(image_file) = ImageFile::new_from_path(&path) { - return Some(image_file); - } - None - })) +pub fn parse_files(paths: &Vec) -> Vec { + paths + .iter() + .flat_map(|item| { + if item.is_dir() { + // If it's a directory, we attempt to read the directory entries + if let Ok(dir) = fs::read_dir(item) { + // Flatten the directory iterator, map each entry to ImageFile, and collect results + dir.flatten() + .filter_map(|entry| { + // Try to create an ImageFile from the entry path + ImageFile::new_from_path(&entry.path()).ok() + }) + .collect::>() // Collect directory entries into a vector + } else { + Vec::new() // If directory read fails, return an empty Vec + } + } else if item.is_file() { + // If it's a file, try to create an ImageFile from it + ImageFile::new_from_path(item).ok().into_iter().collect() + } else { + Vec::new() // If it's neither a file nor a directory, return an empty Vec + } + }) + .collect() } pub fn bar_style() -> ProgressStyle { @@ -68,10 +82,27 @@ pub fn truncate_str(str: &str, size: usize) -> String { assert!(str.len() > 3); if str.len() <= size { - return str.to_string() + return str.to_string(); } - let mut truncated = str[..size].to_string(); + let file_name: Vec = str.chars().rev().collect(); + let file_extension = file_name.iter().position(|c| !c.is_alphanumeric()); + let mut truncated = str[..size - file_extension.unwrap_or(size)].to_string(); + truncated.push_str( + &file_name[file_extension.unwrap_or(0)..file_extension.unwrap_or(size)] + .iter() + .rev() + .cloned() + .take(3) + .collect::(), + ); + let ext = file_extension + .unwrap_or(size) + .to_string() + .chars() + .rev() + .collect::(); + truncated.push_str(&ext); truncated.push_str("..."); truncated -} \ No newline at end of file +}