diff --git a/crates/re_log_types/src/component_types/tensor.rs b/crates/re_log_types/src/component_types/tensor.rs index 3261183f83f6..2605b5c9b84e 100644 --- a/crates/re_log_types/src/component_types/tensor.rs +++ b/crates/re_log_types/src/component_types/tensor.rs @@ -142,7 +142,7 @@ impl ArrowDeserialize for TensorId { /// ), /// ); /// ``` -#[derive(Clone, Debug, PartialEq, ArrowField, ArrowSerialize, ArrowDeserialize)] +#[derive(Clone, PartialEq, ArrowField, ArrowSerialize, ArrowDeserialize)] #[arrow_field(type = "dense")] pub enum TensorData { U8(BinaryBuffer), @@ -198,6 +198,24 @@ impl TensorData { } } +impl std::fmt::Debug for TensorData { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::U8(_) => write!(f, "U8({} bytes)", self.size_in_bytes()), + Self::U16(_) => write!(f, "U16({} bytes)", self.size_in_bytes()), + Self::U32(_) => write!(f, "U32({} bytes)", self.size_in_bytes()), + Self::U64(_) => write!(f, "U64({} bytes)", self.size_in_bytes()), + Self::I8(_) => write!(f, "I8({} bytes)", self.size_in_bytes()), + Self::I16(_) => write!(f, "I16({} bytes)", self.size_in_bytes()), + Self::I32(_) => write!(f, "I32({} bytes)", self.size_in_bytes()), + Self::I64(_) => write!(f, "I64({} bytes)", self.size_in_bytes()), + Self::F32(_) => write!(f, "F32({} bytes)", self.size_in_bytes()), + Self::F64(_) => write!(f, "F64({} bytes)", self.size_in_bytes()), + Self::JPEG(_) => write!(f, "JPEG({} bytes)", self.size_in_bytes()), + } + } +} + /// Flattened `Tensor` data payload /// /// ## Examples diff --git a/crates/re_renderer/shader/rectangle.wgsl b/crates/re_renderer/shader/rectangle.wgsl index 57584df1abb1..2bf5f5517624 100644 --- a/crates/re_renderer/shader/rectangle.wgsl +++ b/crates/re_renderer/shader/rectangle.wgsl @@ -6,9 +6,10 @@ // Keep in sync with mirror in rectangle.rs // Which texture to read from? -const SAMPLE_TYPE_FLOAT = 1u; -const SAMPLE_TYPE_SINT = 2u; -const SAMPLE_TYPE_UINT = 3u; +const SAMPLE_TYPE_FLOAT_FILTER = 1u; +const SAMPLE_TYPE_FLOAT_NOFILTER = 2u; +const SAMPLE_TYPE_SINT_NOFILTER = 3u; +const SAMPLE_TYPE_UINT_NOFILTER = 4u; // How do we do colormapping? const COLOR_MAPPER_OFF = 1u; @@ -67,6 +68,9 @@ var texture_uint: texture_2d; @group(1) @binding(5) var colormap_texture: texture_2d; +@group(1) @binding(6) +var texture_float_filterable: texture_2d; + struct VertexOut { @builtin(position) position: Vec4, @@ -90,12 +94,16 @@ fn vs_main(@builtin(vertex_index) v_idx: u32) -> VertexOut { fn fs_main(in: VertexOut) -> @location(0) Vec4 { // Sample the main texture: var sampled_value: Vec4; - if rect_info.sample_type == SAMPLE_TYPE_FLOAT { - sampled_value = textureSampleLevel(texture_float, texture_sampler, in.texcoord, 0.0); // TODO(emilk): support mipmaps - } else if rect_info.sample_type == SAMPLE_TYPE_SINT { + if rect_info.sample_type == SAMPLE_TYPE_FLOAT_FILTER { + // TODO(emilk): support mipmaps + sampled_value = textureSampleLevel(texture_float_filterable, texture_sampler, in.texcoord, 0.0); + } else if rect_info.sample_type == SAMPLE_TYPE_FLOAT_NOFILTER { + let icoords = IVec2(in.texcoord * Vec2(textureDimensions(texture_float).xy)); + sampled_value = Vec4(textureLoad(texture_float, icoords, 0)); + } else if rect_info.sample_type == SAMPLE_TYPE_SINT_NOFILTER { let icoords = IVec2(in.texcoord * Vec2(textureDimensions(texture_sint).xy)); sampled_value = Vec4(textureLoad(texture_sint, icoords, 0)); - } else if rect_info.sample_type == SAMPLE_TYPE_UINT { + } else if rect_info.sample_type == SAMPLE_TYPE_UINT_NOFILTER { let icoords = IVec2(in.texcoord * Vec2(textureDimensions(texture_uint).xy)); sampled_value = Vec4(textureLoad(texture_uint, icoords, 0)); } else { diff --git a/crates/re_renderer/src/renderer/rectangles.rs b/crates/re_renderer/src/renderer/rectangles.rs index f869f055570b..d002c73fea6c 100644 --- a/crates/re_renderer/src/renderer/rectangles.rs +++ b/crates/re_renderer/src/renderer/rectangles.rs @@ -182,9 +182,10 @@ mod gpu_data { // Keep in sync with mirror in rectangle.wgsl // Which texture to read from? - const SAMPLE_TYPE_FLOAT: u32 = 1; - const SAMPLE_TYPE_SINT: u32 = 2; - const SAMPLE_TYPE_UINT: u32 = 3; + const SAMPLE_TYPE_FLOAT_FILTER: u32 = 1; + const SAMPLE_TYPE_FLOAT_NOFILTER: u32 = 2; + const SAMPLE_TYPE_SINT_NOFILTER: u32 = 3; + const SAMPLE_TYPE_UINT_NOFILTER: u32 = 4; // How do we do colormapping? const COLOR_MAPPER_OFF: u32 = 1; @@ -232,12 +233,18 @@ mod gpu_data { } = &rectangle.colormapped_texture; let sample_type = match texture_info.sample_type { - wgpu::TextureSampleType::Float { .. } => SAMPLE_TYPE_FLOAT, + wgpu::TextureSampleType::Float { .. } => { + if super::is_float_filterable(texture_format) { + SAMPLE_TYPE_FLOAT_FILTER + } else { + SAMPLE_TYPE_FLOAT_NOFILTER + } + } wgpu::TextureSampleType::Depth => { return Err(RectangleError::DepthTexturesNotSupported); } - wgpu::TextureSampleType::Sint => SAMPLE_TYPE_SINT, - wgpu::TextureSampleType::Uint => SAMPLE_TYPE_UINT, + wgpu::TextureSampleType::Sint => SAMPLE_TYPE_SINT_NOFILTER, + wgpu::TextureSampleType::Uint => SAMPLE_TYPE_UINT_NOFILTER, }; let mut colormap_function = 0; @@ -378,14 +385,19 @@ impl RectangleDrawData { )); } - // We set up three texture sources, then instruct the shader to read from at most one of them. - let mut texture_float = ctx.texture_manager_2d.zeroed_texture_float().handle; + // We set up several texture sources, then instruct the shader to read from at most one of them. + let mut texture_float_filterable = ctx.texture_manager_2d.zeroed_texture_float().handle; + let mut texture_float_nofilter = ctx.texture_manager_2d.zeroed_texture_float().handle; let mut texture_sint = ctx.texture_manager_2d.zeroed_texture_sint().handle; let mut texture_uint = ctx.texture_manager_2d.zeroed_texture_uint().handle; match texture_description.sample_type { wgpu::TextureSampleType::Float { .. } => { - texture_float = texture.handle; + if is_float_filterable(&texture_format) { + texture_float_filterable = texture.handle; + } else { + texture_float_nofilter = texture.handle; + } } wgpu::TextureSampleType::Depth => { return Err(RectangleError::DepthTexturesNotSupported); @@ -422,10 +434,11 @@ impl RectangleDrawData { entries: smallvec![ uniform_buffer, BindGroupEntry::Sampler(sampler), - BindGroupEntry::DefaultTextureView(texture_float), + BindGroupEntry::DefaultTextureView(texture_float_nofilter), BindGroupEntry::DefaultTextureView(texture_sint), BindGroupEntry::DefaultTextureView(texture_uint), BindGroupEntry::DefaultTextureView(colormap_texture), + BindGroupEntry::DefaultTextureView(texture_float_filterable), ], layout: rectangle_renderer.bind_group_layout, }, @@ -483,12 +496,12 @@ impl Renderer for RectangleRenderer { ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering), count: None, }, - // float texture: + // float textures without filtering (e.g. R32Float): wgpu::BindGroupLayoutEntry { binding: 2, visibility: wgpu::ShaderStages::FRAGMENT, ty: wgpu::BindingType::Texture { - sample_type: wgpu::TextureSampleType::Float { filterable: true }, + sample_type: wgpu::TextureSampleType::Float { filterable: false }, view_dimension: wgpu::TextureViewDimension::D2, multisampled: false, }, @@ -527,6 +540,17 @@ impl Renderer for RectangleRenderer { }, count: None, }, + // float textures with filtering (e.g. Rgba8UnormSrgb): + wgpu::BindGroupLayoutEntry { + binding: 6, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Texture { + sample_type: wgpu::TextureSampleType::Float { filterable: true }, + view_dimension: wgpu::TextureViewDimension::D2, + multisampled: false, + }, + count: None, + }, ], }, ); @@ -653,3 +677,11 @@ impl Renderer for RectangleRenderer { ] } } + +fn is_float_filterable(format: &wgpu::TextureFormat) -> bool { + format + .describe() + .guaranteed_format_features + .flags + .contains(wgpu::TextureFormatFeatureFlags::FILTERABLE) +} 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 d8455f2bcaa9..09b29a54c5d8 100644 --- a/crates/re_viewer/src/misc/caches/tensor_image_cache.rs +++ b/crates/re_viewer/src/misc/caches/tensor_image_cache.rs @@ -328,8 +328,12 @@ fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result { Ok(ColorImage { size, pixels }) } - (_depth, dtype) => { - anyhow::bail!("Don't know how to turn a tensor of shape={:?} and dtype={dtype:?} into a color image", tensor.shape) + (_depth, _dtype) => { + anyhow::bail!( + "Don't know how to turn a tensor of shape={:?} and dtype={:?} into a color image", + tensor.shape, + tensor.dtype() + ) } } } diff --git a/crates/re_viewer/src/misc/tensor_to_gpu.rs b/crates/re_viewer/src/misc/tensor_to_gpu.rs index 5417a8799f71..f6882cd2c203 100644 --- a/crates/re_viewer/src/misc/tensor_to_gpu.rs +++ b/crates/re_viewer/src/misc/tensor_to_gpu.rs @@ -262,63 +262,114 @@ fn general_texture_creation_desc_from_tensor<'a>( tensor: &'a Tensor, ) -> anyhow::Result> { let [height, width, depth] = height_width_depth(tensor)?; - let (data, format) = match (depth, &tensor.data) { - (1, TensorData::U8(buf)) => (cast_slice_to_cow(buf.as_slice()), TextureFormat::R8Uint), - (1, TensorData::I8(buf)) => (cast_slice_to_cow(buf), TextureFormat::R8Sint), - (1, TensorData::U16(buf)) => (cast_slice_to_cow(buf), TextureFormat::R16Uint), - (1, TensorData::I16(buf)) => (cast_slice_to_cow(buf), TextureFormat::R16Sint), - (1, TensorData::U32(buf)) => (cast_slice_to_cow(buf), TextureFormat::R32Uint), - (1, TensorData::I32(buf)) => (cast_slice_to_cow(buf), TextureFormat::R32Sint), - // (1, TensorData::F16(buf)) => (cast_slice_to_cow(buf), TextureFormat::R16Float), TODO(#854) - (1, TensorData::F32(buf)) => (cast_slice_to_cow(buf), TextureFormat::R32Float), - (1, TensorData::F64(buf)) => (narrow_f64_to_f32s(buf), TextureFormat::R32Float), - - // NOTE: 2-channel images are not supported by the shader yet, but are included here for completeness: - (2, TensorData::U8(buf)) => (cast_slice_to_cow(buf.as_slice()), TextureFormat::Rg8Uint), - (2, TensorData::I8(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rg8Sint), - (2, TensorData::U16(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rg16Uint), - (2, TensorData::I16(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rg16Sint), - (2, TensorData::U32(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rg32Uint), - (2, TensorData::I32(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rg32Sint), - // (2, TensorData::F16(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rg16Float), TODO(#854) - (2, TensorData::F32(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rg32Float), - (2, TensorData::F64(buf)) => (narrow_f64_to_f32s(buf), TextureFormat::Rg32Float), - - // There are no 3-channel textures in wgpu, so we need to pad to 4 channels: - (3, TensorData::U8(buf)) => (pad_and_cast(buf.as_slice(), 0), TextureFormat::Rgba8Uint), - (3, TensorData::I8(buf)) => (pad_and_cast(buf, 0), TextureFormat::Rgba8Sint), - (3, TensorData::U16(buf)) => (pad_and_cast(buf, 0), TextureFormat::Rgba16Uint), - (3, TensorData::I16(buf)) => (pad_and_cast(buf, 0), TextureFormat::Rgba16Sint), - (3, TensorData::U32(buf)) => (pad_and_cast(buf, 0), TextureFormat::Rgba32Uint), - (3, TensorData::I32(buf)) => (pad_and_cast(buf, 0), TextureFormat::Rgba32Sint), - // (3, TensorData::F16(buf)) => (pad_and_cast(buf, 0.0), TextureFormat::Rgba16Float), TODO(#854) - (3, TensorData::F32(buf)) => (pad_and_cast(buf, 0.0), TextureFormat::Rgba32Float), - (3, TensorData::F64(buf)) => { - let pad = 0.0; - let floats: Vec = buf - .chunks_exact(3) - .flat_map(|chunk| [chunk[0] as f32, chunk[1] as f32, chunk[2] as f32, pad]) - .collect(); - ( - pod_collect_to_vec(&floats).into(), - TextureFormat::Rgba32Float, - ) - } - // TODO(emilk): premultiply alpha - (4, TensorData::U8(buf)) => (cast_slice_to_cow(buf.as_slice()), TextureFormat::Rgba8Uint), - (4, TensorData::I8(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rgba8Sint), - (4, TensorData::U16(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rgba16Uint), - (4, TensorData::I16(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rgba16Sint), - (4, TensorData::U32(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rgba32Uint), - (4, TensorData::I32(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rgba32Sint), - // (4, TensorData::F16(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rgba16Float), TODO(#854) - (4, TensorData::F32(buf)) => (cast_slice_to_cow(buf), TextureFormat::Rgba32Float), - (4, TensorData::F64(buf)) => (narrow_f64_to_f32s(buf), TextureFormat::Rgba32Float), - - // TODO(emilk): U64/I64 - (_depth, dtype) => { - anyhow::bail!("Don't know how to turn a tensor of shape={:?} and dtype={dtype:?} into a color image", tensor.shape) + let (data, format) = match depth { + 1 => { + match &tensor.data { + TensorData::U8(buf) => (cast_slice_to_cow(buf.as_slice()), TextureFormat::R8Uint), + TensorData::U16(buf) => (cast_slice_to_cow(buf), TextureFormat::R16Uint), + TensorData::U32(buf) => (cast_slice_to_cow(buf), TextureFormat::R32Uint), + TensorData::U64(buf) => (narrow_u64_to_f32s(buf), TextureFormat::R32Float), // narrowing to f32! + + TensorData::I8(buf) => (cast_slice_to_cow(buf), TextureFormat::R8Sint), + TensorData::I16(buf) => (cast_slice_to_cow(buf), TextureFormat::R16Sint), + TensorData::I32(buf) => (cast_slice_to_cow(buf), TextureFormat::R32Sint), + TensorData::I64(buf) => (narrow_i64_to_f32s(buf), TextureFormat::R32Float), // narrowing to f32! + + // TensorData::F16(buf) => (cast_slice_to_cow(buf), TextureFormat::R16Float), TODO(#854) + TensorData::F32(buf) => (cast_slice_to_cow(buf), TextureFormat::R32Float), + TensorData::F64(buf) => (narrow_f64_to_f32s(buf), TextureFormat::R32Float), // narrowing to f32! + + TensorData::JPEG(_) => { + anyhow::bail!("JPEGs should have been decoded at this point") + } + } + } + 2 => { + // NOTE: 2-channel images are not supported by the shader yet, but are included here for completeness: + match &tensor.data { + TensorData::U8(buf) => (cast_slice_to_cow(buf.as_slice()), TextureFormat::Rg8Uint), + TensorData::U16(buf) => (cast_slice_to_cow(buf), TextureFormat::Rg16Uint), + TensorData::U32(buf) => (cast_slice_to_cow(buf), TextureFormat::Rg32Uint), + TensorData::U64(buf) => (narrow_u64_to_f32s(buf), TextureFormat::Rg32Float), // narrowing to f32! + + TensorData::I8(buf) => (cast_slice_to_cow(buf), TextureFormat::Rg8Sint), + TensorData::I16(buf) => (cast_slice_to_cow(buf), TextureFormat::Rg16Sint), + TensorData::I32(buf) => (cast_slice_to_cow(buf), TextureFormat::Rg32Sint), + TensorData::I64(buf) => (narrow_i64_to_f32s(buf), TextureFormat::Rg32Float), // narrowing to f32! + + // TensorData::F16(buf) => (cast_slice_to_cow(buf), TextureFormat::Rg16Float), TODO(#854) + TensorData::F32(buf) => (cast_slice_to_cow(buf), TextureFormat::Rg32Float), + TensorData::F64(buf) => (narrow_f64_to_f32s(buf), TextureFormat::Rg32Float), // narrowing to f32! + + TensorData::JPEG(_) => { + anyhow::bail!("JPEGs should have been decoded at this point") + } + } + } + 3 => { + // There are no 3-channel textures in wgpu, so we need to pad to 4 channels. + // What should we pad with? It depends on whether or not the shader interprets these as alpha. + // To be safe, we pad with the MAX value of integers, and with 1.0 for floats. + // TODO(emilk): tell the shader to ignore the alpha channel instead! + match &tensor.data { + TensorData::U8(buf) => ( + pad_and_cast(buf.as_slice(), u8::MAX), + TextureFormat::Rgba8Uint, + ), + TensorData::U16(buf) => (pad_and_cast(buf, u16::MAX), TextureFormat::Rgba16Uint), + TensorData::U32(buf) => (pad_and_cast(buf, u32::MAX), TextureFormat::Rgba32Uint), + TensorData::U64(buf) => ( + pad_and_narrow_and_cast(buf, 1.0, |x: u64| x as f32), + TextureFormat::Rgba32Float, + ), + + TensorData::I8(buf) => (pad_and_cast(buf, i8::MAX), TextureFormat::Rgba8Sint), + TensorData::I16(buf) => (pad_and_cast(buf, i16::MAX), TextureFormat::Rgba16Sint), + TensorData::I32(buf) => (pad_and_cast(buf, i32::MAX), TextureFormat::Rgba32Sint), + TensorData::I64(buf) => ( + pad_and_narrow_and_cast(buf, 1.0, |x: i64| x as f32), + TextureFormat::Rgba32Float, + ), + + // TensorData::F16(buf) => (pad_and_cast(buf, 1.0), TextureFormat::Rgba16Float), TODO(#854) + TensorData::F32(buf) => (pad_and_cast(buf, 1.0), TextureFormat::Rgba32Float), + TensorData::F64(buf) => ( + pad_and_narrow_and_cast(buf, 1.0, |x: f64| x as f32), + TextureFormat::Rgba32Float, + ), + + TensorData::JPEG(_) => { + anyhow::bail!("JPEGs should have been decoded at this point") + } + } + } + 4 => { + // TODO(emilk): premultiply alpha, or tell the shader to assume unmultiplied alpha + match &tensor.data { + TensorData::U8(buf) => { + (cast_slice_to_cow(buf.as_slice()), TextureFormat::Rgba8Uint) + } + TensorData::U16(buf) => (cast_slice_to_cow(buf), TextureFormat::Rgba16Uint), + TensorData::U32(buf) => (cast_slice_to_cow(buf), TextureFormat::Rgba32Uint), + TensorData::U64(buf) => (narrow_u64_to_f32s(buf), TextureFormat::Rgba32Float), // narrowing to f32! + + TensorData::I8(buf) => (cast_slice_to_cow(buf), TextureFormat::Rgba8Sint), + TensorData::I16(buf) => (cast_slice_to_cow(buf), TextureFormat::Rgba16Sint), + TensorData::I32(buf) => (cast_slice_to_cow(buf), TextureFormat::Rgba32Sint), + TensorData::I64(buf) => (narrow_i64_to_f32s(buf), TextureFormat::Rgba32Float), // narrowing to f32! + + // TensorData::F16(buf) => (cast_slice_to_cow(buf), TextureFormat::Rgba16Float), TODO(#854) + TensorData::F32(buf) => (cast_slice_to_cow(buf), TextureFormat::Rgba32Float), + TensorData::F64(buf) => (narrow_f64_to_f32s(buf), TextureFormat::Rgba32Float), // narrowing to f32! + + TensorData::JPEG(_) => { + anyhow::bail!("JPEGs should have been decoded at this point") + } + } + } + depth => { + anyhow::bail!("Cannot create texture from tensor of depth {depth}"); } }; @@ -331,7 +382,7 @@ fn general_texture_creation_desc_from_tensor<'a>( }) } -fn get_or_create_texture<'a, Err>( +pub fn get_or_create_texture<'a, Err>( render_ctx: &mut RenderContext, texture_key: u64, try_create_texture_desc: impl FnOnce() -> Result, Err>, @@ -347,6 +398,26 @@ fn cast_slice_to_cow(slice: &[From]) -> Cow<'_, [u8]> { cast_slice(slice).into() } +// wgpu doesn't support u64 textures, so we need to narrow to f32: +fn narrow_u64_to_f32s(slice: &[u64]) -> Cow<'static, [u8]> { + crate::profile_function!(); + let bytes: Vec = slice + .iter() + .flat_map(|&f| (f as f32).to_le_bytes()) + .collect(); + bytes.into() +} + +// wgpu doesn't support i64 textures, so we need to narrow to f32: +fn narrow_i64_to_f32s(slice: &[i64]) -> Cow<'static, [u8]> { + crate::profile_function!(); + let bytes: Vec = slice + .iter() + .flat_map(|&f| (f as f32).to_le_bytes()) + .collect(); + bytes.into() +} + // wgpu doesn't support f64 textures, so we need to narrow to f32: fn narrow_f64_to_f32s(slice: &[f64]) -> Cow<'static, [u8]> { crate::profile_function!(); @@ -371,6 +442,20 @@ fn pad_and_cast(data: &[T], pad: T) -> Cow<'static, [u8]> { bytes.into() } +fn pad_and_narrow_and_cast( + data: &[T], + pad: f32, + narrow: impl Fn(T) -> f32, +) -> Cow<'static, [u8]> { + crate::profile_function!(); + + let floats: Vec = data + .chunks_exact(3) + .flat_map(|chunk| [narrow(chunk[0]), narrow(chunk[1]), narrow(chunk[2]), pad]) + .collect(); + pod_collect_to_vec(&floats).into() +} + // ----------------------------------------------------------------------------; fn height_width_depth(tensor: &Tensor) -> anyhow::Result<[u32; 3]> { diff --git a/examples/python/api_demo/main.py b/examples/python/api_demo/main.py index daaa448677da..aff55a99bab8 100755 --- a/examples/python/api_demo/main.py +++ b/examples/python/api_demo/main.py @@ -12,7 +12,9 @@ import argparse import logging import math +import os +import cv2 import numpy as np import rerun as rr from scipy.spatial.transform import Rotation @@ -277,23 +279,57 @@ def run_extension_component() -> None: ) +def run_image_tensors() -> None: + # Make sure you use a colorful image with alpha! + dir_path = os.path.dirname(os.path.realpath(__file__)) + img_path = f"{dir_path}/../../../crates/re_ui/data/logo_dark_mode.png" + img_bgra = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) + + img_rgba = cv2.cvtColor(img_bgra, cv2.COLOR_BGRA2RGBA) + rr.log_image("img_rgba", img_rgba) + img_rgb = cv2.cvtColor(img_rgba, cv2.COLOR_RGBA2RGB) + rr.log_image("img_rgb", img_rgb) + img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY) + rr.log_image("img_gray", img_gray) + + dtypes = [ + "uint8", + "uint16", + "uint32", + "uint64", + "int8", + "int16", + "int32", + "int64", + "float16", + "float32", + "float64", + ] + + for dtype in dtypes: + rr.log_image(f"img_rgba_{dtype}", img_rgba.astype(dtype)) + rr.log_image(f"img_rgb_{dtype}", img_rgb.astype(dtype)) + rr.log_image(f"img_gray_{dtype}", img_gray.astype(dtype)) + + def main() -> None: demos = { "2d_lines": run_2d_lines, "3d_points": run_3d_points, + "bbox": run_bounding_box, + "extension_components": run_extension_component, + "image_tensors": run_image_tensors, "log_cleared": run_log_cleared, "raw_mesh": raw_mesh, "rects": run_rects, "segmentation": run_segmentation, "text": run_text_logs, "transforms_3d": transforms_rigid_3d, - "bbox": run_bounding_box, - "extension_components": run_extension_component, } parser = argparse.ArgumentParser(description="Logs rich data using the Rerun SDK.") parser.add_argument( - "--demo", type=str, default="all", help="What demo to run", choices=["all"] + list(demos.keys()) + "--demo", type=str, default="most", help="What demo to run", choices=["most", "all"] + list(demos.keys()) ) rr.script_add_args(parser) @@ -301,9 +337,13 @@ def main() -> None: rr.script_setup(args, "api_demo") - if args.demo == "all": - print("Running all demos…") + if args.demo in ["most", "all"]: + print(f"Running {args.demo} demos…") for name, demo in demos.items(): + # Some demos are just a bit… too much + if args.demo == "most" and name in ["image_tensors"]: + continue + logging.info(f"Starting {name}") demo() else: diff --git a/examples/python/api_demo/requirements.txt b/examples/python/api_demo/requirements.txt index 8349c316dcbb..49dd44f9f951 100644 --- a/examples/python/api_demo/requirements.txt +++ b/examples/python/api_demo/requirements.txt @@ -1,3 +1,4 @@ numpy +opencv-python rerun-sdk scipy diff --git a/examples/python/arkitscenes/requirements.txt b/examples/python/arkitscenes/requirements.txt index 132f56a29b4d..0629c41c76b3 100644 --- a/examples/python/arkitscenes/requirements.txt +++ b/examples/python/arkitscenes/requirements.txt @@ -1,7 +1,7 @@ -rerun-sdk numpy -pandas opencv-python -tqdm +pandas +rerun-sdk scipy +tqdm trimesh diff --git a/scripts/run_python_e2e_test.py b/scripts/run_python_e2e_test.py index 78c03ed9c7e4..e654092bf0d8 100755 --- a/scripts/run_python_e2e_test.py +++ b/scripts/run_python_e2e_test.py @@ -16,6 +16,7 @@ import subprocess import sys import time +from typing import List def main() -> None: @@ -39,17 +40,17 @@ def main() -> None: examples = [ # Trivial examples that don't require weird dependencies, or downloading data - "examples/python/api_demo/main.py", - "examples/python/car/main.py", - "examples/python/multithreading/main.py", - "examples/python/plots/main.py", - "examples/python/text_logging/main.py", + ("examples/python/api_demo/main.py", ["--demo", "all"]), + ("examples/python/car/main.py", []), + ("examples/python/multithreading/main.py", []), + ("examples/python/plots/main.py", []), + ("examples/python/text_logging/main.py", []), ] - for example in examples: + for example, args in examples: print("----------------------------------------------------------") print(f"Testing {example}…\n") start_time = time.time() - run_example(example) + run_example(example, args) elapsed = time.time() - start_time print(f"{example} done in {elapsed:.1f} seconds") print() @@ -58,7 +59,7 @@ def main() -> None: print("All tests passed successfully!") -def run_example(example: str) -> None: +def run_example(example: str, args: List[str]) -> None: port = 9752 # sys.executable: the absolute path of the executable binary for the Python interpreter @@ -71,7 +72,7 @@ def run_example(example: str) -> None: ) time.sleep(0.3) # Wait for rerun server to start to remove a logged warning - python_process = subprocess.Popen([python_executable, example, "--connect", "--addr", f"127.0.0.1:{port}"]) + python_process = subprocess.Popen([python_executable, example, "--connect", "--addr", f"127.0.0.1:{port}"] + args) print("Waiting for python process to finish…") returncode = python_process.wait(timeout=30)