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

Test and handle all tensor dtypes as images #1840

Merged
merged 12 commits into from
Apr 13, 2023
Merged

Test and handle all tensor dtypes as images #1840

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
dcfbc1a
Support i64 and u64 tensors
emilk Apr 13, 2023
daada52
api_demo: log all image types
emilk Apr 13, 2023
01f5def
Don't even print out the contents of a tensor
emilk Apr 13, 2023
573cbb8
Handle unfilterable float textures
emilk Apr 13, 2023
dbb4c73
fix typo
emilk Apr 13, 2023
1f02dac
py-format
emilk Apr 13, 2023
54d99b8
Simplify is_float_filterable
emilk Apr 13, 2023
128b895
Add a helper function pad_and_narrow_and_cast
emilk Apr 13, 2023
067a4a9
Reuse existing image
emilk Apr 13, 2023
9fe81a8
Exclude image_tensors demo from default api_demo
emilk Apr 13, 2023
9fe5ff0
Still run all api demos in e2e test
emilk Apr 13, 2023
f973e5f
pyformat
emilk Apr 13, 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
20 changes: 19 additions & 1 deletion crates/re_log_types/src/component_types/tensor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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
Expand Down
22 changes: 15 additions & 7 deletions crates/re_renderer/shader/rectangle.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -67,6 +68,9 @@ var texture_uint: texture_2d<u32>;
@group(1) @binding(5)
var colormap_texture: texture_2d<f32>;

@group(1) @binding(6)
var texture_float_filterable: texture_2d<f32>;


struct VertexOut {
@builtin(position) position: Vec4,
Expand All @@ -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 {
Expand Down
135 changes: 123 additions & 12 deletions crates/re_renderer/src/renderer/rectangles.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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,
},
Expand Down Expand Up @@ -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,
},
Expand Down Expand Up @@ -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,
},
],
},
);
Expand Down Expand Up @@ -653,3 +677,90 @@ impl Renderer for RectangleRenderer {
]
}
}

fn is_float_filterable(format: &wgpu::TextureFormat) -> bool {
// based on https://gpuweb.github.io/gpuweb/#texture-format-caps,
// TODO(emilk): add this to the wgpu crate
emilk marked this conversation as resolved.
Show resolved Hide resolved
match format {
// Norms:
wgpu::TextureFormat::R8Unorm |
wgpu::TextureFormat::R8Snorm |
wgpu::TextureFormat::R16Unorm |
wgpu::TextureFormat::R16Snorm |
wgpu::TextureFormat::Rg8Unorm |
wgpu::TextureFormat::Rg8Snorm |
wgpu::TextureFormat::Rg16Unorm |
wgpu::TextureFormat::Rg16Snorm |
wgpu::TextureFormat::Rgba8Unorm |
wgpu::TextureFormat::Rgba8UnormSrgb |
wgpu::TextureFormat::Rgba8Snorm |
wgpu::TextureFormat::Bgra8Unorm |
wgpu::TextureFormat::Bgra8UnormSrgb |
wgpu::TextureFormat::Rgb10a2Unorm |
wgpu::TextureFormat::Rgba16Snorm |
wgpu::TextureFormat::Rgba16Unorm |
wgpu::TextureFormat::Depth16Unorm |
wgpu::TextureFormat::Bc1RgbaUnorm |
wgpu::TextureFormat::Bc1RgbaUnormSrgb |
wgpu::TextureFormat::Bc2RgbaUnorm |
wgpu::TextureFormat::Bc2RgbaUnormSrgb |
wgpu::TextureFormat::Bc3RgbaUnorm |
wgpu::TextureFormat::Bc3RgbaUnormSrgb |
wgpu::TextureFormat::Bc4RUnorm |
wgpu::TextureFormat::Bc4RSnorm |
wgpu::TextureFormat::Bc5RgUnorm |
wgpu::TextureFormat::Bc5RgSnorm |
wgpu::TextureFormat::Bc7RgbaUnorm |
wgpu::TextureFormat::Bc7RgbaUnormSrgb |
wgpu::TextureFormat::Etc2Rgb8Unorm |
wgpu::TextureFormat::Etc2Rgb8UnormSrgb |
wgpu::TextureFormat::Etc2Rgb8A1Unorm |
wgpu::TextureFormat::Etc2Rgb8A1UnormSrgb |
wgpu::TextureFormat::Etc2Rgba8Unorm |
wgpu::TextureFormat::Etc2Rgba8UnormSrgb |
wgpu::TextureFormat::EacR11Unorm |
wgpu::TextureFormat::EacR11Snorm |
wgpu::TextureFormat::EacRg11Unorm |
wgpu::TextureFormat::EacRg11Snorm |
wgpu::TextureFormat::Astc { .. } |

// Floats:
wgpu::TextureFormat::R16Float |
wgpu::TextureFormat::Rg16Float |
wgpu::TextureFormat::Rgba16Float => true,

// Requires "float32-filterable" extension:
wgpu::TextureFormat::R32Float |
wgpu::TextureFormat::Rg32Float |
wgpu::TextureFormat::Rgba32Float |

// Non-floats:
wgpu::TextureFormat::R8Uint |
wgpu::TextureFormat::R8Sint |
wgpu::TextureFormat::R16Uint |
wgpu::TextureFormat::R16Sint |
wgpu::TextureFormat::Rg8Uint |
wgpu::TextureFormat::Rg8Sint |
wgpu::TextureFormat::R32Uint |
wgpu::TextureFormat::R32Sint |
wgpu::TextureFormat::Rg16Uint |
wgpu::TextureFormat::Rg16Sint |
wgpu::TextureFormat::Rgba8Uint |
wgpu::TextureFormat::Rgba8Sint |
wgpu::TextureFormat::Rgb9e5Ufloat |
wgpu::TextureFormat::Rg11b10Float |
wgpu::TextureFormat::Rg32Uint |
wgpu::TextureFormat::Rg32Sint |
wgpu::TextureFormat::Rgba16Uint |
wgpu::TextureFormat::Rgba16Sint |
wgpu::TextureFormat::Rgba32Uint |
wgpu::TextureFormat::Rgba32Sint |
wgpu::TextureFormat::Stencil8 |
wgpu::TextureFormat::Depth24Plus |
wgpu::TextureFormat::Depth24PlusStencil8 |
wgpu::TextureFormat::Depth32Float |
wgpu::TextureFormat::Depth32FloatStencil8 |
wgpu::TextureFormat::Bc6hRgbUfloat |
wgpu::TextureFormat::Bc6hRgbSfloat => false,
}
}
8 changes: 6 additions & 2 deletions crates/re_viewer/src/misc/caches/tensor_image_cache.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -328,8 +328,12 @@ fn color_tensor_as_color_image(tensor: &Tensor) -> anyhow::Result<ColorImage> {
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()
)
}
}
}
Expand Down
Loading