Rust で jpeg をリサイズして mozjpeg で高圧縮化

Rust で jpeg をリサイズして mozjpeg で高圧縮化（その2）という記事を書いたのでそちらを参照してください。

環境

toolchain: 1.33.0-i686-pc-windows-msvc
(1.33.0 以降の stable だと mozjpeg crate がビルドできないようです)
1.33.0 でもビルド出来ます。Rust の内部処理が変わって、mozjpeg 側でエラーがあると panic せずにスレッドが落ちるようです。

目的

デジカメ画像を PC で表示する程度なら 1280px もあれば十分なのでリサイズして、高圧縮 jpeg ライブラリの mozjpeg を使ってファイルサイズをできるだけ小さくしようという試みです。Rust で。

リサイズ

Rust で画像を扱うなら、image crate ですね。JPEG・PNG・GIF などが扱えて、リサイズ・回転などの画像処理も可能です。

resize関数は jpeg ファイルへのパスを引数に取って、リサイズ後の RGB 画像のバイト列（Bitmap データ）と、幅・高さを返します。幅・高さともに 1280px より小さい場合は読み込んだ画像のデータをそのまま返すようにしています。

use image::{self, GenericImageView, FilterType};

const TARGET_SIZE: usize = 1280;

fn resize(path: &Path) -> Result<(Vec<u8>, usize, usize), String> {
    let img = image::open(path).map_err(|e| e.to_string())?;
    let width = img.width() as usize;
    let height = img.height() as usize;

    if width > TARGET_SIZE || height > TARGET_SIZE {
        let (target_width, target_height) = 
            if width > height {
                let ratio: f32 = TARGET_SIZE as f32 / width as f32;
                (TARGET_SIZE, (height as f32 * ratio) as usize)
            } else {
                let ratio: f32 = TARGET_SIZE as f32 / height as f32;
                ((width as f32 * ratio) as usize, TARGET_SIZE)
            };
        let resized_img = img.resize(
                target_width as u32,
                target_height as u32,
                FilterType::Lanczos3);
        Ok((resized_img.to_rgb().to_vec(), resized_img.width() as usize, resized_img.height() as usize))
    } else {
        Ok((img.to_rgb().to_vec(), width, height))
    }
}

圧縮

mozjpeg crate を使います。compress関数はリサイズ後の画像データ・幅・高さを引数に取って、圧縮後の画像データを返すようにしています。

set_qualityで圧縮品質を指定できます。70 でも画像の劣化は感じません。write_scanlinesには画像の幅の 3 倍（1pixel に RGB の 3bytes）ずつを画像データから slice で抜き出して渡し、それを画像の高さ回繰り返します。例えば 3x3 の画像なら、write_scanlines(&data[0..9])、write_scanlines(&data[9..18])、write_scanlines(&data[18..27])を実行することになります。

use mozjpeg::{Compress, ColorSpace, ScanMode};

fn compress(resized_img_data: Vec<u8>, target_width: usize, target_height: usize) -> Result<Vec<u8>, String> {
    let mut comp = Compress::new(ColorSpace::JCS_RGB);
    comp.set_scan_optimization_mode(ScanMode::AllComponentsTogether);
    comp.set_quality(70.0);

    comp.set_size(target_width, target_height);

    comp.set_mem_dest();
    comp.start_compress();

    let mut line = 0;
    loop {
        if line > target_height - 1 {
            break;
        }
        comp.write_scanlines(&resized_img_data[line * target_width * 3..(line + 1) * target_width * 3]);
        line += 1;
    }
    comp.finish_compress();

    let compressed = comp.data_to_vec()
        .map_err(|_| "data_to_vec failed".to_string())?;
    Ok(compressed)
}

main 関数

.exe に画像ファイルを D&D したら compressed フォルダを作って、そこに圧縮された画像が保存されるようにします。

use std::fs::{self, File};
use std::io::{self, BufWriter, Write, Read};
use std::env;

fn main() {
    let target_dir = match env::args().nth(1) {
        Some(v) => PathBuf::from(v).parent().unwrap().join("compressed"),
        None => return,
    };
    if !target_dir.exists() {
        fs::create_dir(&target_dir).unwrap();
    }

    let source_files = env::args()
        .skip(1)
        .map(PathBuf::from)
        .filter(|p| p.is_file() && p.file_name().is_some() && p.extension().is_some())
        .collect::<Vec<PathBuf>>();

    source_files.iter().for_each(|path| {
        match process(&path, &target_dir) {
            Ok(file_name) => println!("{} is compressed.", file_name),
            Err((file_name, err)) => println!("{} FAILED due to {}", file_name, err),
        }
    });
    pause();
}

fn process(path: &Path, target_dir: &Path) -> Result<String, (String, String)> {
    let file_name = path.file_name().unwrap().to_string_lossy().to_string();
    let extension = path.extension().unwrap().to_string_lossy().to_lowercase();
    if extension.ne("jpg") && extension.ne("jpeg") {
        return Err((file_name, "not jpeg file".into()))
    }

    let (resized_img_data, target_width, target_height) = match resize(path) {
        Ok(v) => v,
        Err(e) => return Err((file_name, e.to_string())),
    };

    let compressed_img_data = match compress(resized_img_data, target_width, target_height) {
        Ok(v) => v,
        Err(e) => return Err((file_name, e.to_string())),
    };

    let target_file = target_dir.join(&file_name);
    let mut file = BufWriter::new(File::create(target_file)
        .map_err(|e| (file_name.clone(), e.to_string()))?);
    file.write_all(&compressed_img_data)
        .map_err(|e| (file_name.clone(), e.to_string()))?;

    Ok(file_name)
}

fn pause() {
    let mut stdin = io::stdin();
    let mut stdout = io::stdout();

    write!(stdout, "Press any key to continue...").unwrap();
    stdout.flush().unwrap();

    let _ = stdin.read(&mut [0u8]).unwrap();
}

並列化

圧縮を並列に実行して、高速化しましょう。Rust で並列処理なら rayon crate が鉄板。

Cargo.toml にrayon = "1.0"を追加、main.rs にuse rayon::prelude::*;を追加、main関数のsource_files.iter()をsource_files.par_iter()に書き換えるだけ。これだけで CPU コア数分のスレッドを作成して処理を並列に実行してくれます。

まとめ

use std::fs::{self, File};
use std::io::{self, BufWriter, Write, Read};
use std::env;
use std::path::{Path, PathBuf};

use image::{self, GenericImageView, FilterType};
use mozjpeg::{Compress, ColorSpace, ScanMode};
use rayon::prelude::*;

const TARGET_SIZE: usize = 1280;

fn main() {
    let target_dir = match env::args().nth(1) {
        Some(v) => PathBuf::from(v).parent().unwrap().join("compressed"),
        None => return,
    };
    if !target_dir.exists() {
        fs::create_dir(&target_dir).unwrap();
    }

    let source_files = env::args()
        .skip(1)
        .map(PathBuf::from)
        .filter(|p| p.is_file() && p.file_name().is_some() && p.extension().is_some())
        .collect::<Vec<PathBuf>>();

    source_files.par_iter().for_each(|path| {
        match process(&path, &target_dir) {
            Ok(file_name) => println!("{} is compressed.", file_name),
            Err((file_name, err)) => println!("{} FAILED due to {}", file_name, err),
        }
    });
    pause();
}

fn process(path: &Path, target_dir: &Path) -> Result<String, (String, String)> {
    let file_name = path.file_name().unwrap().to_string_lossy().to_string();
    let extension = path.extension().unwrap().to_string_lossy().to_lowercase();
    if extension.ne("jpg") && extension.ne("jpeg") {
        return Err((file_name, "not jpeg file".into()))
    }

    let (resized_img_data, target_width, target_height) = match resize(path) {
        Ok(v) => v,
        Err(e) => return Err((file_name, e.to_string())),
    };

    let compressed_img_data = match compress(resized_img_data, target_width, target_height) {
        Ok(v) => v,
        Err(e) => return Err((file_name, e.to_string())),
    };

    let target_file = target_dir.join(&file_name);
    let mut file = BufWriter::new(File::create(target_file)
        .map_err(|e| (file_name.clone(), e.to_string()))?);
    file.write_all(&compressed_img_data)
        .map_err(|e| (file_name.clone(), e.to_string()))?;

    Ok(file_name)
}

fn resize(path: &Path) -> Result<(Vec<u8>, usize, usize), String> {
    let img = image::open(path).map_err(|e| e.to_string())?;
    let width = img.width() as usize;
    let height = img.height() as usize;

    if width > TARGET_SIZE || height > TARGET_SIZE {
        let (target_width, target_height) = 
            if width > height {
                let ratio: f32 = TARGET_SIZE as f32 / width as f32;
                (TARGET_SIZE, (height as f32 * ratio) as usize)
            } else {
                let ratio: f32 = TARGET_SIZE as f32 / height as f32;
                ((width as f32 * ratio) as usize, TARGET_SIZE)
            };
        let resized_img = img.resize(
                target_width as u32,
                target_height as u32,
                FilterType::Lanczos3);
        Ok((resized_img.to_rgb().to_vec(), resized_img.width() as usize, resized_img.height() as usize))
    } else {
        Ok((img.to_rgb().to_vec(), width, height))
    }
}

fn compress(resized_img_data: Vec<u8>, target_width: usize, target_height: usize) -> Result<Vec<u8>, String> {
    let mut comp = Compress::new(ColorSpace::JCS_RGB);
    comp.set_scan_optimization_mode(ScanMode::AllComponentsTogether);
    comp.set_quality(70.0);

    comp.set_size(target_width, target_height);

    comp.set_mem_dest();
    comp.start_compress();

    let mut line = 0;
    loop {
        if line > target_height - 1 {
            break;
        }
        comp.write_scanlines(&resized_img_data[line * target_width * 3..(line + 1) * target_width * 3]);
        line += 1;
    }
    comp.finish_compress();

    let compressed = comp.data_to_vec()
        .map_err(|_| "data_to_vec failed".to_string())?;
    Ok(compressed)
}

fn pause() {
    let mut stdin = io::stdin();
    let mut stdout = io::stdout();

    write!(stdout, "Press any key to continue...").unwrap();
    stdout.flush().unwrap();

    let _ = stdin.read(&mut [0u8]).unwrap();
}