Add Tensor::from_image_file and Tensor::from_image_bytes (#2097)

* Add Tensor::from_image_file and Tensor::from_image_bytes

* Support PNG in log_image_file

* remove duplicated cfg

Co-authored-by: Andreas Reich <andreas@rerun.io>

* Guess format even if extension is unknown

* Always turn on the "image" feature of re_log_types

* Python SDK: always support PNG images

---------

Co-authored-by: Andreas Reich <andreas@rerun.io>

crates/re_log_types/src/component_types/tensor.rs

@@ -776,6 +776,13 @@ pub enum TensorImageLoadError {
         expected: Vec<TensorDimension>,
         found: Vec<TensorDimension>,
     },
+
+    #[cfg(not(target_arch = "wasm32"))]
+    #[error("Unsupported file extension '{extension}' for file {path:?}")]
+    UnknownExtension {
+        extension: String,
+        path: std::path::PathBuf,
+    },
 }
 
 #[cfg(feature = "image")]
@@ -828,21 +835,79 @@ impl Tensor {
 
 #[cfg(feature = "image")]
 impl Tensor {
+    /// Construct a tensor from the contents of an image file on disk.
+    ///
+    /// JPEGs will be kept encoded, left to the viewer to decode on-the-fly.
+    /// Other images types will be decoded directly.
+    ///
+    /// Requires the `image` feature.
+    #[cfg(not(target_arch = "wasm32"))]
+    pub fn from_image_file(path: &std::path::Path) -> Result<Self, TensorImageLoadError> {
+        crate::profile_function!(path.to_string_lossy());
+
+        let img_bytes = {
+            crate::profile_scope!("fs::read");
+            std::fs::read(path)?
+        };
+
+        let img_format = if let Some(extension) = path.extension() {
+            if let Some(format) = image::ImageFormat::from_extension(extension) {
+                format
+            } else {
+                image::guess_format(&img_bytes)?
+            }
+        } else {
+            image::guess_format(&img_bytes)?
+        };
+
+        Self::from_image_bytes(img_bytes, img_format)
+    }
+
     /// Construct a tensor from the contents of a JPEG file on disk.
     ///
     /// Requires the `image` feature.
     #[cfg(not(target_arch = "wasm32"))]
+    pub fn from_jpeg_file(path: &std::path::Path) -> Result<Self, TensorImageLoadError> {
+        crate::profile_function!(path.to_string_lossy());
+        let jpeg_bytes = {
+            crate::profile_scope!("fs::read");
+            std::fs::read(path)?
+        };
+        Self::from_jpeg_bytes(jpeg_bytes)
+    }
+
+    #[deprecated = "Renamed 'from_jpeg_file'"]
+    #[cfg(not(target_arch = "wasm32"))]
     pub fn tensor_from_jpeg_file(
         image_path: impl AsRef<std::path::Path>,
     ) -> Result<Self, TensorImageLoadError> {
-        let jpeg_bytes = std::fs::read(image_path)?;
-        Self::tensor_from_jpeg_bytes(jpeg_bytes)
+        Self::from_jpeg_file(image_path.as_ref())
+    }
+
+    /// Construct a tensor from the contents of an image file.
+    ///
+    /// JPEGs will be kept encoded, left to the viewer to decode on-the-fly.
+    /// Other images types will be decoded directly.
+    ///
+    /// Requires the `image` feature.
+    pub fn from_image_bytes(
+        bytes: Vec<u8>,
+        format: image::ImageFormat,
+    ) -> Result<Self, TensorImageLoadError> {
+        crate::profile_function!(format!("{format:?}"));
+        if format == image::ImageFormat::Jpeg {
+            Self::from_jpeg_bytes(bytes)
+        } else {
+            let image = image::load_from_memory_with_format(&bytes, format)?;
+            Self::from_image(image)
+        }
     }
 
     /// Construct a tensor from the contents of a JPEG file.
     ///
     /// Requires the `image` feature.
-    pub fn tensor_from_jpeg_bytes(jpeg_bytes: Vec<u8>) -> Result<Self, TensorImageLoadError> {
+    pub fn from_jpeg_bytes(jpeg_bytes: Vec<u8>) -> Result<Self, TensorImageLoadError> {
+        crate::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 {
@@ -866,6 +931,12 @@ impl Tensor {
         })
     }
 
+    #[deprecated = "Renamed 'from_jpeg_bytes'"]
+    #[cfg(not(target_arch = "wasm32"))]
+    pub fn tensor_from_jpeg_bytes(jpeg_bytes: Vec<u8>) -> Result<Self, TensorImageLoadError> {
+        Self::from_jpeg_bytes(jpeg_bytes)
+    }
+
     /// Construct a tensor from something that can be turned into a [`image::DynamicImage`].
     ///
     /// Requires the `image` feature.
@@ -1059,6 +1130,8 @@ impl DecodedTensor {
     pub fn from_dynamic_image(
         image: image::DynamicImage,
     ) -> Result<DecodedTensor, TensorImageLoadError> {
+        crate::profile_function!();
+
         let (w, h) = (image.width(), image.height());
 
         let (depth, data) = match image {

examples/rust/objectron/src/main.rs

@@ -154,7 +154,7 @@ fn log_baseline_objects(
 
 fn log_video_frame(rec_stream: &RecordingStream, ar_frame: &ArFrame) -> anyhow::Result<()> {
     let image_path = ar_frame.dir.join(format!("video/{}.jpg", ar_frame.index));
-    let tensor = rerun::components::Tensor::tensor_from_jpeg_file(image_path)?;
+    let tensor = rerun::components::Tensor::from_jpeg_file(&image_path)?;
 
     MsgSender::new("world/camera/video")
         .with_timepoint(ar_frame.timepoint.clone())

rerun_py/Cargo.toml

@@ -42,7 +42,7 @@ web_viewer = ["rerun/web_viewer", "dep:re_web_viewer_server", "dep:re_ws_comms"]
 re_build_info.workspace = true
 re_error.workspace = true
 re_log.workspace = true
-re_log_types = { workspace = true, features = ["glam"] }
+re_log_types = { workspace = true, features = ["glam", "image"] }
 re_memory.workspace = true
 rerun = { workspace = true, features = ["analytics", "server", "sdk"] }
 re_web_viewer_server = { workspace = true, optional = true }
@@ -51,7 +51,10 @@ re_ws_comms = { workspace = true, optional = true }
 arrow2 = { workspace = true, features = ["io_ipc", "io_print"] }
 document-features = "0.2"
 glam.workspace = true
-image = { workspace = true, default-features = false, features = ["jpeg"] }
+image = { workspace = true, default-features = false, features = [
+  "jpeg",
+  "png",
+] }
 itertools = { workspace = true }
 macaw.workspace = true
 mimalloc = { workspace = true, features = ["local_dynamic_tls"] }

rerun_py/rerun_sdk/rerun/log/file.py

@@ -39,6 +39,9 @@ class ImageFormat(Enum):
     JPEG = "jpeg"
     """JPEG format."""
 
+    PNG = "png"
+    """PNG format."""
+
 
 @log_decorator
 def log_mesh_file(
@@ -112,11 +115,15 @@ def log_image_file(
     """
     Log an image file given its contents or path on disk.
 
-    Only JPEGs are supported right now.
-
     You must pass either `img_bytes` or `img_path`.
 
-    If no `img_format` is specified, we will try and guess it.
+    Only JPEGs and PNGs are supported right now.
+
+    JPEGs will be stored compressed, saving memory,
+    whilst PNGs will currently be decoded before they are logged.
+    This may change in the future.
+
+    If no `img_format` is specified, rerun will try to guess it.
 
     Parameters
     ----------

rerun_py/src/python_bridge.rs

@@ -2,7 +2,7 @@
 #![allow(clippy::borrow_deref_ref)] // False positive due to #[pufunction] macro
 #![allow(unsafe_op_in_unsafe_fn)] // False positive due to #[pufunction] macro
 
-use std::{borrow::Cow, io::Cursor, path::PathBuf};
+use std::{borrow::Cow, path::PathBuf};
 
 use itertools::izip;
 use pyo3::{
@@ -1000,46 +1000,11 @@ fn log_image_file(
         }
     };
 
-    use image::ImageDecoder as _;
-    let (w, h) = match img_format {
-        image::ImageFormat::Jpeg => {
-            use image::codecs::jpeg::JpegDecoder;
-            let jpeg = JpegDecoder::new(Cursor::new(&img_bytes))
-                .map_err(|err| PyTypeError::new_err(err.to_string()))?;
-
-            let color_format = jpeg.color_type();
-            if !matches!(color_format, image::ColorType::Rgb8) {
-                // TODO(emilk): support gray-scale jpeg aswell
-                return Err(PyTypeError::new_err(format!(
-                    "Unsupported color format {color_format:?}. \
-                    Expected one of: RGB8"
-                )));
-            }
-
-            jpeg.dimensions()
-        }
-        _ => {
-            return Err(PyTypeError::new_err(format!(
-                "Unsupported image format {img_format:?}. \
-                Expected one of: JPEG"
-            )))
-        }
-    };
+    let tensor = Tensor::from_image_bytes(img_bytes, img_format)
+        .map_err(|err| PyTypeError::new_err(err.to_string()))?;
 
     let time_point = time(timeless, data_stream);
 
-    let tensor = re_log_types::component_types::Tensor {
-        tensor_id: TensorId::random(),
-        shape: vec![
-            TensorDimension::height(h as _),
-            TensorDimension::width(w as _),
-            TensorDimension::depth(3),
-        ],
-        data: re_log_types::component_types::TensorData::JPEG(img_bytes.into()),
-        meaning: re_log_types::component_types::TensorDataMeaning::Unknown,
-        meter: None,
-    };
-
     let row = DataRow::from_cells1(
         RowId::random(),
         entity_path,