Created
June 24, 2025 04:13
-
-
Save cynthia2006/898bbfc5e580ab1c3ace6b1184a6ee4c to your computer and use it in GitHub Desktop.
Fast conversion of pictures into JXL
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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(()) | |
| } |
Author
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Can you understand any of this? No? Well, that's Rust for you.