Depth map support (#26)

* Avoid double-allocating when logging tensors

* Visualize single-channel float tensors by converting their range to u16

* Update to maturin 0.13.1

* remove --no-sdist flag from maturin build (it's now default)

* Show object type in context panel

* Fix bug when serializing events (missing type messages)

* Show TypeMsg in log table view

* Log depth image "meter" unit

* fix wasm32 compilation

.github/workflows/rust.yml

@@ -181,9 +181,9 @@ jobs:
       - uses: actions/checkout@v3
       - uses: messense/maturin-action@v1
         with:
-          maturin-version: "0.12.20"
+          maturin-version: "0.13.1"
           command: build
-          args: -m crates/re_sdk_python/Cargo.toml --release --no-sdist -o dist
+          args: -m crates/re_sdk_python/Cargo.toml --release -o dist
       - name: Upload wheels
         uses: actions/upload-artifact@v2
         with:
@@ -197,9 +197,9 @@ jobs:
       - uses: actions/checkout@v3
       - uses: messense/maturin-action@v1
         with:
-          maturin-version: "0.12.20"
+          maturin-version: "0.13.1"
           command: build
-          args: -m crates/re_sdk_python/Cargo.toml --release --no-sdist -o dist --universal2
+          args: -m crates/re_sdk_python/Cargo.toml --release -o dist --universal2
       - name: Upload wheels
         uses: actions/upload-artifact@v2
         with:
@@ -220,6 +220,6 @@ jobs:
         env:
           MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
         with:
-          maturin-version: "0.12.20"
+          maturin-version: "0.13.1"
           command: upload
           args: -m crates/re_sdk_python/Cargo.toml --skip-existing *

crates/re_data_store/src/objects.rs

@@ -70,6 +70,13 @@ impl<'s, T: Clone + Copy + std::fmt::Debug> ObjectVec<'s, T> {
 #[derive(Copy, Clone, Debug)]
 pub struct Image<'s> {
     pub tensor: &'s re_log_types::Tensor,
+
+    /// If this is a depth map, how long is a meter?
+    ///
+    /// For instance, with a `u16` dtype one might have
+    /// `meter == 1000.0` for millimeter precision
+    /// up to a ~65m range.
+    pub meter: Option<f32>,
 }
 
 impl<'s> Image<'s> {
@@ -80,24 +87,28 @@ impl<'s> Image<'s> {
     ) {
         crate::profile_function!();
 
-        visit_type_data_2(
+        visit_type_data_3(
             obj_store,
             &FieldName::from("tensor"),
             time_query,
-            ("space", "color"),
+            ("space", "color", "meter"),
             |obj_path: &ObjPath,
              log_id: &LogId,
              tensor: &re_log_types::Tensor,
              space: Option<&ObjPath>,
-             color: Option<&[u8; 4]>| {
+             color: Option<&[u8; 4]>,
+             meter: Option<&f32>| {
                 out.image.0.push(Object {
                     props: ObjectProps {
                         log_id,
                         space,
                         color: color.copied(),
                         obj_path,
                     },
-                    data: Image { tensor },
+                    data: Image {
+                        tensor,
+                        meter: meter.copied(),
+                    },
                 });
             },
         );

crates/re_log_types/src/objects.rs

@@ -23,7 +23,7 @@ impl ObjectType {
         match self {
             Self::Space => &["up"],
 
-            Self::Image => &["space", "color", "tensor"],
+            Self::Image => &["space", "color", "tensor", "meter"],
             Self::Point2D => &["space", "color", "pos", "radius"],
             Self::BBox2D => &["space", "color", "bbox", "stroke_width"],
             Self::LineSegments2D => &["space", "color", "line_segments", "stroke_width"],

crates/re_sdk_python/pyproject.toml

@@ -1,6 +1,6 @@
 [build-system]
 build-backend = "maturin"
-requires = ["maturin>=0.12,<0.13"]
+requires = ["maturin==0.13,<0.14"]
 
 [project]
 classifiers = [

crates/re_sdk_python/python/rerun_sdk/__init__.py

@@ -6,7 +6,7 @@
 print("rerun_sdk initialized")
 
 
-def log_points(name, positions, colors):
+def log_points(obj_path, positions, colors):
     if colors is not None:
         # Rust expects colors in 0-255 uint8
         if colors.dtype in ['float32', 'float64']:
@@ -21,10 +21,10 @@ def log_points(name, positions, colors):
 
     positions.astype('float32')
 
-    log_points_rs(name, positions, colors)
+    log_points_rs(obj_path, positions, colors)
 
 
-def log_image(name, image):
+def log_image(obj_path, image):
     # Catch some errors early:
     if len(image.shape) < 2 or 3 < len(image.shape):
         raise TypeError(f"Expected image, got array of shape {image.shape}")
@@ -33,19 +33,36 @@ def log_image(name, image):
         depth = image.shape[2]
         if depth not in (1, 3, 4):
             raise TypeError(
-                f"Expected image depth of  of 1 (gray), 3 (RGB) or 4 (RGBA), got array of shape {image.shape}")
+                f"Expected image depth of 1 (gray), 3 (RGB) or 4 (RGBA). Instead got array of shape {image.shape}")
 
-    log_tensor(name, image)
+    log_tensor(obj_path, image)
 
 
-def log_tensor(name, image):
+def log_depth_image(obj_path, image, meter=None):
+    """
+    meter: How long is a meter in the given dtype?
+           For instance: with uint16, perhaps meter=1000 which would mean
+           you have millimeter precision and a range of up to ~65 meters (2^16 / 1000).
+    """
+    # Catch some errors early:
+    if len(image.shape) != 2:
+        raise TypeError(
+            f"Expected 2D depth image, got array of shape {image.shape}")
+
+    log_tensor(obj_path, image)
+
+    if meter != None:
+        log_f32(obj_path, "meter", meter)
+
+
+def log_tensor(obj_path, image):
     if image.dtype == 'uint8':
-        log_tensor_u8(name, image)
+        log_tensor_u8(obj_path, image)
     elif image.dtype == 'uint16':
-        log_tensor_u16(name, image)
+        log_tensor_u16(obj_path, image)
     elif image.dtype == 'float32':
-        log_tensor_f32(name, image)
+        log_tensor_f32(obj_path, image)
     elif image.dtype == 'float64':
-        log_tensor_f32(name, image.astype('float32'))
+        log_tensor_f32(obj_path, image.astype('float32'))
     else:
         raise TypeError(f"Unsupported dtype: {image.dtype}")

crates/re_sdk_python/requirements.txt

@@ -1,3 +1,3 @@
-maturin
-numpy
-opencv-python
+maturin==0.13.1
+numpy==1.23.1
+opencv-python==4.6.0.66

crates/re_sdk_python/src/python_bridge.rs

@@ -22,6 +22,7 @@ fn rerun_sdk(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
         m.add_function(wrap_pyfunction!(show, m)?)?;
     }
 
+    m.add_function(wrap_pyfunction!(log_f32, m)?)?;
     m.add_function(wrap_pyfunction!(log_point2d, m)?)?;
     m.add_function(wrap_pyfunction!(log_points_rs, m)?)?;
 
@@ -68,10 +69,22 @@ fn show() {
 }
 
 #[pyfunction]
-fn log_point2d(name: &str, x: f32, y: f32) {
+fn log_f32(obj_path: &str, field_name: &str, value: f32) {
+    let obj_path = ObjPath::from(obj_path); // TODO(emilk): pass in proper obj path somehow
     let mut sdk = Sdk::global();
+    sdk.send(LogMsg::DataMsg(DataMsg {
+        id: LogId::random(),
+        time_point: time_point(),
+        data_path: DataPath::new(obj_path, field_name.into()),
+        data: re_log_types::LoggedData::Single(Data::F32(value)),
+    }));
+}
 
-    let obj_path = ObjPath::from(name); // TODO(emilk): pass in proper obj path somehow
+#[pyfunction]
+fn log_point2d(obj_path: &str, x: f32, y: f32) {
+    let mut sdk = Sdk::global();
+
+    let obj_path = ObjPath::from(obj_path); // TODO(emilk): pass in proper obj path somehow
     sdk.register_type(obj_path.obj_type_path(), ObjectType::Point2D);
     let data_path = DataPath::new(obj_path, "pos".into());
 
@@ -91,7 +104,7 @@ fn log_point2d(name: &str, x: f32, y: f32) {
 /// * `colors.len() == positions.len()`: a color per point
 #[pyfunction]
 fn log_points_rs(
-    name: &str,
+    obj_path: &str,
     positions: numpy::PyReadonlyArray2<'_, f64>,
     colors: numpy::PyReadonlyArray2<'_, u8>,
 ) -> PyResult<()> {
@@ -111,7 +124,7 @@ fn log_points_rs(
 
     let mut sdk = Sdk::global();
 
-    let root_path = ObjPathBuilder::from(name); // TODO(emilk): pass in proper obj path somehow
+    let root_path = ObjPathBuilder::from(obj_path); // TODO(emilk): pass in proper obj path somehow
     let point_path = ObjPath::from(&root_path / ObjPathComp::Index(Index::Placeholder));
 
     let mut type_path = root_path.obj_type_path();
@@ -219,29 +232,29 @@ fn log_points_rs(
 
 #[allow(clippy::needless_pass_by_value)]
 #[pyfunction]
-fn log_tensor_u8(name: &str, img: numpy::PyReadonlyArrayDyn<'_, u8>) {
-    log_tensor(name, img);
+fn log_tensor_u8(obj_path: &str, img: numpy::PyReadonlyArrayDyn<'_, u8>) {
+    log_tensor(obj_path, img);
 }
 
 #[allow(clippy::needless_pass_by_value)]
 #[pyfunction]
-fn log_tensor_u16(name: &str, img: numpy::PyReadonlyArrayDyn<'_, u16>) {
-    log_tensor(name, img);
+fn log_tensor_u16(obj_path: &str, img: numpy::PyReadonlyArrayDyn<'_, u16>) {
+    log_tensor(obj_path, img);
 }
 
 #[allow(clippy::needless_pass_by_value)]
 #[pyfunction]
-fn log_tensor_f32(name: &str, img: numpy::PyReadonlyArrayDyn<'_, f32>) {
-    log_tensor(name, img);
+fn log_tensor_f32(obj_path: &str, img: numpy::PyReadonlyArrayDyn<'_, f32>) {
+    log_tensor(obj_path, img);
 }
 
 fn log_tensor<T: TensorDataTypeTrait + numpy::Element + bytemuck::Pod>(
-    name: &str,
+    obj_path: &str,
     img: numpy::PyReadonlyArrayDyn<'_, T>,
 ) {
     let mut sdk = Sdk::global();
 
-    let obj_path = ObjPath::from(name); // TODO(emilk): pass in proper obj path somehow
+    let obj_path = ObjPath::from(obj_path); // TODO(emilk): pass in proper obj path somehow
     sdk.register_type(obj_path.obj_type_path(), ObjectType::Image);
 
     let data = Data::Tensor(to_rerun_tensor(&img));
@@ -267,12 +280,13 @@ fn time_point() -> TimePoint {
 fn to_rerun_tensor<T: TensorDataTypeTrait + numpy::Element + bytemuck::Pod>(
     img: &numpy::PyReadonlyArrayDyn<'_, T>,
 ) -> re_log_types::Tensor {
+    let vec = img.to_owned_array().into_raw_vec();
+    let vec = bytemuck::allocation::try_cast_vec(vec)
+        .unwrap_or_else(|(_err, vec)| bytemuck::allocation::pod_collect_to_vec(&vec));
+
     re_log_types::Tensor {
         shape: img.shape().iter().map(|&d| d as u64).collect(),
         dtype: T::DTYPE,
-        // TODO(emilk): avoid double-allocating here
-        data: TensorData::Dense(bytemuck::allocation::pod_collect_to_vec(
-            &img.to_owned_array().into_raw_vec(),
-        )),
+        data: TensorData::Dense(vec),
     }
 }

crates/re_sdk_python/test.py

@@ -15,14 +15,14 @@ def log(args):
 
     print(rerun.info())
 
-    if True:
+    if False:
         image = cv2.imread('crates/re_viewer/data/logo_dark_mode.png',
                            cv2.IMREAD_UNCHANGED)
         rerun.log_image("logo", image)
 
-    if False:
+    if True:
         depth_img = cv2.imread('depth_image.pgm', cv2.IMREAD_UNCHANGED)
-        rerun.log_image("depth", depth_img)
+        rerun.log_depth_image("depth", depth_img, meter=1000)
 
     if False:
         for i in range(64):
@@ -33,17 +33,20 @@ def log(args):
             rerun.log_point2d(f"point2d_{i}", x, y)
 
     if False:
-        pos3 = []
+        positions = []
         for i in range(1000):
             angle = 6.28 * i / 64
             r = 1.0
             x = r * math.cos(angle) + 18.0
             y = r * math.sin(angle) + 16.0
             z = i / 64.0
-            pos3.append([x, y, z])
-        pos3 = np.array(pos3)
+            positions.append([x, y, z])
+        positions = np.array(positions)
+
+        # Same for all points, but you can also have a different color for each point:
         colors = np.array([[200, 0, 100, 200]])
-        rerun.log_points(f"point3d", pos3, colors)
+
+        rerun.log_points(f"point3d", positions, colors)
 
     if args.connect:
         time.sleep(1.0)  # HACK: give rerun time to send it all

crates/re_viewer/src/misc/image_cache.rs

@@ -84,6 +84,44 @@ fn tensor_to_dynamic_image(tensor: &Tensor) -> anyhow::Result<DynamicImage> {
                 image::RgbaImage::from_raw(width, height, bytes.clone())
                     .context("Bad Rgba8")
                     .map(DynamicImage::ImageRgba8)
+            } else if depth == 1 && tensor.dtype == TensorDataType::F32 {
+                // Maybe a depth map?
+                if let TensorData::Dense(bytes) = &tensor.data {
+                    if let Ok(floats) = bytemuck::try_cast_slice(bytes) {
+                        // 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 dept 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 floats {
+                            min = min.min(float);
+                            max = max.max(float);
+                        }
+
+                        if min < max && min.is_finite() && max.is_finite() {
+                            let ints = floats
+                                .iter()
+                                .map(|&float| {
+                                    let int = egui::remap(float, min..=max, 0.0..=65535.0);
+                                    int as u16
+                                })
+                                .collect();
+
+                            return Gray16Image::from_raw(width, height, ints)
+                                .context("Bad Luminance16")
+                                .map(DynamicImage::ImageLuma16);
+                        }
+                    }
+                }
+
+                anyhow::bail!(
+                    "Don't know how to turn a tensor of shape={:?} and dtype={:?} into an image",
+                    shape,
+                    tensor.dtype
+                )
             } else {
                 anyhow::bail!(
                     "Don't know how to turn a tensor of shape={:?} and dtype={:?} into an image",