image_compressor 1.5.3

//! Module that contains things related to compressing an image.
//!
//! # Compressor
//!
//! The `compress_to_jpg` function resizes the given image and compresses it by a certain percentage.
//! # Examples
//! ```rust,no_run
//! use std::path::PathBuf;
//! use image_compressor::compressor::Compressor;
//! use image_compressor::Factor;
//!
//! let source_file = PathBuf::from("source").join("file1.jpg");
//! let dest_dir = PathBuf::from("dest");
//!
//! let mut compressor = Compressor::new(source_file, dest_dir);
//! compressor.set_factor(Factor::new(80., 0.8));
//! compressor.compress_to_jpg();
//! ```

use image::imageops::FilterType;
use image::{ImageError, ImageFormat, ImageReader, Limits};
use mozjpeg::{ColorSpace, Compress, ScanMode};
use std::error::Error;
use std::fs::File;
use std::io::{BufReader, BufWriter, ErrorKind, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::{fs, io};

/// Factor struct that used for setting quality and resize ratio in the new image.
///
/// The [`Compressor`] and [`FolderCompressor`](super::FolderCompressor) need `Factor` for compressing images.
///
/// So, to create a new `Compressor` or `FolderCompressor` instance
/// you need to define a new `Factor` instance contains the quality ratio of image and file size ratio to compress.
///
/// The recommended range of quality is 60 to 80.
#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
pub struct Factor {
    /// Quality of the new compressed image.
    /// Values range from 0 to 100 in float.
    quality: f32,

    /// Ratio for resize the new compressed image.
    /// Values range from 0 to 1 in float.
    size_ratio: f32,
}

impl Factor {
    /// Create a new `Factor` instance.
    /// The `quality` range from 0 to 100 in float,
    /// and `size_ratio` range from 0 to 1 in float.
    ///
    /// # Panics
    ///
    /// - If the quality value is 0 or less.
    /// - If the quality value exceeds 100.
    /// - If the size ratio value is 0 or less.
    /// - If the size ratio value exceeds 1.
    pub fn new(quality: f32, size_ratio: f32) -> Self {
        if (quality > 0. && quality <= 100.) && (size_ratio > 0. && size_ratio <= 1.) {
            Self {
                quality,
                size_ratio,
            }
        } else {
            panic!("Wrong Factor argument!");
        }
    }

    /// Getter for `quality` of `Factor`.
    pub fn quality(&self) -> f32 {
        self.quality
    }

    /// Getter for `size_ratio` of `Factor`.
    pub fn size_ratio(&self) -> f32 {
        self.size_ratio
    }
}

impl Default for Factor {
    fn default() -> Self {
        Self {
            quality: 80.,
            size_ratio: 0.8,
        }
    }
}

/// Compressor struct.
pub struct Compressor<O: AsRef<Path>, D: AsRef<Path>> {
    factor: Factor,
    source_path: O,
    dest_path: D,
    delete_source: bool,
    memory_limit: Option<u64>,
}

impl<O: AsRef<Path>, D: AsRef<Path>> Compressor<O, D> {
    /// Create a new `Compressor` instance.
    pub fn new(source_path: O, dest_dir_path: D) -> Self {
        Compressor {
            factor: Factor::default(),
            source_path,
            dest_path: dest_dir_path,
            delete_source: false,
            memory_limit: None,
        }
    }

    /// Set factor for the new compressed image.
    pub fn set_factor(&mut self, factor: Factor) {
        self.factor = factor;
    }

    /// Set the maximum number of bytes the decoder may allocate while reading the source image.
    ///
    /// By default there is no limit, so images of any resolution can be decoded.
    /// Set a limit when decoding images from an untrusted source, where a maliciously
    /// crafted file could otherwise exhaust memory.
    ///
    /// # Examples
    /// ```
    /// use image_compressor::compressor::Compressor;
    /// use std::path::Path;
    ///
    /// let mut comp = Compressor::new(Path::new("source.png"), Path::new("dest"));
    /// comp.set_memory_limit(512 * 1024 * 1024);   // 512 MiB
    /// ```
    pub fn set_memory_limit(&mut self, bytes: u64) {
        self.memory_limit = Some(bytes);
    }

    /// Sets whether the program deletes the source file.
    pub fn set_delete_source(&mut self, to_delete: bool) {
        self.delete_source = to_delete;
    }

    /// Compress the image to jpg format.
    /// The new image will be saved in the destination directory.
    fn compress(
        &self,
        img: image::DynamicImage,
        target_width: usize,
        target_height: usize,
        quality: f32,
    ) -> Result<Vec<u8>, Box<dyn Error>> {
        let mut comp = Compress::new(ColorSpace::JCS_RGB);
        comp.set_scan_optimization_mode(ScanMode::Auto);
        comp.set_quality(quality);

        comp.set_size(target_width, target_height);

        comp.set_optimize_scans(true);
        let mut comp = comp.start_compress(Vec::new())?;

        let mut line = 0;
        let img_vec = img.to_rgb8().into_vec();
        while line < target_height {
            comp.write_scanlines(&img_vec[line * target_width * 3..(line + 1) * target_width * 3])?;
            line += 1;
        }
        let compressed = comp.finish()?;
        Ok(compressed)
    }

    /// Resize the image vector.
    fn resize(
        &self,
        img: image::DynamicImage,
        resize_ratio: f32,
    ) -> Result<(image::DynamicImage, usize, usize), Box<dyn Error>> {
        let width = img.width() as usize;
        let height = img.height() as usize;

        let width = width as f32 * resize_ratio;
        let height = height as f32 * resize_ratio;

        let resized_img = img.resize(width as u32, height as u32, FilterType::Triangle);

        let resized_width = resized_img.width() as usize;
        let resized_height = resized_img.height() as usize;

        Ok((resized_img, resized_width, resized_height))
    }

    /// Guess actual image format
    fn guess_image_format(&self, source_file_path: &Path) -> Result<ImageFormat, ImageError> {
        let mut file = File::open(source_file_path)?;
        let _ = file.seek(SeekFrom::Start(0));
        let mut buf = vec![0; 10];
        let _ = file.read_exact(&mut buf);

        image::guess_format(buf.as_slice())
    }

    /// Compress a file.
    ///
    /// Compress the given image file and save it to target_dir.
    /// If the extension of the given image file is not jpg or jpeg, convert the image to jpg file.
    /// If the image module can not open the file, such as pdf, mp4, etc., just copy it to target_dir.
    /// Compress quality and resize ratio calculate based on file size of the image.
    /// For a continuous multithreading process, every single error doesn't occur panic or exception and just print error message with return Ok.
    ///
    /// If the flag to delete the source is true, the function delete the source file.
    pub fn compress_to_jpg(&self) -> Result<PathBuf, Box<dyn Error>> {
        let source_file_path = self.source_path.as_ref();
        let target_dir = self.dest_path.as_ref();

        let file_name = match source_file_path.file_name() {
            Some(e) => e.to_str().unwrap_or(""),
            None => "",
        };

        let file_stem = source_file_path.file_stem().unwrap();

        let mut target_file_name = PathBuf::from(file_stem);
        target_file_name.set_extension("jpg");
        let target_file = target_dir.join(&target_file_name);
        if target_file.is_file() {
            return Err(Box::new(io::Error::new(
                ErrorKind::AlreadyExists,
                format!(
                    "A file with the same name exists: {}",
                    target_file.file_name().unwrap().to_str().unwrap()
                ),
            )));
        }

        let Ok(guessed_format) = self.guess_image_format(source_file_path) else {
            return Err(Box::new(io::Error::new(
                ErrorKind::InvalidInput,
                "Unrecognized image format",
            )));
        };

        // `ImageReader` rather than `image::load`, because `image::load` applies
        // `Limits::default()`, which caps decoder allocations at 512 MiB and so rejects
        // high resolution images outright. See issue #19.
        let mut reader = ImageReader::with_format(
            BufReader::new(File::open(source_file_path)?),
            guessed_format,
        );
        let mut limits = Limits::no_limits();
        limits.max_alloc = self.memory_limit;
        reader.limits(limits);

        let image_vec = match reader.decode() {
            Ok(p) => p,
            Err(e) => {
                let m = format!(
                    "Cannot open file {} as image. Just copy it: {}",
                    file_name, e
                );
                fs::copy(source_file_path, target_dir.join(file_name))?;
                return Err(Box::new(io::Error::new(ErrorKind::InvalidData, m)));
            }
        };

        let (resized_img_data, target_width, target_height) =
            self.resize(image_vec, self.factor.size_ratio())?;
        let compressed_img_data = match self.compress(
            resized_img_data,
            target_width,
            target_height,
            self.factor.quality(),
        ) {
            Ok(p) => p,
            Err(e) => {
                let m = format!("Cannot compress file {}: {}", file_name, e);
                return Err(Box::new(io::Error::new(ErrorKind::InvalidData, m)));
            }
        };

        let mut file = BufWriter::new(File::create(&target_file)?);
        file.write_all(&compressed_img_data)?;

        // Delete the source file when the flag is true.
        if self.delete_source {
            fs::remove_file(&self.source_path)?;
        }
        Ok(target_file)
    }
}

#[cfg(test)]
mod tests {
    use std::fs;

    use super::*;

    use image::ImageBuffer;
    use std::path::{Path, PathBuf};

    /// Create test directory and an image file in it.
    fn setup<T: AsRef<Path>>(test_name: T) -> (PathBuf, Vec<PathBuf>) {
        let test_dir = test_name.as_ref().to_path_buf();
        if test_dir.is_dir() {
            fs::remove_dir_all(&test_dir).unwrap();
        }
        fs::create_dir_all(&test_dir).unwrap();

        const WIDTH: u32 = 256;
        const HEIGHT: u32 = 256;
        let img_stripe = ImageBuffer::from_fn(WIDTH, HEIGHT, |x, _| {
            if x % 10 == 0 {
                image::Luma([0u8])
            } else {
                image::Luma([255u8])
            }
        });
        let stripe_path = test_dir.join("img_stripe.png");
        img_stripe.save(&stripe_path).unwrap();
        let img_rgb = ImageBuffer::from_fn(WIDTH, HEIGHT, |x, y| {
            image::Rgb([(x * 7) as u8, (y * 13) as u8, (x * y) as u8])
        });
        let rgb_path = test_dir.join("img_rgb.gif");
        img_rgb.save(&rgb_path).unwrap();

        (test_dir, vec![stripe_path, rgb_path])
    }

    fn cleanup<T: AsRef<Path>>(test_dir: T) {
        if test_dir.as_ref().is_dir() {
            fs::remove_dir_all(&test_dir).unwrap();
        }
    }

    #[test]
    fn skip_wrong_ext_test() {
        let (test_dir, _) = setup("skip_wrong_ext_test_dir");
        let txt_data = "Hello, World!";
        let mut txt_path = PathBuf::from(&test_dir).join("skip_wrong_ext_test.txt");
        let mut txt_file = File::create(&txt_path).unwrap();
        write!(txt_file, "{}", txt_data).unwrap();

        let compressor = Compressor::new(&txt_path, &test_dir);
        assert!(compressor.compress_to_jpg().is_err());
        assert!(txt_path.is_file());
        txt_path.set_extension("jpg");
        assert!(!txt_path.is_file());
        cleanup(test_dir);
    }

    #[test]
    fn compress_to_jpg_test() {
        let (test_dir, mut test_images) = setup("compress_to_jpg_test");

        let dest_dir = PathBuf::from("compress_to_jpg_dest_dir");
        fs::create_dir_all(&dest_dir).unwrap();

        for test_image in &test_images {
            let mut compressor = Compressor::new(test_image, &dest_dir);
            compressor.set_factor(Factor::new(80., 1.0));
            compressor.compress_to_jpg().unwrap();
        }
        test_images = test_images
            .iter()
            .map(|image| dest_dir.join(image.file_name().unwrap()))
            .collect();
        for new_image in &test_images {
            let mut new_test_image = new_image.clone();
            new_test_image.set_extension("jpg");
            assert!(new_test_image.is_file());
        }
        cleanup(test_dir);
        cleanup(dest_dir);
    }

    #[test]
    fn compress_to_jpg_with_delete_test() {
        let (test_dir, mut test_images) = setup("compress_to_jpg_with_delete_test");

        let dest_dir = PathBuf::from("compress_to_jpg_with_delete_dest_dir");
        fs::create_dir_all(&dest_dir).unwrap();

        for test_image in &test_images {
            let mut compressor = Compressor::new(test_image, &dest_dir);
            compressor.set_delete_source(true);
            compressor.compress_to_jpg().unwrap();
        }
        for test_image in &test_images {
            assert!(!test_image.is_file());
        }
        test_images = test_images
            .iter()
            .map(|image| dest_dir.join(image.file_name().unwrap()))
            .collect();
        for new_image in &test_images {
            let mut new_test_image = new_image.clone();
            new_test_image.set_extension("jpg");
            assert!(new_test_image.is_file());
        }
        cleanup(test_dir);
        cleanup(dest_dir);
    }

    /// The decoder must be unlimited by default, and must honour a limit once one is set.
    /// Guards against the `image::load` default of 512 MiB coming back. See issue #19.
    #[test]
    fn memory_limit_test() {
        let (test_dir, test_images) = setup("memory_limit_test");
        let dest_dir = PathBuf::from("memory_limit_test_dest_dir");
        fs::create_dir_all(&dest_dir).unwrap();

        // A 256x256 image needs far more than 16 bytes to decode.
        let mut limited = Compressor::new(&test_images[0], &dest_dir);
        limited.set_memory_limit(16);
        assert!(limited.compress_to_jpg().is_err());

        let unlimited = Compressor::new(&test_images[0], &dest_dir);
        assert!(unlimited.compress_to_jpg().is_ok());

        cleanup(test_dir);
        cleanup(dest_dir);
    }
}