Replace image crate jpeg decoder with zune-jpeg (#2376)

### Why
It is much faster.

On my Mac, it takes objectron JPEG decoding from 15 to 10 ms. On the web
the difference is smaller: 35ms -> 29ms. See
https://github.com/emilk/image-decode-bench for more.

It can also decode directly to RGBA, which means we don't need to spend
cycles later padding RGB to RGBA. This takes the entire `objectron` demo
`App::update` down from 21 ms to 14 ms.

Sadly, the `gltf` crate still use the `"jpeg"` feature of the `image`
crate, so weneed to pay for the compilation and code bloat of two jpeg
decoders. The `.wasm` of the web-viewer increases by 23kB
(pre-compression) because of this PR.

This PR also adds support for monochrome JPEGs.

### Related
* image-rs/image#1845
* image-rs/image#1876

### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)

<!-- This line will get updated when the PR build summary job finishes.
-->
PR Build Summary: https://build.rerun.io/pr/2376

<!-- pr-link-docs:start -->
Docs preview: https://rerun.io/preview/1aae8d4/docs
Examples preview: https://rerun.io/preview/1aae8d4/examples
<!-- pr-link-docs:end -->

crates/re_components/Cargo.toml

@@ -29,8 +29,8 @@ ecolor = ["dep:ecolor"]
 ## Add support for some math operations using [`glam`](https://crates.io/crates/glam/).
 glam = ["dep:glam"]
 
-## Integration with the [`image`](https://crates.io/crates/image/) crate.
-image = ["dep:ecolor", "dep:image"]
+## Integration with the [`image`](https://crates.io/crates/image/) crate, plus JPEG support..
+image = ["dep:ecolor", "dep:image", "dep:zune-core", "dep:zune-jpeg"]
 
 ## Enable (de)serialization using serde.
 serde = ["dep:serde", "half/serde", "re_log_types/serde"]
@@ -62,11 +62,11 @@ uuid = { version = "1.1", features = ["serde", "v4", "js"] }
 # Optional dependencies:
 ecolor = { workspace = true, optional = true }
 glam = { workspace = true, optional = true }
-image = { workspace = true, optional = true, default-features = false, features = [
-  "jpeg",
-] }
+image = { workspace = true, optional = true, default-features = false }
 rand = { version = "0.8", optional = true }
 serde = { version = "1", optional = true, features = ["derive", "rc"] }
+zune-core = { version = "0.2", optional = true }
+zune-jpeg = { version = "0.3", optional = true }
 
 # Native dependencies:
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]

crates/re_components/src/tensor.rs

@@ -765,20 +765,20 @@ pub enum TensorImageLoadError {
     #[error(transparent)]
     Image(std::sync::Arc<image::ImageError>),
 
-    #[error("Unsupported JPEG color type: {0:?}. Only RGB Jpegs are supported")]
-    UnsupportedJpegColorType(image::ColorType),
+    #[error("Expected a HxW, HxWx1 or HxWx3 tensor, but got {0:?}")]
+    UnexpectedJpegShape(Vec<TensorDimension>),
 
     #[error("Unsupported color type: {0:?}. We support 8-bit, 16-bit, and f32 images, and RGB, RGBA, Luminance, and Luminance-Alpha.")]
     UnsupportedImageColorType(image::ColorType),
 
     #[error("Failed to load file: {0}")]
     ReadError(std::sync::Arc<std::io::Error>),
 
-    #[error("The encoded tensor did not match its metadata {expected:?} != {found:?}")]
-    InvalidMetaData {
-        expected: Vec<TensorDimension>,
-        found: Vec<TensorDimension>,
-    },
+    #[error("The encoded tensor shape did not match its metadata {expected:?} != {found:?}")]
+    InvalidMetaData { expected: Vec<u64>, found: Vec<u64> },
+
+    #[error(transparent)]
+    JpegDecode(#[from] zune_jpeg::errors::DecodeErrors),
 }
 
 #[cfg(feature = "image")]
@@ -899,20 +899,17 @@ impl Tensor {
         }
     }
 
-    /// Construct a tensor from the contents of a JPEG file.
+    /// Construct a tensor from the contents of a JPEG file, without decoding it now.
     ///
     /// Requires the `image` feature.
     pub fn from_jpeg_bytes(jpeg_bytes: Vec<u8>) -> Result<Self, TensorImageLoadError> {
         re_tracing::profile_function!();
-        use image::ImageDecoder as _;
-        let jpeg = image::codecs::jpeg::JpegDecoder::new(std::io::Cursor::new(&jpeg_bytes))?;
-        if jpeg.color_type() != image::ColorType::Rgb8 {
-            // TODO(emilk): support gray-scale jpeg as well
-            return Err(TensorImageLoadError::UnsupportedJpegColorType(
-                jpeg.color_type(),
-            ));
-        }
-        let (w, h) = jpeg.dimensions();
+
+        use zune_jpeg::JpegDecoder;
+
+        let mut decoder = JpegDecoder::new(&jpeg_bytes);
+        decoder.decode_headers()?;
+        let (w, h) = decoder.dimensions().unwrap(); // Can't fail after a successful decode_headers
 
         Ok(Self {
             tensor_id: TensorId::random(),
@@ -1181,8 +1178,6 @@ impl DecodedTensor {
     }
 
     pub fn try_decode(maybe_encoded_tensor: Tensor) -> Result<Self, TensorImageLoadError> {
-        re_tracing::profile_function!();
-
         match &maybe_encoded_tensor.data {
             TensorData::U8(_)
             | TensorData::U16(_)
@@ -1196,26 +1191,85 @@ impl DecodedTensor {
             | TensorData::F32(_)
             | TensorData::F64(_) => Ok(Self(maybe_encoded_tensor)),
 
-            TensorData::JPEG(buf) => {
-                use image::io::Reader as ImageReader;
-                let mut reader = ImageReader::new(std::io::Cursor::new(buf.as_slice()));
-                reader.set_format(image::ImageFormat::Jpeg);
-                let img = {
-                    re_tracing::profile_scope!("decode_jpeg");
-                    reader.decode()?
-                };
-                let decoded_tensor = DecodedTensor::from_image(img)?;
-                if decoded_tensor.shape() == maybe_encoded_tensor.shape() {
-                    Ok(decoded_tensor)
-                } else {
-                    Err(TensorImageLoadError::InvalidMetaData {
-                        expected: maybe_encoded_tensor.shape().into(),
-                        found: decoded_tensor.shape().into(),
-                    })
-                }
+            TensorData::JPEG(jpeg_bytes) => {
+                re_log::debug!(
+                    "Decoding JPEG image of shape {:?}",
+                    maybe_encoded_tensor.shape()
+                );
+
+                let [h, w, c] = maybe_encoded_tensor
+                    .image_height_width_channels()
+                    .ok_or_else(|| {
+                        TensorImageLoadError::UnexpectedJpegShape(
+                            maybe_encoded_tensor.shape().to_vec(),
+                        )
+                    })?;
+
+                Self::decode_jpeg_bytes(jpeg_bytes, [h, w, c])
             }
         }
     }
+
+    pub fn decode_jpeg_bytes(
+        jpeg_bytes: &Buffer<u8>,
+        [expected_height, expected_width, expected_channels]: [u64; 3],
+    ) -> Result<DecodedTensor, TensorImageLoadError> {
+        re_tracing::profile_function!();
+
+        re_log::debug!("Decoding {expected_width}x{expected_height} JPEG");
+
+        use zune_core::colorspace::ColorSpace;
+        use zune_core::options::DecoderOptions;
+        use zune_jpeg::JpegDecoder;
+
+        let mut options = DecoderOptions::default();
+
+        let depth = if expected_channels == 1 {
+            options = options.jpeg_set_out_colorspace(ColorSpace::Luma);
+            1
+        } else {
+            // We decode to RGBA directly so we don't need to pad to four bytes later when uploading to GPU.
+            options = options.jpeg_set_out_colorspace(ColorSpace::RGBA);
+            4
+        };
+
+        let mut decoder = JpegDecoder::new_with_options(options, jpeg_bytes);
+        let pixels = decoder.decode()?;
+        let (w, h) = decoder.dimensions().unwrap(); // Can't fail after a successful decode
+
+        let (w, h) = (w as u64, h as u64);
+
+        if w != expected_width || h != expected_height {
+            return Err(TensorImageLoadError::InvalidMetaData {
+                expected: [expected_height, expected_width, expected_channels].into(),
+                found: [h, w, depth].into(),
+            });
+        }
+
+        if pixels.len() as u64 != w * h * depth {
+            return Err(zune_jpeg::errors::DecodeErrors::Format(format!(
+                "Bug in zune-jpeg: Expected {w}x{h}x{depth}={} bytes, got {}",
+                w * h * depth,
+                pixels.len()
+            ))
+            .into());
+        }
+
+        let tensor = Tensor {
+            tensor_id: TensorId::random(),
+            shape: vec![
+                TensorDimension::height(h),
+                TensorDimension::width(w),
+                TensorDimension::depth(depth),
+            ],
+            data: TensorData::U8(pixels.into()),
+            meaning: TensorDataMeaning::Unknown,
+            meter: None,
+        };
+        let decoded_tensor = DecodedTensor(tensor);
+
+        Ok(decoded_tensor)
+    }
 }
 
 impl AsRef<Tensor> for DecodedTensor {

crates/re_log/src/lib.rs

@@ -40,6 +40,12 @@ pub mod external {
     pub use log;
 }
 
+/// Never log anything less serious than a `ERROR` from these crates.
+const CRATES_AT_ERROR_LEVEL: [&str; 1] = [
+    // Waiting for https://github.com/etemesi254/zune-image/issues/131 to be released
+    "zune_jpeg",
+];
+
 /// Never log anything less serious than a `WARN` from these crates.
 const CRATES_AT_WARN_LEVEL: [&str; 3] = [
     // wgpu crates spam a lot on info level, which is really annoying
@@ -57,6 +63,13 @@ const CRATES_FORCED_TO_INFO: [&str; 4] = [
 
 /// Should we log this message given the filter?
 fn is_log_enabled(filter: log::LevelFilter, metadata: &log::Metadata<'_>) -> bool {
+    if CRATES_AT_ERROR_LEVEL
+        .iter()
+        .any(|crate_name| metadata.target().starts_with(crate_name))
+    {
+        return metadata.level() <= log::LevelFilter::Error;
+    }
+
     if CRATES_AT_WARN_LEVEL
         .iter()
         .any(|crate_name| metadata.target().starts_with(crate_name))

crates/re_log/src/setup.rs

@@ -7,6 +7,11 @@
 fn log_filter() -> String {
     let mut rust_log = std::env::var("RUST_LOG").unwrap_or_else(|_| "info".to_owned());
 
+    for crate_name in crate::CRATES_AT_ERROR_LEVEL {
+        if !rust_log.contains(&format!("{crate_name}=")) {
+            rust_log += &format!(",{crate_name}=error");
+        }
+    }
     for crate_name in crate::CRATES_AT_WARN_LEVEL {
         if !rust_log.contains(&format!("{crate_name}=")) {
             rust_log += &format!(",{crate_name}=warn");

crates/re_sdk/Cargo.toml

@@ -27,7 +27,7 @@ demo = []
 ## Add support for some math operations using [`glam`](https://crates.io/crates/glam/).
 glam = ["re_components/glam"]
 
-## Integration with the [`image`](https://crates.io/crates/image/) crate.
+## Integration with the [`image`](https://crates.io/crates/image/) crate, plus JPEG support..
 image = ["re_components/image"]
 
 

crates/re_viewer/Cargo.toml

@@ -83,10 +83,7 @@ eframe = { workspace = true, default-features = false, features = [
 ] }
 egui.workspace = true
 egui-wgpu.workspace = true
-image = { workspace = true, default-features = false, features = [
-  "jpeg",
-  "png",
-] }
+image = { workspace = true, default-features = false, features = ["png"] }
 itertools = { workspace = true }
 poll-promise = "0.2"
 rfd.workspace = true

crates/re_viewport/Cargo.toml

@@ -40,10 +40,7 @@ egui_tiles.workspace = true
 egui.workspace = true
 enumset.workspace = true
 glam.workspace = true
-image = { workspace = true, default-features = false, features = [
-  "jpeg",
-  "png",
-] }
+image = { workspace = true, default-features = false, features = ["png"] }
 itertools.workspace = true
 nohash-hasher = "0.2"
 serde = "1.0"

crates/rerun/Cargo.toml

@@ -37,7 +37,7 @@ demo = ["re_sdk?/demo"]
 ## Only relevant if feature `sdk` is enabled.
 glam = ["re_sdk?/glam"]
 
-## Integration with the [`image`](https://crates.io/crates/image/) crate.
+## Integration with the [`image`](https://crates.io/crates/image/) crate, plus JPEG support..
 image = ["re_sdk?/image"]
 
 ## Support spawning a native viewer.

rerun_py/Cargo.toml

@@ -56,10 +56,7 @@ re_viewport.workspace = true
 
 arrow2 = { workspace = true, features = ["io_ipc", "io_print"] }
 document-features = "0.2"
-image = { workspace = true, default-features = false, features = [
-  "jpeg",
-  "png",
-] }
+image = { workspace = true, default-features = false, features = ["png"] }
 itertools = { workspace = true }
 mimalloc = { workspace = true, features = ["local_dynamic_tls"] }
 numpy = { version = "0.19.0", features = ["half"] }