Fast conversion of pictures into JXL

xl-convert.rs

use std::{
    collections::VecDeque,
    error::Error,
    fs::File,
    path::PathBuf,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc, Mutex,
    },
};

use clap::Parser;

use jpegxl_rs::{
    self as jxl,
    encode::{EncoderFrame, EncoderResult},
};
use jpegxl_sys::encoder::encode::JxlEncoderFrameSettingId;
use png::{self, ColorType, DecodeOptions, Transformations};
use turbojpeg as tj;

use threadpool::ThreadPool;

#[derive(Parser)]
struct Opts {
    indir: PathBuf,
    outdir: PathBuf,
}

struct EncodeJob {
    frame: Vec<u8>,
    width: u32,
    height: u32,
    channels: u8,
    // Mainly for better debug messages.
    in_file: PathBuf,
    out_file: PathBuf,
}

const N_CONCURRENT: usize = 8;
const N_JOBS: usize = 16;
const JOB_QUEUE: usize = 32;

fn main() -> Result<(), Box<dyn Error>> {
    let opts = Opts::parse();

    // TODO Make errors more user-friendly.
    assert!(opts.indir.try_exists()?);
    assert!(opts.outdir.try_exists()?);

    let jobs = Arc::new(Mutex::new(VecDeque::<EncodeJob>::with_capacity(JOB_QUEUE)));
    let flush_mode = Arc::new(AtomicBool::new(false));
    let encoder_pool = ThreadPool::new(N_CONCURRENT);

    for _ in 0..N_JOBS {
        let jobs_clone = Arc::clone(&jobs);
        let flush_mode_clone = Arc::clone(&flush_mode);

        encoder_pool.execute(move || {
            let runner = jxl::ThreadsRunner::default();
            let mut jxl_encoder = jxl::encoder_builder()
                .use_container(false)
                .parallel_runner(&runner)
                .build()
                .unwrap();

            jxl_encoder
                .set_frame_option(JxlEncoderFrameSettingId::Modular, 1)
                .unwrap();

            loop {
                let mut jobs_clone = jobs_clone.lock().unwrap();
                let job = jobs_clone.pop_front();
                let all_jobs_finished = jobs_clone.is_empty();
                drop(jobs_clone);

                match job {
                    Some(job) => {
                        jxl_encoder.has_alpha = job.channels == 4;

                        let result: EncoderResult<u8> = match jxl_encoder
                            .encode_frame(
                                &EncoderFrame::new(&job.frame).num_channels(job.channels as u32),
                                job.width as u32,
                                job.height as u32,
                            ) {
                                Ok(value) => value,
                                Err(err) => {
                                    eprintln!("Couldn't encode {} into JXL because of {}",
                                              job.in_file.to_str().unwrap(), err);
                                    continue;
                                }
                            };

                        std::fs::write(&job.out_file, result.data).unwrap();
                        println!("Finished {}", job.out_file.to_str().unwrap());
                    }
                    None => {}
                }

                // All of our data had been enqueued, now finish encoding all things.
                if flush_mode_clone.load(Ordering::Relaxed) && all_jobs_finished {
                    break;
                }
            }
        });
    }

    let mut jpeg_decoder = tj::Decompressor::new()?;
    let mut in_dir_reader = std::fs::read_dir(opts.indir)?;

    loop {
        let jobs_lock = jobs.lock().unwrap();
        if jobs_lock.len() >= JOB_QUEUE {
            continue;
        }
        drop(jobs_lock);

        let in_file = match in_dir_reader.next() {
            Some(elem) => elem,
            None => break
        }?.path();

        let mut out_file = opts.outdir.clone();
        out_file.push(in_file.file_name().unwrap());
        out_file.set_extension("jxl");

        if out_file.exists() {
            println!("Skipping {} as it already exists at {}",
                     in_file.to_str().unwrap(), out_file.to_str().unwrap());
            continue;
        }

        match in_file.extension().unwrap().to_str().unwrap() {
            "jpg" => {
                let jpeg_data = std::fs::read(&in_file)?;
                let jpeg_info = match jpeg_decoder.read_header(&jpeg_data) {
                    Ok(value) => value,
                    Err(err) => {
                        eprintln!("Couldn't read JPEG header for {} because of {}",
                                  in_file.to_str().unwrap(), err);
                        continue;
                    }
                };

                let mut image = tj::Image {
                    format: tj::PixelFormat::RGB,
                    width: jpeg_info.width,
                    height: jpeg_info.height,
                    pitch: jpeg_info.width * 3,
                    pixels: vec![0; jpeg_info.width * jpeg_info.height * 3],
                };

                if let Err(err) = jpeg_decoder.decompress(&jpeg_data, image.as_deref_mut()) {
                    eprintln!("Couldn't decode JPEG data for {} because of {}",
                              in_file.to_str().unwrap(), err);
                    continue;
                }

                println!("Enqueuing {} for encode", in_file.to_str().unwrap());
                let mut jobs = jobs.lock().unwrap();
                jobs.push_front(
                    EncodeJob {
                        in_file,
                        out_file,
                        frame: image.pixels,
                        width: image.width as u32,
                        height: image.height as u32,
                        channels: 3,
                    }
                );
                drop(jobs);
            }
            "png" => {
                let mut decode_opts = DecodeOptions::default();
                decode_opts.set_ignore_checksums(true);

                let mut png_decoder =
                    png::Decoder::new_with_options(File::open(&in_file)?, decode_opts);
                png_decoder.set_transformations(Transformations::STRIP_16 | Transformations::EXPAND);

                let mut reader = match png_decoder.read_info() {
                    Ok(value) => value,
                    Err(err) => {
                        eprintln!("Couldn't read PNG header for {} because of {}",
                                  in_file.to_str().unwrap(), err);
                        continue;
                    }
                };

                let mut buf = vec![0; reader.output_buffer_size()];
                let info = match reader.next_frame(&mut buf) {
                    Ok(value) => value,
                    Err(err) => {
                        eprintln!("Couldn't decode PNG data for {} because of {}",
                                  in_file.to_str().unwrap(), err);
                        continue;
                    }
                };

                let num_channels = if info.color_type == ColorType::Rgba {
                    4
                } else {
                    3
                };

                buf.shrink_to(info.buffer_size());

                println!("Enqueuing {} for encode", in_file.to_str().unwrap());
                let mut jobs = jobs.lock().unwrap();
                jobs.push_front(
                    EncodeJob {
                        in_file,
                        out_file,
                        frame: buf,
                        width: info.width,
                        height: info.height,
                        channels: num_channels,
                    }
                );
                drop(jobs);
            }
            _ => {
                println!("Ignored {}", in_file.to_str().unwrap());
            }
        }
    }

    flush_mode.store(true, Ordering::Relaxed);
    encoder_pool.join();

    Ok(())
}