wehub-resource-sync
207 lines
7.3 KiB
Python
207 lines
7.3 KiB
Python
from __future__ import annotations
|
|
|
|
import itertools
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import rerun as rr
|
|
import torch
|
|
from rerun.components import DepthMeter, ImageFormat
|
|
from rerun.datatypes import ChannelDatatype, Float32Like
|
|
from rerun.error_utils import RerunWarning
|
|
|
|
rng = np.random.default_rng(12345)
|
|
RANDOM_IMAGE_SOURCE = rng.uniform(0.0, 1.0, (10, 20))
|
|
|
|
IMAGE_INPUTS: list[Any] = [
|
|
RANDOM_IMAGE_SOURCE,
|
|
RANDOM_IMAGE_SOURCE,
|
|
]
|
|
|
|
METER_INPUTS: list[Float32Like] = [1000, DepthMeter(1000)]
|
|
|
|
|
|
def depth_image_expected() -> Any:
|
|
return rr.DepthImage(RANDOM_IMAGE_SOURCE, meter=1000)
|
|
|
|
|
|
def test_depth_image() -> None:
|
|
ranges = [None, [0.0, 1.0], (1000, 1000)]
|
|
|
|
for img, meter, depth_range in itertools.zip_longest(IMAGE_INPUTS, METER_INPUTS, ranges):
|
|
if img is None:
|
|
img = IMAGE_INPUTS[0]
|
|
|
|
print(
|
|
f"rr.DepthImage(\n {img}\n meter={meter!r}\n depth_range={depth_range!r}\n)",
|
|
)
|
|
arch = rr.DepthImage(img, meter=meter, depth_range=depth_range)
|
|
|
|
assert arch.buffer == rr.components.ImageBufferBatch._converter(img.tobytes())
|
|
assert arch.format == rr.components.ImageFormatBatch._converter(
|
|
ImageFormat(
|
|
width=img.shape[1],
|
|
height=img.shape[0],
|
|
channel_datatype=ChannelDatatype.from_np_dtype(img.dtype),
|
|
),
|
|
)
|
|
assert arch.meter == rr.components.DepthMeterBatch._converter(meter)
|
|
assert arch.depth_range == rr.components.ValueRangeBatch._converter(depth_range)
|
|
|
|
|
|
GOOD_IMAGE_INPUTS: list[Any] = [
|
|
# Mono
|
|
rng.uniform(0.0, 1.0, (10, 20)),
|
|
# Assorted Extra Dimensions
|
|
rng.uniform(0.0, 1.0, (1, 10, 20)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 1)),
|
|
torch.rand(10, 20, 1),
|
|
]
|
|
|
|
BAD_IMAGE_INPUTS: list[Any] = [
|
|
rng.uniform(0.0, 1.0, (10, 20, 3)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 4)),
|
|
rng.uniform(0.0, 1.0, (10,)),
|
|
rng.uniform(0.0, 1.0, (1, 10, 20, 3)),
|
|
rng.uniform(0.0, 1.0, (1, 10, 20, 4)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 3, 1)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 4, 1)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 2)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 5)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 3, 2)),
|
|
]
|
|
|
|
|
|
def test_depth_image_shapes() -> None:
|
|
import rerun as rr
|
|
|
|
rr.set_strict_mode(True)
|
|
|
|
for img in GOOD_IMAGE_INPUTS:
|
|
rr.DepthImage(img)
|
|
|
|
for img in BAD_IMAGE_INPUTS:
|
|
with pytest.raises(ValueError):
|
|
rr.DepthImage(img)
|
|
|
|
|
|
def _compressed_blob_size(encoded_depth: Any) -> int:
|
|
"""Extract the byte size of the PNG blob from an EncodedDepthImage."""
|
|
return len(encoded_depth.blob.as_arrow_array()[0].as_py())
|
|
|
|
|
|
def test_depth_image_compress() -> None:
|
|
rr.set_strict_mode(False)
|
|
|
|
# U16 supported (most common depth format)
|
|
depth_data = np.asarray(rng.uniform(0, 65535, (10, 20)), dtype=np.uint16)
|
|
compressed = rr.DepthImage(depth_data, meter=1000).compress()
|
|
assert type(compressed) is rr.EncodedDepthImage
|
|
|
|
# U8 supported
|
|
depth_data = np.asarray(rng.uniform(0, 255, (10, 20)), dtype=np.uint8)
|
|
compressed = rr.DepthImage(depth_data).compress()
|
|
assert type(compressed) is rr.EncodedDepthImage
|
|
|
|
# F32 not supported
|
|
with pytest.warns(RerunWarning) as warnings:
|
|
depth_data = np.asarray(rng.uniform(0, 1, (10, 20)), dtype=np.float32)
|
|
compressed = rr.DepthImage(depth_data).compress()
|
|
|
|
assert len(warnings) == 1
|
|
assert "Cannot PNG compress a depth image of datatype" in str(warnings[0])
|
|
assert type(compressed) is rr.DepthImage
|
|
|
|
# U32 not supported
|
|
with pytest.warns(RerunWarning) as warnings:
|
|
depth_data = np.asarray(rng.uniform(0, 65535, (10, 20)), dtype=np.uint32)
|
|
compressed = rr.DepthImage(depth_data).compress()
|
|
|
|
assert len(warnings) == 1
|
|
assert "Cannot PNG compress a depth image of datatype" in str(warnings[0])
|
|
assert type(compressed) is rr.DepthImage
|
|
|
|
|
|
def test_depth_image_compress_reduces_size() -> None:
|
|
"""Verify that PNG compression actually reduces data size for realistic depth images."""
|
|
rr.set_strict_mode(True)
|
|
|
|
# Smooth gradient (simulates a flat wall receding) — highly compressible
|
|
rows, cols = 480, 640
|
|
gradient_u16 = np.tile(np.linspace(500, 10000, cols, dtype=np.uint16), (rows, 1))
|
|
raw_size = gradient_u16.nbytes
|
|
compressed = rr.DepthImage(gradient_u16, meter=1000).compress()
|
|
assert type(compressed) is rr.EncodedDepthImage
|
|
compressed_size = _compressed_blob_size(compressed)
|
|
assert compressed_size < raw_size, f"PNG should be smaller than raw for a gradient: {compressed_size} >= {raw_size}"
|
|
|
|
# Constant depth (e.g. flat floor) — maximally compressible
|
|
constant_u16 = np.full((rows, cols), 3000, dtype=np.uint16)
|
|
raw_size = constant_u16.nbytes
|
|
compressed = rr.DepthImage(constant_u16).compress()
|
|
compressed_size = _compressed_blob_size(compressed)
|
|
assert compressed_size < raw_size, (
|
|
f"PNG should be smaller than raw for constant data: {compressed_size} >= {raw_size}"
|
|
)
|
|
# Constant data should compress very aggressively (>90% reduction)
|
|
assert compressed_size < raw_size * 0.1, (
|
|
f"Constant image should compress to <10% of raw: {compressed_size} vs {raw_size}"
|
|
)
|
|
|
|
# U8 gradient
|
|
gradient_u8 = np.tile(np.linspace(0, 255, cols, dtype=np.uint8), (rows, 1))
|
|
raw_size = gradient_u8.nbytes
|
|
compressed = rr.DepthImage(gradient_u8).compress()
|
|
compressed_size = _compressed_blob_size(compressed)
|
|
assert compressed_size < raw_size, (
|
|
f"PNG should be smaller than raw for U8 gradient: {compressed_size} >= {raw_size}"
|
|
)
|
|
|
|
# Stepped depth (simulates discrete depth planes) — should compress well
|
|
stepped_u16 = np.zeros((rows, cols), dtype=np.uint16)
|
|
for i in range(4):
|
|
stepped_u16[i * (rows // 4) : (i + 1) * (rows // 4), :] = 1000 * (i + 1)
|
|
raw_size = stepped_u16.nbytes
|
|
compressed = rr.DepthImage(stepped_u16).compress()
|
|
compressed_size = _compressed_blob_size(compressed)
|
|
assert compressed_size < raw_size, (
|
|
f"PNG should be smaller than raw for stepped data: {compressed_size} >= {raw_size}"
|
|
)
|
|
|
|
|
|
def test_depth_image_compress_level() -> None:
|
|
"""Verify that compress_level parameter affects output size."""
|
|
rr.set_strict_mode(True)
|
|
|
|
rows, cols = 480, 640
|
|
gradient_u16 = np.tile(np.linspace(500, 10000, cols, dtype=np.uint16), (rows, 1))
|
|
|
|
size_level_0 = _compressed_blob_size(rr.DepthImage(gradient_u16).compress(compress_level=0))
|
|
size_level_9 = _compressed_blob_size(rr.DepthImage(gradient_u16).compress(compress_level=9))
|
|
|
|
assert size_level_9 < size_level_0, (
|
|
f"Level 9 should produce smaller output than level 0: {size_level_9} >= {size_level_0}"
|
|
)
|
|
|
|
|
|
def test_depth_image_compress_preserves_fields() -> None:
|
|
rr.set_strict_mode(True)
|
|
|
|
depth_data = np.asarray(rng.uniform(0, 65535, (10, 20)), dtype=np.uint16)
|
|
original = rr.DepthImage(
|
|
depth_data,
|
|
meter=1000,
|
|
depth_range=[100.0, 60000.0],
|
|
point_fill_ratio=0.5,
|
|
draw_order=1.0,
|
|
)
|
|
compressed = original.compress()
|
|
|
|
assert type(compressed) is rr.EncodedDepthImage
|
|
assert compressed.meter is not None
|
|
assert compressed.depth_range is not None
|
|
assert compressed.point_fill_ratio is not None
|
|
assert compressed.draw_order is not None
|
|
assert compressed.media_type is not None
|