Skip to content

Commit

Permalink
Port SegmentationImage to the new image archetype style (#6928)
Browse files Browse the repository at this point in the history
### What
* Part of #6386
* Part of #6844
* Very similar to #6915

### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested the web demo (if applicable):
* Using examples from latest `main` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/6928?manifest_url=https://app.rerun.io/version/main/examples_manifest.json)
* Using full set of examples from `nightly` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/6928?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json)
* [x] The PR title and labels are set such as to maximize their
usefulness for the next release's CHANGELOG
* [x] If applicable, add a new check to the [release
checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)!
* [x] If have noted any breaking changes to the log API in
`CHANGELOG.md` and the migration guide

- [PR Build Summary](https://build.rerun.io/pr/6928)
- [Recent benchmark results](https://build.rerun.io/graphs/crates.html)
- [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)

To run all checks from `main`, comment on the PR with `@rerun-bot
full-check`.
  • Loading branch information
emilk authored Jul 18, 2024
1 parent 82abe40 commit a8bf4dd
Show file tree
Hide file tree
Showing 37 changed files with 414 additions and 534 deletions.
2 changes: 1 addition & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
* 3D transform APIs: Previously, the transform component was represented as one of several variants (an Arrow union, `enum` in Rust) depending on how the transform was expressed. Instead, there are now several components for translation/scale/rotation/matrices that can live side-by-side in the 3D transform archetype.
* Python: `NV12/YUY2` are now logged with the new `ImageChromaDownsampled`
* [`ImageEncoded`](https://rerun.io/docs/reference/types/archetypes/image_encoded?speculative-link):s `format` parameter has been replaced with `media_type` (MIME)
* [`DepthImage`](https://rerun.io/docs/reference/types/archetypes/depth_image) is no longer encoded as a tensor, and expects its shape in `[width, height]` order
* [`DepthImage`](https://rerun.io/docs/reference/types/archetypes/depth_image) and [`SegmentationImage`](https://rerun.io/docs/reference/types/archetypes/segmentation_image) are no longer encoded as a tensors, and expects its shape in `[width, height]` order

🧳 Migration guide: http://rerun.io/docs/reference/migration/migration-0-18?speculative-link

Expand Down
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
namespace rerun.archetypes;


/// A depth image.
/// A depth image, i.e. as captured by a depth camera.
///
/// The shape of the [components.TensorData] must be mappable to an `HxW` tensor.
/// Each pixel corresponds to a depth value in units specified by `meter`.
/// Each pixel corresponds to a depth value in units specified by [components.DepthMeter].
///
/// \cpp Since the underlying `rerun::datatypes::TensorData` uses `rerun::Collection` internally,
/// \cpp data can be passed in without a copy from raw pointers or by reference from `std::vector`/`std::array`/c-arrays.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,11 @@ namespace rerun.archetypes;

/// An image made up of integer [components.ClassId]s.
///
/// The shape of the [components.TensorData] must be mappable to an `HxW` tensor.
/// Each pixel corresponds to a [components.ClassId] that will be mapped to a color based on annotation context.
///
/// In the case of floating point images, the label will be looked up based on rounding to the nearest
/// integer value.
///
/// Leading and trailing unit-dimensions are ignored, so that
/// `1x640x480x1` is treated as a `640x480` image.
///
/// See also [archetypes.AnnotationContext] to associate each class with a color and a label.
///
/// \cpp Since the underlying `rerun::datatypes::TensorData` uses `rerun::Collection` internally,
Expand All @@ -27,8 +23,14 @@ table SegmentationImage (
) {
// --- Required ---

/// The image data. Should always be a 2-dimensional tensor.
data: rerun.components.TensorData ("attr.rerun.component_required", order: 1000);
/// The raw image data.
data: rerun.components.Blob ("attr.rerun.component_required", order: 1000);

/// The size of the image.
resolution: rerun.components.Resolution2D ("attr.rerun.component_required", order: 1500);

/// The data type of the segmentation image data (U16, U32, …).
data_type: rerun.components.ChannelDataType ("attr.rerun.component_required", order: 2000);

// --- Optional ---

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@ namespace rerun.components;
/// Specified what color components are present in an [archetypes.Image].
///
/// This combined with [components.ChannelDataType] determines the pixel format of an image.
enum ColorModel: byte {
enum ColorModel: byte (
"attr.docs.unreleased"
) {
/// Grayscale luminance intencity/brightness/value, sometimes called `Y`
L (default),

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ namespace rerun.components;

/// The width and height of a 2D image.
struct Resolution2D (
"attr.docs.unreleased",
"attr.python.aliases": "npt.NDArray[np.int], Sequence[int], Tuple[int, int]",
"attr.python.array_aliases": "npt.NDArray[np.int], Sequence[int]",
"attr.rust.derive": "Default, Copy, PartialEq, Eq, Hash, bytemuck::Pod, bytemuck::Zeroable",
Expand Down
5 changes: 2 additions & 3 deletions crates/store/re_types/src/archetypes/depth_image.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

82 changes: 66 additions & 16 deletions crates/store/re_types/src/archetypes/segmentation_image.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

80 changes: 33 additions & 47 deletions crates/store/re_types/src/archetypes/segmentation_image_ext.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use crate::{
datatypes::TensorData,
components::{ChannelDataType, Resolution2D},
datatypes::{Blob, TensorBuffer, TensorData},
image::{find_non_empty_dim_indices, ImageConstructionError},
};

Expand All @@ -16,62 +17,47 @@ impl SegmentationImage {
where
<T as TryInto<TensorData>>::Error: std::error::Error,
{
let mut data: TensorData = data
let tensor_data: TensorData = data
.try_into()
.map_err(ImageConstructionError::TensorDataConversion)?;

let non_empty_dim_inds = find_non_empty_dim_indices(&data.shape);
let non_empty_dim_inds = find_non_empty_dim_indices(&tensor_data.shape);

match non_empty_dim_inds.len() {
2 => {
assign_if_none(&mut data.shape[non_empty_dim_inds[0]].name, "height");
assign_if_none(&mut data.shape[non_empty_dim_inds[1]].name, "width");
if non_empty_dim_inds.len() != 2 {
return Err(ImageConstructionError::BadImageShape(tensor_data.shape));
}

let (blob, data_type) = match tensor_data.buffer {
TensorBuffer::U8(buffer) => (Blob(buffer), ChannelDataType::U8),
TensorBuffer::U16(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::U16),
TensorBuffer::U32(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::U32),
TensorBuffer::U64(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::U64),
TensorBuffer::I8(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::I8),
TensorBuffer::I16(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::I16),
TensorBuffer::I32(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::I32),
TensorBuffer::I64(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::I64),
TensorBuffer::F16(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::F16),
TensorBuffer::F32(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::F32),
TensorBuffer::F64(buffer) => (Blob(buffer.cast_to_u8()), ChannelDataType::F64),
TensorBuffer::Nv12(_) | TensorBuffer::Yuy2(_) => {
return Err(ImageConstructionError::ChromaDownsamplingNotSupported);
}
_ => return Err(ImageConstructionError::BadImageShape(data.shape)),
};

let (height, width) = (
&tensor_data.shape[non_empty_dim_inds[0]],
&tensor_data.shape[non_empty_dim_inds[1]],
);
let height = height.size as u32;
let width = width.size as u32;
let resolution = Resolution2D::from([width, height]);

Ok(Self {
data: data.into(),
data: blob.into(),
resolution,
data_type,
draw_order: None,
opacity: None,
})
}
}

fn assign_if_none(name: &mut Option<::re_types_core::ArrowString>, new_name: &str) {
if name.is_none() {
*name = Some(new_name.into());
}
}

// ----------------------------------------------------------------------------
// Make it possible to create an ArrayView directly from an Image.

macro_rules! forward_array_views {
($type:ty, $alias:ty) => {
impl<'a> TryFrom<&'a $alias> for ::ndarray::ArrayViewD<'a, $type> {
type Error = crate::tensor_data::TensorCastError;

#[inline]
fn try_from(value: &'a $alias) -> Result<Self, Self::Error> {
(&value.data.0).try_into()
}
}
};
}

forward_array_views!(u8, SegmentationImage);
forward_array_views!(u16, SegmentationImage);
forward_array_views!(u32, SegmentationImage);
forward_array_views!(u64, SegmentationImage);

forward_array_views!(i8, SegmentationImage);
forward_array_views!(i16, SegmentationImage);
forward_array_views!(i32, SegmentationImage);
forward_array_views!(i64, SegmentationImage);

forward_array_views!(half::f16, SegmentationImage);
forward_array_views!(f32, SegmentationImage);
forward_array_views!(f64, SegmentationImage);

// ----------------------------------------------------------------------------
Loading

0 comments on commit a8bf4dd

Please sign in to comment.