Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Direct conversion to dynamic image from the non-classic Tensor #1455

Merged
merged 2 commits into from
Mar 1, 2023
Merged

Direct conversion to dynamic image from the non-classic Tensor #1455

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3acfcf3
Direct conversion to dynamic image from the non-classic Tensor
jleibs Feb 28, 2023
bc5317a
fix weird cargo fmt
jleibs Mar 1, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 52 additions & 77 deletions crates/re_viewer/src/misc/caches/tensor_image_cache.rs
View file
Open in desktop
Original file line number Original file line Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@ use egui::{Color32, ColorImage};
use egui_extras::RetainedImage; use egui_extras::RetainedImage;
use image::DynamicImage; use image::DynamicImage;
use re_log_types::{ use re_log_types::{
component_types::{self, ClassId, TensorDataMeaning, TensorTrait}, component_types::{self, ClassId, TensorData, TensorDataMeaning, TensorTrait},
ClassicTensor, MsgId, TensorDataStore, TensorDataType, ClassicTensor, MsgId,
}; };
use re_renderer::{ use re_renderer::{
resource_managers::{GpuTexture2DHandle, Texture2DCreationDesc}, resource_managers::{GpuTexture2DHandle, Texture2DCreationDesc},
Expand Down Expand Up @@ -240,22 +240,21 @@ pub trait AsDynamicImage: component_types::TensorTrait {
fn as_dynamic_image(&self, annotations: &Arc<Annotations>) -> anyhow::Result<DynamicImage>; fn as_dynamic_image(&self, annotations: &Arc<Annotations>) -> anyhow::Result<DynamicImage>;
} }


impl AsDynamicImage for component_types::Tensor { impl AsDynamicImage for ClassicTensor {
fn as_dynamic_image(&self, annotations: &Arc<Annotations>) -> anyhow::Result<DynamicImage> { fn as_dynamic_image(&self, _annotations: &Arc<Annotations>) -> anyhow::Result<DynamicImage> {
crate::profile_function!(); anyhow::bail!("ClassicTensor is deprecated")
let classic_tensor = ClassicTensor::from(self);
classic_tensor.as_dynamic_image(annotations)
} }
} }


impl AsDynamicImage for ClassicTensor { impl AsDynamicImage for component_types::Tensor {
fn as_dynamic_image(&self, annotations: &Arc<Annotations>) -> anyhow::Result<DynamicImage> { fn as_dynamic_image(&self, annotations: &Arc<Annotations>) -> anyhow::Result<DynamicImage> {
use anyhow::Context as _; use anyhow::Context as _;
let tensor = self; let tensor = self;


crate::profile_function!(format!( crate::profile_function!(format!(
"dtype: {}, meaning: {:?}", "dtype: {}, meaning: {:?}",
tensor.dtype, tensor.meaning tensor.dtype(),
tensor.meaning
)); ));


let shape = &tensor.shape(); let shape = &tensor.shape();
Expand Down Expand Up @@ -287,15 +286,8 @@ impl AsDynamicImage for ClassicTensor {
use egui::epaint::ecolor::gamma_u8_from_linear_f32; use egui::epaint::ecolor::gamma_u8_from_linear_f32;
use egui::epaint::ecolor::linear_u8_from_linear_f32; use egui::epaint::ecolor::linear_u8_from_linear_f32;


match &tensor.data { match (depth, &tensor.data, tensor.meaning) {
TensorDataStore::Dense(bytes) => { (1, TensorData::U8(buf), TensorDataMeaning::ClassId) => {
anyhow::ensure!(
bytes.len() as u64 == tensor.len() * tensor.dtype().size(),
"Tensor data length doesn't match tensor shape and dtype"
);

match (depth, tensor.dtype, tensor.meaning) {
(1, TensorDataType::U8, TensorDataMeaning::ClassId) => {
// Apply annotation mapping to raw bytes interpreted as u8 // Apply annotation mapping to raw bytes interpreted as u8
let color_lookup: Vec<[u8; 4]> = (0..256) let color_lookup: Vec<[u8; 4]> = (0..256)
.map(|id| { .map(|id| {
Expand All @@ -306,7 +298,7 @@ impl AsDynamicImage for ClassicTensor {
.to_array() .to_array()
}) })
.collect(); .collect();
let color_bytes = bytes let color_bytes = buf
.iter() .iter()
.flat_map(|p: &u8| color_lookup[*p as usize]) .flat_map(|p: &u8| color_lookup[*p as usize])
.collect(); .collect();
Expand All @@ -315,10 +307,10 @@ impl AsDynamicImage for ClassicTensor {
.context("Bad RGBA8") .context("Bad RGBA8")
.map(DynamicImage::ImageRgba8) .map(DynamicImage::ImageRgba8)
} }
(1, TensorDataType::U16, TensorDataMeaning::ClassId) => { (1, TensorData::U16(buf), TensorDataMeaning::ClassId) => {
// Apply annotations mapping to bytes interpreted as u16 // Apply annotations mapping to bytes interpreted as u16
let mut color_lookup: ahash::HashMap<u16, [u8; 4]> = Default::default(); let mut color_lookup: ahash::HashMap<u16, [u8; 4]> = Default::default();
let color_bytes = bytemuck::cast_slice(bytes) let color_bytes = buf
.iter() .iter()
.flat_map(|id: &u16| { .flat_map(|id: &u16| {
*color_lookup.entry(*id).or_insert_with(|| { *color_lookup.entry(*id).or_insert_with(|| {
Expand All @@ -335,24 +327,22 @@ impl AsDynamicImage for ClassicTensor {
.context("Bad RGBA8") .context("Bad RGBA8")
.map(DynamicImage::ImageRgba8) .map(DynamicImage::ImageRgba8)
} }
(1, TensorDataType::U8, _) => { (1, TensorData::U8(buf), _) => {
// TODO(emilk): we should read some meta-data to check if this is luminance or alpha. // TODO(emilk): we should read some meta-data to check if this is luminance or alpha.
image::GrayImage::from_raw(width, height, bytes.to_vec()) image::GrayImage::from_raw(width, height, buf.clone())
.context("Bad Luminance8") .context("Bad Luminance8")
.map(DynamicImage::ImageLuma8) .map(DynamicImage::ImageLuma8)
} }
(1, TensorDataType::U16, _) => { (1, TensorData::U16(buf), _) => {
// TODO(emilk): we should read some meta-data to check if this is luminance or alpha. // TODO(emilk): we should read some meta-data to check if this is luminance or alpha.
Gray16Image::from_raw(width, height, bytemuck::cast_slice(bytes).to_vec()) Gray16Image::from_raw(width, height, buf.to_vec())
.context("Bad Luminance16") .context("Bad Luminance16")
.map(DynamicImage::ImageLuma16) .map(DynamicImage::ImageLuma16)
} }
(1, TensorDataType::F32, TensorDataMeaning::Depth) => { (1, TensorData::F32(buf), TensorDataMeaning::Depth) => {
if bytes.is_empty() { if buf.is_empty() {
Ok(DynamicImage::ImageLuma16(Gray16Image::default())) Ok(DynamicImage::ImageLuma16(Gray16Image::default()))
} else { } else {
let floats = bytemuck::cast_slice(bytes);

// Convert to u16 so we can put them in an image. // Convert to u16 so we can put them in an image.
// TODO(emilk): Eventually we want a renderer that can show f32 images natively. // 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 // One big downside of the approach below is that if we have two depth images
Expand All @@ -361,9 +351,9 @@ impl AsDynamicImage for ClassicTensor {


let mut min = f32::INFINITY; let mut min = f32::INFINITY;
let mut max = f32::NEG_INFINITY; let mut max = f32::NEG_INFINITY;
for &float in floats { for float in buf.iter() {
min = min.min(float); min = min.min(*float);
max = max.max(float); max = max.max(*float);
} }


anyhow::ensure!( anyhow::ensure!(
Expand All @@ -373,16 +363,14 @@ impl AsDynamicImage for ClassicTensor {


if min == max { if min == max {
// Uniform image. We can't remap it to a 0-1 range, so do whatever: // Uniform image. We can't remap it to a 0-1 range, so do whatever:
let ints = floats.iter().map(|&float| float as u16).collect(); let ints = buf.iter().map(|&float| float as u16).collect();
Gray16Image::from_raw(width, height, ints) Gray16Image::from_raw(width, height, ints)
.context("Bad Luminance16") .context("Bad Luminance16")
.map(DynamicImage::ImageLuma16) .map(DynamicImage::ImageLuma16)
} else { } else {
let ints = floats let ints = buf
.iter() .iter()
.map(|&float| { .map(|&float| egui::remap(float, min..=max, 0.0..=65535.0) as u16)
egui::remap(float, min..=max, 0.0..=65535.0) as u16
})
.collect(); .collect();


Gray16Image::from_raw(width, height, ints) Gray16Image::from_raw(width, height, ints)
Expand All @@ -391,26 +379,21 @@ impl AsDynamicImage for ClassicTensor {
} }
} }
} }
(1, TensorDataType::F32, _) => { (1, TensorData::F32(buf), _) => {
let l: &[f32] = bytemuck::cast_slice(bytes); let l = buf.as_slice();
let colors: Vec<u8> = let colors: Vec<u8> = l.iter().copied().map(linear_u8_from_linear_f32).collect();
l.iter().copied().map(linear_u8_from_linear_f32).collect();
image::GrayImage::from_raw(width, height, colors) image::GrayImage::from_raw(width, height, colors)
.context("Bad Luminance f32") .context("Bad Luminance f32")
.map(DynamicImage::ImageLuma8) .map(DynamicImage::ImageLuma8)
} }
(3, TensorDataType::U8, _) => { (3, TensorData::U8(buf), _) => image::RgbImage::from_raw(width, height, buf.clone())
image::RgbImage::from_raw(width, height, bytes.to_vec())
.context("Bad RGB8") .context("Bad RGB8")
.map(DynamicImage::ImageRgb8) .map(DynamicImage::ImageRgb8),
} (3, TensorData::U16(buf), _) => Rgb16Image::from_raw(width, height, buf.to_vec())
(3, TensorDataType::U16, _) => {
Rgb16Image::from_raw(width, height, bytemuck::cast_slice(bytes).to_vec())
.context("Bad RGB16 image") .context("Bad RGB16 image")
.map(DynamicImage::ImageRgb16) .map(DynamicImage::ImageRgb16),
} (3, TensorData::F32(buf), _) => {
(3, TensorDataType::F32, _) => { let rgb: &[[f32; 3]] = bytemuck::cast_slice(buf.as_slice());
let rgb: &[[f32; 3]] = bytemuck::cast_slice(bytes);
let colors: Vec<u8> = rgb let colors: Vec<u8> = rgb
.iter() .iter()
.flat_map(|&[r, g, b]| { .flat_map(|&[r, g, b]| {
Expand All @@ -425,18 +408,14 @@ impl AsDynamicImage for ClassicTensor {
.map(DynamicImage::ImageRgb8) .map(DynamicImage::ImageRgb8)
} }


(4, TensorDataType::U8, _) => { (4, TensorData::U8(buf), _) => image::RgbaImage::from_raw(width, height, buf.clone())
image::RgbaImage::from_raw(width, height, bytes.to_vec())
.context("Bad RGBA8") .context("Bad RGBA8")
.map(DynamicImage::ImageRgba8) .map(DynamicImage::ImageRgba8),
} (4, TensorData::U16(buf), _) => Rgba16Image::from_raw(width, height, buf.to_vec())
(4, TensorDataType::U16, _) => {
Rgba16Image::from_raw(width, height, bytemuck::cast_slice(bytes).to_vec())
.context("Bad RGBA16 image") .context("Bad RGBA16 image")
.map(DynamicImage::ImageRgba16) .map(DynamicImage::ImageRgba16),
} (4, TensorData::F32(buf), _) => {
(4, TensorDataType::F32, _) => { let rgba: &[[f32; 4]] = bytemuck::cast_slice(buf.as_slice());
let rgba: &[[f32; 4]] = bytemuck::cast_slice(bytes);
let colors: Vec<u8> = rgba let colors: Vec<u8> = rgba
.iter() .iter()
.flat_map(|&[r, g, b, a]| { .flat_map(|&[r, g, b, a]| {
Expand All @@ -451,23 +430,9 @@ impl AsDynamicImage for ClassicTensor {
.context("Bad RGBA f32") .context("Bad RGBA f32")
.map(DynamicImage::ImageRgba8) .map(DynamicImage::ImageRgba8)
} }
(_depth, dtype, meaning @ TensorDataMeaning::ClassId) => { (depth, TensorData::JPEG(buf), _) => {
anyhow::bail!(
"Shape={shape:?} and dtype={dtype:?} is incompatible with meaning={meaning:?}"
)
}

(_depth, dtype, _) => {
anyhow::bail!(
"Don't know how to turn a tensor of shape={shape:?} and dtype={dtype:?} into an image"
)
}
}
}

TensorDataStore::Jpeg(bytes) => {
use image::io::Reader as ImageReader; use image::io::Reader as ImageReader;
let mut reader = ImageReader::new(std::io::Cursor::new(bytes)); let mut reader = ImageReader::new(std::io::Cursor::new(buf));
reader.set_format(image::ImageFormat::Jpeg); reader.set_format(image::ImageFormat::Jpeg);
// TODO(emilk): handle grayscale JPEG:s (depth == 1) // TODO(emilk): handle grayscale JPEG:s (depth == 1)
let img = { let img = {
Expand All @@ -486,6 +451,16 @@ impl AsDynamicImage for ClassicTensor {


Ok(DynamicImage::ImageRgb8(img)) Ok(DynamicImage::ImageRgb8(img))
} }

(_depth, dtype, meaning @ TensorDataMeaning::ClassId) => {
anyhow::bail!(
"Shape={shape:?} and dtype={dtype:?} is incompatible with meaning={meaning:?}"
)
}

(_depth, dtype, _) => {
anyhow::bail!("Don't know how to turn a tensor of shape={shape:?} and dtype={dtype:?} into an image")
}
} }
} }
} }
Expand Down