🦀 (Rust): Support I;16 and F pixel values decoding (#73)

.github/workflows/test.yml

@@ -60,7 +60,7 @@ jobs:
       - name: Test with pytest
         run: |
           source venv/bin/activate
-          pip install -r requirements-dev.txt
+          pip install -e .[dev]
           pytest test/ --junitxml=junit/test-results-${{ matrix.python-version }}.xml
 
       - name: Upload pytest test results

pyproject.toml

@@ -21,10 +21,13 @@ dependencies = [
     "Pillow",
 ]
 
+[project.optional-dependencies]
+dev = ["numpy", "pytest"]
+
 [project.urls]
 "Homepage" = "https://github.com/Isotr0py/pillow-jpegxl-plugin"
 "Bug Tracker" = "https://github.com/Isotr0py/pillow-jpegxl-plugin/issues"
 "Releases" = "https://github.com/Isotr0py/pillow-jpegxl-plugin/releases"
 
 [tool.maturin]
-features = ["pyo3/extension-module"]
+features = ["pyo3/extension-module", "vendored"]

src/decode.rs

@@ -1,6 +1,6 @@
 use std::borrow::Cow;
 
-use pyo3::exceptions::PyRuntimeError;
+use pyo3::exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
 
 use jpegxl_rs::decode::{Data, Metadata, Pixels};
@@ -24,49 +24,44 @@ struct ImageInfo {
 }
 
 impl ImageInfo {
-    fn from(item: Metadata) -> ImageInfo {
+    fn from(item: Metadata, pixel: &Data) -> ImageInfo {
+        let pixel_type = match &pixel {
+            Data::Pixels(pixels) => Some(pixels),
+            Data::Jpeg(_) => None,
+        };
         ImageInfo {
-            mode: Self::mode(item.num_color_channels, item.has_alpha_channel),
+            mode: Self::mode(item.num_color_channels, item.has_alpha_channel, pixel_type).unwrap(),
             width: item.width,
             height: item.height,
             num_channels: item.num_color_channels,
             has_alpha_channel: item.has_alpha_channel,
         }
     }
 
-    fn mode(num_channels: u32, has_alpha_channel: bool) -> String {
-        match (num_channels, has_alpha_channel) {
+    fn mode(
+        num_channels: u32,
+        has_alpha_channel: bool,
+        pixel_type: Option<&Pixels>,
+    ) -> PyResult<String> {
+        let mode = match (num_channels, has_alpha_channel) {
             (1, false) => "L".to_string(),
             (1, true) => "LA".to_string(),
             (3, false) => "RGB".to_string(),
             (3, true) => "RGBA".to_string(),
-            _ => panic!("Unsupported number of channels"),
-        }
-    }
-}
-
-pub fn convert_pixels(pixels: Pixels) -> Vec<u8> {
-    let mut result = Vec::new();
-    match pixels {
-        Pixels::Uint8(pixels) => {
-            for pixel in pixels {
-                result.push(pixel);
-            }
-        }
-        Pixels::Uint16(pixels) => {
-            for pixel in pixels {
-                result.push((pixel >> 8) as u8);
-                result.push(pixel as u8);
+            _ => return Err(PyNotImplementedError::new_err("Unsupported color mode")),
+        };
+        if let Some(Pixels::Uint16(_)) = pixel_type {
+            if mode == "L" {
+                return Ok("I;16".to_string());
             }
         }
-        Pixels::Float(pixels) => {
-            for pixel in pixels {
-                result.push((pixel * 255.0) as u8);
+        if let Some(Pixels::Float(_)) = pixel_type {
+            if mode == "L" {
+                return Ok("F".to_string());
             }
         }
-        Pixels::Float16(_) => panic!("Float16 is not supported yet"),
+        Ok(mode)
     }
-    result
 }
 
 #[pyclass(module = "pillow_jxl")]
@@ -96,6 +91,75 @@ impl Decoder {
     }
 }
 
+impl Decoder {
+    fn pixels_to_bytes_8bit(&self, pixels: Pixels) -> PyResult<Vec<u8>> {
+        // Convert pixels to bytes with 8-bit casting
+        let mut result = Vec::new();
+        match pixels {
+            Pixels::Uint8(pixels) => {
+                return Ok(pixels);
+            }
+            Pixels::Uint16(pixels) => {
+                for pixel in pixels {
+                    result.push((pixel >> 8) as u8);
+                }
+            }
+            Pixels::Float(pixels) => {
+                for pixel in pixels {
+                    result.push((pixel * 255.0) as u8);
+                }
+            }
+            Pixels::Float16(_) => {
+                return Err(PyNotImplementedError::new_err(
+                    "Float16 is not supported yet",
+                ))
+            }
+        }
+        Ok(result)
+    }
+
+    fn pixels_to_bytes(&self, pixels: Pixels) -> PyResult<Vec<u8>> {
+        // Convert pixels to bytes without casting
+        let mut result = Vec::new();
+        match pixels {
+            Pixels::Uint8(pixels) => {
+                return Ok(pixels);
+            }
+            Pixels::Uint16(pixels) => {
+                for pixel in pixels {
+                    let pix_bytes = pixel.to_ne_bytes();
+                    for byte in pix_bytes.iter() {
+                        result.push(*byte);
+                    }
+                }
+            }
+            Pixels::Float(pixels) => {
+                for pixel in pixels {
+                    let pix_bytes = pixel.to_ne_bytes();
+                    for byte in pix_bytes.iter() {
+                        result.push(*byte);
+                    }
+                }
+            }
+            Pixels::Float16(_) => {
+                return Err(PyNotImplementedError::new_err(
+                    "Float16 is not supported yet",
+                ))
+            }
+        }
+        Ok(result)
+    }
+
+    fn convert_pil_pixels(&self, pixels: Pixels, num_channels: u32) -> PyResult<Vec<u8>> {
+        let result = match num_channels {
+            1 => self.pixels_to_bytes(pixels)?,
+            3 => self.pixels_to_bytes_8bit(pixels)?,
+            _ => return Err(PyValueError::new_err("image color channels must be 1 or 3")),
+        };
+        Ok(result)
+    }
+}
+
 impl Decoder {
     fn call_inner(&self, data: &[u8]) -> PyResult<(bool, ImageInfo, Cow<'_, [u8]>, Cow<'_, [u8]>)> {
         let parallel_runner = ThreadsRunner::new(
@@ -113,20 +177,16 @@ impl Decoder {
             .build()
             .map_err(to_pyjxlerror)?;
         let (info, img) = decoder.reconstruct(&data).map_err(to_pyjxlerror)?;
-        let (jpeg, img) = match img {
-            Data::Jpeg(x) => (true, x),
-            Data::Pixels(x) => (false, convert_pixels(x)),
-        };
         let icc_profile: Vec<u8> = match &info.icc_profile {
             Some(x) => x.to_vec(),
             None => Vec::new(),
         };
-        Ok((
-            jpeg,
-            ImageInfo::from(info),
-            Cow::Owned(img),
-            Cow::Owned(icc_profile),
-        ))
+        let img_info = ImageInfo::from(info, &img);
+        let (jpeg, img) = match img {
+            Data::Jpeg(x) => (true, x),
+            Data::Pixels(x) => (false, self.convert_pil_pixels(x, img_info.num_channels)?),
+        };
+        Ok((jpeg, img_info, Cow::Owned(img), Cow::Owned(icc_profile)))
     }
 }
 

test/test_plugin.py

@@ -1,18 +1,44 @@
 import tempfile
 
 import pytest
+import numpy as np
 from PIL import Image
 
 import pillow_jxl
 
 
 def test_decode():
-    img = Image.open("test/images/sample.jxl")
+    img_jxl = Image.open("test/images/sample.jxl")
+    img_png = Image.open("test/images/sample.png")
 
-    assert img.size == (40, 50)
-    assert img.mode == "RGBA"
-    assert not img.is_animated
-    assert img.n_frames == 1
+    assert img_jxl.size == img_png.size
+    assert img_jxl.mode == img_png.mode == "RGBA"
+    assert not img_jxl.is_animated
+    assert img_jxl.n_frames == 1
+    assert np.allclose(np.array(img_jxl), np.array(img_png))
+
+
+def test_decode_I16():
+    img_jxl = Image.open("test/images/sample_grey.jxl")
+    img_png = Image.open("test/images/sample_grey.png")
+
+    assert img_jxl.size == img_png.size
+    assert img_jxl.mode == img_png.mode == "I;16"
+    assert not img_jxl.is_animated
+    assert img_jxl.n_frames == 1
+    # we need to use atol=1 here otherwise the test will fail on MacOS
+    assert np.allclose(np.array(img_jxl), np.array(img_png), atol=1)
+
+
+def test_decode_F():
+    img_jxl = Image.open("test/images/sample_float.jxl")
+    img_ppm = Image.open("test/images/sample_float.ppm")
+
+    assert img_jxl.size == img_ppm.size
+    assert img_jxl.mode == img_ppm.mode == "F"
+    assert not img_jxl.is_animated
+    assert img_jxl.n_frames == 1
+    assert np.allclose(np.array(img_jxl), np.array(img_ppm), atol=3e-2)
 
 
 @pytest.mark.parametrize("image", ["test/images/sample.png", "test/images/sample.jpg"])