From 041b47401ba09276e3a724f0bfc1e0dbf707ae1c Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com> Date: Thu, 23 Mar 2023 12:46:08 +0100 Subject: [PATCH 1/5] Refactor how tensor image cache turns tensors into images --- .../src/component_types/tensor.rs | 62 ++-- .../src/misc/caches/tensor_image_cache.rs | 274 +++++++++++------- 2 files changed, 214 insertions(+), 122 deletions(-) diff --git a/crates/re_log_types/src/component_types/tensor.rs b/crates/re_log_types/src/component_types/tensor.rs index 4b678232d7d2..502ab2b73c4c 100644 --- a/crates/re_log_types/src/component_types/tensor.rs +++ b/crates/re_log_types/src/component_types/tensor.rs @@ -174,6 +174,42 @@ pub enum TensorData { JPEG(BinaryBuffer), } +impl TensorData { + pub fn dtype(&self) -> TensorDataType { + match self { + Self::U8(_) | Self::JPEG(_) => TensorDataType::U8, + Self::U16(_) => TensorDataType::U16, + Self::U32(_) => TensorDataType::U32, + Self::U64(_) => TensorDataType::U64, + Self::I8(_) => TensorDataType::I8, + Self::I16(_) => TensorDataType::I16, + Self::I32(_) => TensorDataType::I32, + Self::I64(_) => TensorDataType::I64, + Self::F32(_) => TensorDataType::F32, + Self::F64(_) => TensorDataType::F64, + } + } + + pub fn size_in_bytes(&self) -> usize { + match self { + Self::U8(buf) | Self::JPEG(buf) => buf.0.len(), + Self::U16(buf) => buf.len(), + Self::U32(buf) => buf.len(), + Self::U64(buf) => buf.len(), + Self::I8(buf) => buf.len(), + Self::I16(buf) => buf.len(), + Self::I32(buf) => buf.len(), + Self::I64(buf) => buf.len(), + Self::F32(buf) => buf.len(), + Self::F64(buf) => buf.len(), + } + } + + pub fn is_empty(&self) -> bool { + self.size_in_bytes() == 0 + } +} + /// Flattened `Tensor` data payload /// /// ## Examples @@ -394,33 +430,11 @@ impl TensorTrait for Tensor { } fn dtype(&self) -> TensorDataType { - match &self.data { - TensorData::U8(_) | TensorData::JPEG(_) => TensorDataType::U8, - TensorData::U16(_) => TensorDataType::U16, - TensorData::U32(_) => TensorDataType::U32, - TensorData::U64(_) => TensorDataType::U64, - TensorData::I8(_) => TensorDataType::I8, - TensorData::I16(_) => TensorDataType::I16, - TensorData::I32(_) => TensorDataType::I32, - TensorData::I64(_) => TensorDataType::I64, - TensorData::F32(_) => TensorDataType::F32, - TensorData::F64(_) => TensorDataType::F64, - } + self.data.dtype() } fn size_in_bytes(&self) -> usize { - match &self.data { - TensorData::U8(buf) | TensorData::JPEG(buf) => buf.0.len(), - TensorData::U16(buf) => buf.len(), - TensorData::U32(buf) => buf.len(), - TensorData::U64(buf) => buf.len(), - TensorData::I8(buf) => buf.len(), - TensorData::I16(buf) => buf.len(), - TensorData::I32(buf) => buf.len(), - TensorData::I64(buf) => buf.len(), - TensorData::F32(buf) => buf.len(), - TensorData::F64(buf) => buf.len(), - } + self.data.size_in_bytes() } } diff --git a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs index fba14a2af73e..3a99774256c6 100644 --- a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs +++ b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs @@ -239,12 +239,20 @@ impl CachedImage { } fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::Result<ColorImage> { + match tensor.meaning { + TensorDataMeaning::Unknown => color_tensor_as_color_image(tensor), + TensorDataMeaning::ClassId => class_id_tensor_as_color_image(tensor, annotations), + TensorDataMeaning::Depth => depth_tensor_as_color_image(tensor), + } +} + +fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { use anyhow::Context as _; crate::profile_function!(format!( - "dtype: {}, meaning: {:?}", + "dtype: {}, shape: {:?}", tensor.dtype(), - tensor.meaning + tensor.shape() )); let shape = &tensor.shape(); @@ -274,43 +282,8 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R let size = [width as _, height as _]; - match (depth, &tensor.data, tensor.meaning) { - (1, TensorData::U8(buf), TensorDataMeaning::ClassId) => { - // Apply annotation mapping to raw bytes interpreted as u8 - let color_lookup: Vec<Color32> = (0..256) - .map(|id| { - annotations - .class_description(Some(ClassId(id))) - .annotation_info() - .color(None, DefaultColor::TransparentBlack) - }) - .collect(); - let pixels: Vec<Color32> = buf - .0 - .iter() - .map(|p: &u8| color_lookup[*p as usize]) - .collect(); - crate::profile_scope!("from_raw"); - Ok(ColorImage { size, pixels }) - } - (1, TensorData::U16(buf), TensorDataMeaning::ClassId) => { - // Apply annotations mapping to bytes interpreted as u16 - let mut color_lookup: ahash::HashMap<u16, Color32> = Default::default(); - let pixels = buf - .iter() - .map(|id: &u16| { - *color_lookup.entry(*id).or_insert_with(|| { - annotations - .class_description(Some(ClassId(*id))) - .annotation_info() - .color(None, DefaultColor::TransparentBlack) - }) - }) - .collect(); - crate::profile_scope!("from_raw"); - Ok(ColorImage { size, pixels }) - } - (1, TensorData::U8(buf), _) => { + match (depth, &tensor.data) { + (1, TensorData::U8(buf)) => { // TODO(emilk): we should read some meta-data to check if this is luminance or alpha. let pixels = buf .0 @@ -319,7 +292,7 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R .collect(); Ok(ColorImage { size, pixels }) } - (1, TensorData::U16(buf), _) => { + (1, TensorData::U16(buf)) => { // TODO(emilk): we should read some meta-data to check if this is luminance or alpha. let pixels = buf .iter() @@ -328,50 +301,7 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R Ok(ColorImage { size, pixels }) } - (1, TensorData::F32(buf), TensorDataMeaning::Depth) => { - if buf.is_empty() { - Ok(ColorImage::default()) - } else { - // Convert to u16 so we can put them in an image. - // TODO(emilk): Eventually we want a renderer that can show f32 images natively. - // One big downside of the approach below is that if we have two depth images - // in the same range, they cannot be visually compared with each other, - // because their individual max-depths will be scaled to 65535. - - let mut min = f32::INFINITY; - let mut max = f32::NEG_INFINITY; - for float in buf.iter() { - min = min.min(*float); - max = max.max(*float); - } - - anyhow::ensure!( - min.is_finite() && max.is_finite(), - "Depth image had non-finite values" - ); - - let ints: Vec<u16> = if min == max { - // Uniform image. We can't remap it to a 0-1 range, so do whatever: - buf.iter().map(|&float| float as u16).collect() - } else { - buf.iter() - .map(|&float| egui::remap(float, min..=max, 0.0..=65535.0) as u16) - .collect() - }; - - let pixels = ints - .iter() - .map(|pixel| { - let [r, g, b, _] = - re_renderer::colormap_turbo_srgb((*pixel as f32) / (u16::MAX as f32)); - egui::Color32::from_rgb(r, g, b) - }) - .collect(); - - Ok(ColorImage { size, pixels }) - } - } - (1, TensorData::F32(buf), _) => { + (1, TensorData::F32(buf)) => { let pixels = buf .iter() .map(|pixel| Color32::from_gray(linear_u8_from_linear_f32(*pixel))) @@ -379,13 +309,13 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R Ok(ColorImage { size, pixels }) } - (3, TensorData::U8(buf), _) => Ok(ColorImage::from_rgb(size, buf.0.as_slice())), - (3, TensorData::U16(buf), _) => { + (3, TensorData::U8(buf)) => Ok(ColorImage::from_rgb(size, buf.0.as_slice())), + (3, TensorData::U16(buf)) => { let u8_buf: Vec<u8> = buf.iter().map(|pixel| (*pixel / 256) as u8).collect(); Ok(ColorImage::from_rgb(size, &u8_buf)) } - (3, TensorData::F32(buf), _) => { + (3, TensorData::F32(buf)) => { let rgb: &[[f32; 3]] = bytemuck::cast_slice(buf.as_slice()); let pixels: Vec<Color32> = rgb .iter() @@ -400,15 +330,13 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R Ok(ColorImage { size, pixels }) } - (4, TensorData::U8(buf), _) => { - Ok(ColorImage::from_rgba_unmultiplied(size, buf.0.as_slice())) - } - (4, TensorData::U16(buf), _) => { + (4, TensorData::U8(buf)) => Ok(ColorImage::from_rgba_unmultiplied(size, buf.0.as_slice())), + (4, TensorData::U16(buf)) => { let u8_buf: Vec<u8> = buf.iter().map(|pixel| (*pixel / 256) as u8).collect(); Ok(ColorImage::from_rgba_unmultiplied(size, &u8_buf)) } - (4, TensorData::F32(buf), _) => { + (4, TensorData::F32(buf)) => { let rgba: &[[f32; 4]] = bytemuck::cast_slice(buf.as_slice()); let pixels: Vec<Color32> = rgba .iter() @@ -423,18 +351,168 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R Ok(ColorImage { size, pixels }) } - (_, TensorData::JPEG(_), _) => { - anyhow::bail!("JPEG tensor should have been decoded before using TensorImageCache") + + (_depth, dtype) => { + anyhow::bail!("Don't know how to turn a tensor of shape={shape:?} and dtype={dtype:?} into a color image") } + } +} + +fn class_id_tensor_as_color_image( + tensor: &Tensor, + annotations: &Annotations, +) -> anyhow::Result<ColorImage> { + use anyhow::Context as _; + + crate::profile_function!(format!( + "dtype: {}, shape: {:?}", + tensor.dtype(), + tensor.shape() + )); + + let shape = &tensor.shape(); - (_depth, dtype, meaning @ TensorDataMeaning::ClassId) => { + anyhow::ensure!( + shape.len() == 2 || shape.len() == 3, + "Expected a 2D or 3D tensor, got {shape:?}", + ); + + let [height, width] = [ + u32::try_from(shape[0].size).context("tensor too large")?, + u32::try_from(shape[1].size).context("tensor too large")?, + ]; + let depth = if shape.len() == 2 { 1 } else { shape[2].size }; + + anyhow::ensure!( + depth == 1, + "Cannot apply annotations to tensor of shape {shape:?}" + ); + debug_assert!( + tensor.is_shaped_like_an_image(), + "We should make the same checks above, but with actual error messages" + ); + let size = [width as _, height as _]; + + match &tensor.data { + TensorData::U8(buf) => { + // Apply annotation mapping to raw bytes interpreted as u8 + let color_lookup: Vec<Color32> = (0..256) + .map(|id| { + annotations + .class_description(Some(ClassId(id))) + .annotation_info() + .color(None, DefaultColor::TransparentBlack) + }) + .collect(); + let pixels: Vec<Color32> = buf + .0 + .iter() + .map(|p: &u8| color_lookup[*p as usize]) + .collect(); + crate::profile_scope!("from_raw"); + Ok(ColorImage { size, pixels }) + } + TensorData::U16(buf) => { + // Apply annotations mapping to bytes interpreted as u16 + let mut color_lookup: ahash::HashMap<u16, Color32> = Default::default(); + let pixels = buf + .iter() + .map(|id: &u16| { + *color_lookup.entry(*id).or_insert_with(|| { + annotations + .class_description(Some(ClassId(*id))) + .annotation_info() + .color(None, DefaultColor::TransparentBlack) + }) + }) + .collect(); + crate::profile_scope!("from_raw"); + Ok(ColorImage { size, pixels }) + } + _ => { anyhow::bail!( - "Shape={shape:?} and dtype={dtype:?} is incompatible with meaning={meaning:?}" + "Cannot apply annotations to tensor of dtype {}", + tensor.dtype() ) } + } +} + +fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { + if tensor.data.is_empty() { + return Ok(ColorImage::default()); + } + + use anyhow::Context as _; + + crate::profile_function!(format!( + "dtype: {}, shape: {:?}", + tensor.dtype(), + tensor.shape() + )); - (_depth, dtype, _) => { - anyhow::bail!("Don't know how to turn a tensor of shape={shape:?} and dtype={dtype:?} into an image") + let shape = &tensor.shape(); + + anyhow::ensure!( + shape.len() == 2 || shape.len() == 3, + "Expected a 2D or 3D tensor, got {shape:?}", + ); + + let [height, width] = [ + u32::try_from(shape[0].size).context("tensor too large")?, + u32::try_from(shape[1].size).context("tensor too large")?, + ]; + let depth = if shape.len() == 2 { 1 } else { shape[2].size }; + + anyhow::ensure!(depth == 1, "Depth tensor of shape {shape:?}"); + debug_assert!( + tensor.is_shaped_like_an_image(), + "We should make the same checks above, but with actual error messages" + ); + let size = [width as _, height as _]; + + match &tensor.data { + TensorData::F32(buf) => { + // Convert to u16 so we can put them in an image. + // TODO(emilk): Eventually we want a renderer that can show f32 images natively. + // One big downside of the approach below is that if we have two depth images + // in the same range, they cannot be visually compared with each other, + // because their individual max-depths will be scaled to 65535. + + let mut min = f32::INFINITY; + let mut max = f32::NEG_INFINITY; + for float in buf.iter() { + min = min.min(*float); + max = max.max(*float); + } + + anyhow::ensure!( + min.is_finite() && max.is_finite(), + "Depth image had non-finite values" + ); + + let ints: Vec<u16> = if min == max { + // Uniform image. We can't remap it to a 0-1 range, so do whatever: + buf.iter().map(|&float| float as u16).collect() + } else { + buf.iter() + .map(|&float| egui::remap(float, min..=max, 0.0..=65535.0) as u16) + .collect() + }; + + let pixels = ints + .iter() + .map(|pixel| { + let [r, g, b, _] = + re_renderer::colormap_turbo_srgb((*pixel as f32) / (u16::MAX as f32)); + egui::Color32::from_rgb(r, g, b) + }) + .collect(); + + Ok(ColorImage { size, pixels }) + } + _ => { + color_tensor_as_color_image(tensor) // TODO(emilk): support more depth dtypes } } } From 5fe43c7d2778b77d7a07cde3cb0cd4734b7e0403 Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com> Date: Thu, 23 Mar 2023 12:53:32 +0100 Subject: [PATCH 2/5] simplify the code --- .../src/misc/caches/tensor_image_cache.rs | 70 ++++++------------- 1 file changed, 20 insertions(+), 50 deletions(-) diff --git a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs index 3a99774256c6..29cfd93489b2 100644 --- a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs +++ b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs @@ -246,15 +246,9 @@ fn apply_color_map(tensor: &Tensor, annotations: &Arc<Annotations>) -> anyhow::R } } -fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { +fn height_width_depth(tensor: &Tensor) -> anyhow::Result<[u32; 3]> { use anyhow::Context as _; - crate::profile_function!(format!( - "dtype: {}, shape: {:?}", - tensor.dtype(), - tensor.shape() - )); - let shape = &tensor.shape(); anyhow::ensure!( @@ -277,8 +271,19 @@ fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { "We should make the same checks above, but with actual error messages" ); - use egui::epaint::ecolor::gamma_u8_from_linear_f32; - use egui::epaint::ecolor::linear_u8_from_linear_f32; + Ok([height, width, depth as u32]) +} + +fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { + crate::profile_function!(format!( + "dtype: {}, shape: {:?}", + tensor.dtype(), + tensor.shape() + )); + + let [height, width, depth] = height_width_depth(tensor)?; + + use egui::epaint::ecolor::{gamma_u8_from_linear_f32, linear_u8_from_linear_f32}; let size = [width as _, height as _]; @@ -353,7 +358,7 @@ fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { } (_depth, dtype) => { - anyhow::bail!("Don't know how to turn a tensor of shape={shape:?} and dtype={dtype:?} into a color image") + anyhow::bail!("Don't know how to turn a tensor of shape={:?} and dtype={dtype:?} into a color image", tensor.shape) } } } @@ -362,34 +367,17 @@ fn class_id_tensor_as_color_image( tensor: &Tensor, annotations: &Annotations, ) -> anyhow::Result<ColorImage> { - use anyhow::Context as _; - crate::profile_function!(format!( "dtype: {}, shape: {:?}", tensor.dtype(), tensor.shape() )); - let shape = &tensor.shape(); - - anyhow::ensure!( - shape.len() == 2 || shape.len() == 3, - "Expected a 2D or 3D tensor, got {shape:?}", - ); - - let [height, width] = [ - u32::try_from(shape[0].size).context("tensor too large")?, - u32::try_from(shape[1].size).context("tensor too large")?, - ]; - let depth = if shape.len() == 2 { 1 } else { shape[2].size }; - + let [height, width, depth] = height_width_depth(tensor)?; anyhow::ensure!( depth == 1, - "Cannot apply annotations to tensor of shape {shape:?}" - ); - debug_assert!( - tensor.is_shaped_like_an_image(), - "We should make the same checks above, but with actual error messages" + "Cannot apply annotations to tensor of shape {:?}", + tensor.shape ); let size = [width as _, height as _]; @@ -443,32 +431,14 @@ fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { return Ok(ColorImage::default()); } - use anyhow::Context as _; - crate::profile_function!(format!( "dtype: {}, shape: {:?}", tensor.dtype(), tensor.shape() )); - let shape = &tensor.shape(); - - anyhow::ensure!( - shape.len() == 2 || shape.len() == 3, - "Expected a 2D or 3D tensor, got {shape:?}", - ); - - let [height, width] = [ - u32::try_from(shape[0].size).context("tensor too large")?, - u32::try_from(shape[1].size).context("tensor too large")?, - ]; - let depth = if shape.len() == 2 { 1 } else { shape[2].size }; - - anyhow::ensure!(depth == 1, "Depth tensor of shape {shape:?}"); - debug_assert!( - tensor.is_shaped_like_an_image(), - "We should make the same checks above, but with actual error messages" - ); + let [height, width, depth] = height_width_depth(tensor)?; + anyhow::ensure!(depth == 1, "Depth tensor of shape {:?}", tensor.shape); let size = [width as _, height as _]; match &tensor.data { From 33227cefcfd505fa6b280ecf73af0e445e64b16b Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com> Date: Thu, 23 Mar 2023 14:44:52 +0100 Subject: [PATCH 3/5] Refactor how we color map depth images --- crates/re_log_types/src/data.rs | 32 +++++++++ .../src/misc/caches/tensor_image_cache.rs | 72 ++++++++++--------- 2 files changed, 70 insertions(+), 34 deletions(-) diff --git a/crates/re_log_types/src/data.rs b/crates/re_log_types/src/data.rs index 7d59384cdac8..779adfff4cf6 100644 --- a/crates/re_log_types/src/data.rs +++ b/crates/re_log_types/src/data.rs @@ -93,6 +93,38 @@ impl TensorDataType { Self::F64 => std::mem::size_of::<f64>() as _, } } + + pub fn is_float(&self) -> bool { + match self { + Self::U8 + | Self::U16 + | Self::U32 + | Self::U64 + | Self::I8 + | Self::I16 + | Self::I32 + | Self::I64 => false, + Self::F16 | Self::F32 | Self::F64 => true, + } + } + + pub fn max_value(&self) -> f64 { + match self { + Self::U8 => u8::MAX as _, + Self::U16 => u16::MAX as _, + Self::U32 => u32::MAX as _, + Self::U64 => u64::MAX as _, + + Self::I8 => i8::MAX as _, + Self::I16 => i16::MAX as _, + Self::I32 => i32::MAX as _, + Self::I64 => i64::MAX as _, + + Self::F16 => f16::MAX.into(), + Self::F32 => f32::MAX as _, + Self::F64 => f64::MAX, + } + } } impl std::fmt::Display for TensorDataType { diff --git a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs index 29cfd93489b2..2f3c97406f4b 100644 --- a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs +++ b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs @@ -12,7 +12,10 @@ use re_renderer::{ RenderContext, }; -use crate::ui::{Annotations, DefaultColor, MISSING_ANNOTATIONS}; +use crate::{ + misc::caches::TensorStats, + ui::{Annotations, DefaultColor, MISSING_ANNOTATIONS}, +}; // --- @@ -431,6 +434,12 @@ fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { return Ok(ColorImage::default()); } + // This function applies color mapping to a depth image. + // We are planning on moving this to the GPU: https://github.com/rerun-io/rerun/issues/1612 + // One big downside of the approach below is that if we have two depth images + // in the same range, they cannot be visually compared with each other, + // because their individual max-depths will be scaled to 65535. + crate::profile_function!(format!( "dtype: {}, shape: {:?}", tensor.dtype(), @@ -441,42 +450,37 @@ fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { anyhow::ensure!(depth == 1, "Depth tensor of shape {:?}", tensor.shape); let size = [width as _, height as _]; + let range = TensorStats::new(tensor).range.ok_or(anyhow::anyhow!( + "Depth image had no range!? Was this compressed?" + ))?; + + let (mut min, mut max) = range; + + anyhow::ensure!( + min.is_finite() && max.is_finite(), + "Depth image had non-finite values" + ); + + if min == max { + // Uniform image. We can't remap it to a 0-1 range, so do whatever: + min = 0.0; + max = if tensor.dtype().is_float() { + 1.0 + } else { + tensor.dtype().max_value() + }; + } + + fn colormap(t: f32) -> egui::Color32 { + let [r, g, b, _] = re_renderer::colormap_turbo_srgb(t); + egui::Color32::from_rgb(r, g, b) + } + match &tensor.data { TensorData::F32(buf) => { - // Convert to u16 so we can put them in an image. - // TODO(emilk): Eventually we want a renderer that can show f32 images natively. - // One big downside of the approach below is that if we have two depth images - // in the same range, they cannot be visually compared with each other, - // because their individual max-depths will be scaled to 65535. - - let mut min = f32::INFINITY; - let mut max = f32::NEG_INFINITY; - for float in buf.iter() { - min = min.min(*float); - max = max.max(*float); - } - - anyhow::ensure!( - min.is_finite() && max.is_finite(), - "Depth image had non-finite values" - ); - - let ints: Vec<u16> = if min == max { - // Uniform image. We can't remap it to a 0-1 range, so do whatever: - buf.iter().map(|&float| float as u16).collect() - } else { - buf.iter() - .map(|&float| egui::remap(float, min..=max, 0.0..=65535.0) as u16) - .collect() - }; - - let pixels = ints + let pixels = buf .iter() - .map(|pixel| { - let [r, g, b, _] = - re_renderer::colormap_turbo_srgb((*pixel as f32) / (u16::MAX as f32)); - egui::Color32::from_rgb(r, g, b) - }) + .map(|&float| colormap(egui::remap(float, min as f32..=max as f32, 0.0..=1.0))) .collect(); Ok(ColorImage { size, pixels }) From 59e3201d084c38dd44375a10e9be772fefff2e76 Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com> Date: Thu, 23 Mar 2023 14:51:35 +0100 Subject: [PATCH 4/5] Apply colormaps to all types of depth images --- .../component_types/arrow_convert_shims.rs | 5 ++ .../src/misc/caches/tensor_image_cache.rs | 54 +++++++++++++++---- 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/crates/re_log_types/src/component_types/arrow_convert_shims.rs b/crates/re_log_types/src/component_types/arrow_convert_shims.rs index 050575101229..8d196842169d 100644 --- a/crates/re_log_types/src/component_types/arrow_convert_shims.rs +++ b/crates/re_log_types/src/component_types/arrow_convert_shims.rs @@ -27,6 +27,11 @@ impl BinaryBuffer { pub fn as_slice(&self) -> &[u8] { self.0.as_slice() } + + #[inline] + pub fn iter(&self) -> impl Iterator<Item = &u8> { + self.0.iter() + } } impl Index<usize> for BinaryBuffer { diff --git a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs index 2f3c97406f4b..7ecb0bb1b0d7 100644 --- a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs +++ b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs @@ -471,22 +471,58 @@ fn depth_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> { }; } - fn colormap(t: f32) -> egui::Color32 { + let colormap = |value: f64| { + let t = egui::remap(value, min..=max, 0.0..=1.0) as f32; let [r, g, b, _] = re_renderer::colormap_turbo_srgb(t); egui::Color32::from_rgb(r, g, b) - } + }; match &tensor.data { - TensorData::F32(buf) => { - let pixels = buf - .iter() - .map(|&float| colormap(egui::remap(float, min as f32..=max as f32, 0.0..=1.0))) - .collect(); + TensorData::U8(buf) => { + let pixels = buf.iter().map(|&value| colormap(value as _)).collect(); + Ok(ColorImage { size, pixels }) + } + TensorData::U16(buf) => { + let pixels = buf.iter().map(|&value| colormap(value as _)).collect(); + Ok(ColorImage { size, pixels }) + } + TensorData::U32(buf) => { + let pixels = buf.iter().map(|&value| colormap(value as _)).collect(); + Ok(ColorImage { size, pixels }) + } + TensorData::U64(buf) => { + let pixels = buf.iter().map(|&value| colormap(value as _)).collect(); + Ok(ColorImage { size, pixels }) + } + TensorData::I8(buf) => { + let pixels = buf.iter().map(|&value| colormap(value as _)).collect(); Ok(ColorImage { size, pixels }) } - _ => { - color_tensor_as_color_image(tensor) // TODO(emilk): support more depth dtypes + TensorData::I16(buf) => { + let pixels = buf.iter().map(|&value| colormap(value as _)).collect(); + Ok(ColorImage { size, pixels }) + } + TensorData::I32(buf) => { + let pixels = buf.iter().map(|&value| colormap(value as _)).collect(); + Ok(ColorImage { size, pixels }) + } + TensorData::I64(buf) => { + let pixels = buf.iter().map(|&value| colormap(value as _)).collect(); + Ok(ColorImage { size, pixels }) + } + + TensorData::F32(buf) => { + let pixels = buf.iter().map(|&value| colormap(value as _)).collect(); + Ok(ColorImage { size, pixels }) + } + TensorData::F64(buf) => { + let pixels = buf.iter().map(|&value| colormap(value)).collect(); + Ok(ColorImage { size, pixels }) + } + + TensorData::JPEG(_) => { + anyhow::bail!("Cannot apply colormap to JPEG image") } } } From df0c0314a5bb4ccc82cd38fc1b3b32d8247fb7ce Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com> Date: Thu, 23 Mar 2023 14:54:27 +0100 Subject: [PATCH 5/5] Add a comment --- crates/re_viewer/src/misc/caches/mod.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/re_viewer/src/misc/caches/mod.rs b/crates/re_viewer/src/misc/caches/mod.rs index f0363a68b5d8..db61a56109e2 100644 --- a/crates/re_viewer/src/misc/caches/mod.rs +++ b/crates/re_viewer/src/misc/caches/mod.rs @@ -53,6 +53,7 @@ impl Caches { } pub struct TensorStats { + /// This will currently only be `None` for jpeg-encoded tensors. pub range: Option<(f64, f64)>, }