🚀 (Encode): add effort argument (#24)

🚀 (Encode): add effort argument
Isotr0py · Feb 26, 2024 · 2c7f025 · 2c7f025
2 parents b968267 + b7abd4b
commit 2c7f025
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git a/pillow_jxl/JpegXLImagePlugin.py b/pillow_jxl/JpegXLImagePlugin.py
@@ -80,13 +80,15 @@ def _save(im, fp, filename, save_all=False):
     lossless = info.get("lossless", False)
     quality = 0 if lossless else 1
     decoding_speed = info.get("decoding_speed", 0)
+    effort = info.get("effort", 7)
     use_container = info.get("use_container", True)
 
     enc = Encoder(
         mode=im.mode,
         lossless=lossless,
         quality=quality,
         decoding_speed=decoding_speed,
+        effort=effort,
         use_container=use_container,
     )
     # FIXME (Isotr0py): im.filename maybe None if parse stream

diff --git a/src/lib.rs b/src/lib.rs
@@ -2,7 +2,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 
 use jpegxl_rs::decode::{Data, Metadata, Pixels};
-use jpegxl_rs::encode::{ColorEncoding, EncoderFrame, EncoderResult};
+use jpegxl_rs::encode::{ColorEncoding, EncoderSpeed, EncoderFrame, EncoderResult};
 use jpegxl_rs::parallel::threads_runner::ThreadsRunner;
 use jpegxl_rs::{decoder_builder, encoder_builder};
 // it works even if the item is not documented:
@@ -15,20 +15,22 @@ struct Encoder {
     lossless: bool,
     quality: f32,
     decoding_speed: i64,
+    effort: u32,
     use_container: bool,
     use_original_profile: bool,
 }
 
 #[pymethods]
 impl Encoder {
     #[new]
-    #[pyo3(signature = (mode, parallel=true, lossless=false, quality=1.0, decoding_speed=0, use_container=true, use_original_profile=false))]
+    #[pyo3(signature = (mode, parallel=true, lossless=false, quality=1.0, decoding_speed=0, effort=7, use_container=true, use_original_profile=false))]
     fn new(
         mode: &str,
         parallel: bool,
         lossless: bool,
         quality: f32,
         decoding_speed: i64,
+        effort: u32,
         use_container: bool,
         use_original_profile: bool,
     ) -> Self {
@@ -52,6 +54,7 @@ impl Encoder {
                 0...4 => decoding_speed,
                 _ => panic!("Decoding speed must be between 0 and 4"),
             },
+            effort: effort,
             use_container: use_container,
             use_original_profile: match lossless {
                 true => true,
@@ -91,6 +94,18 @@ impl Encoder {
             3 | 4 => ColorEncoding::Srgb,
             _ => panic!("Invalid num channels"),
         };
+        encoder.speed = match self.effort {
+            1 => EncoderSpeed::Lightning,
+            2 => EncoderSpeed::Thunder,
+            3 => EncoderSpeed::Falcon,
+            4 => EncoderSpeed::Cheetah,
+            5 => EncoderSpeed::Hare,
+            6 => EncoderSpeed::Wombat,
+            7 => EncoderSpeed::Squirrel,
+            8 => EncoderSpeed::Kitten,
+            9 => EncoderSpeed::Tortoise,
+            _ => panic!("Invalid effort"),
+        };
         let buffer: EncoderResult<u8> = match jpeg_encode {
             true => encoder.encode_jpeg(&data).unwrap(),
             false => {
@@ -103,8 +118,8 @@ impl Encoder {
 
     fn __repr__(&self) -> PyResult<String> {
         Ok(format!(
-            "Encoder(parallel={}, has_alpha={}, lossless={}, quality={}, decoding_speed={})",
-            self.parallel, self.has_alpha, self.lossless, self.quality, self.decoding_speed
+            "Encoder(parallel={}, has_alpha={}, lossless={}, quality={}, decoding_speed={}, effort={})",
+            self.parallel, self.has_alpha, self.lossless, self.quality, self.decoding_speed, self.effort
         ))
     }
 }