From 0e8f8c562d195dfb528a783bd09cf111a88534df Mon Sep 17 00:00:00 2001 From: Alan Liddell Date: Fri, 8 Mar 2024 17:15:43 -0500 Subject: [PATCH] Update bindings for Zarr driver --- python/acquire/acquire.pyi | 20 +-- src/storage.rs | 52 ++++++- tests/test_basic.py | 8 +- tests/test_zarr.py | 290 +++++++++++++++++++++++++++++++++---- 4 files changed, 327 insertions(+), 43 deletions(-) diff --git a/python/acquire/acquire.pyi b/python/acquire/acquire.pyi index fa50483..6d18d38 100644 --- a/python/acquire/acquire.pyi +++ b/python/acquire/acquire.pyi @@ -69,16 +69,6 @@ class Capabilities: def __init__(self, *args: None, **kwargs: Any) -> None: ... def dict(self) -> Dict[str, Any]: ... -@final -class StorageDimension: - name: str - kind: DimensionType - array_size_px: int - chunk_size_px: int - shard_size_chunks: int - - def dict(self) -> Dict[str, Any]: ... - @final class DimensionType: Space: ClassVar[DimensionType] @@ -323,6 +313,16 @@ class StorageCapabilities: def dict(self) -> Dict[str, Any]: ... +@final +class StorageDimension: + name: str + kind: DimensionType + array_size_px: int + chunk_size_px: int + shard_size_chunks: int + + def dict(self) -> Dict[str, Any]: ... + @final class StorageProperties: external_metadata_json: Optional[str] diff --git a/src/storage.rs b/src/storage.rs index 17621f1..aba74df 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -160,7 +160,7 @@ impl TryFrom for StorageProperties { }; let mut acquisition_dimensions: Vec> = Default::default(); - for i in 1..value.acquisition_dimensions.size { + for i in 0..value.acquisition_dimensions.size { acquisition_dimensions.push(Python::with_gil(|py| { Py::new( py, @@ -294,6 +294,27 @@ impl Default for capi::StorageProperties { } } +unsafe extern "C" fn init_storage_dimension_array( + data: *mut *mut capi::StorageDimension, + size: usize, +) { + if data.is_null() { + return; + } + + *data = std::alloc::alloc(std::alloc::Layout::array::(size).unwrap()) + as *mut capi::StorageDimension; +} + +unsafe extern "C" fn destroy_storage_dimension_array(data: *mut capi::StorageDimension) { + if !data.is_null() { + std::alloc::dealloc( + data as *mut u8, + std::alloc::Layout::array::(1).unwrap(), + ); + } +} + impl TryFrom<&StorageProperties> for capi::StorageProperties { type Error = anyhow::Error; @@ -338,13 +359,40 @@ impl TryFrom<&StorageProperties> for capi::StorageProperties { }, ) == 1 } { - Err(anyhow::anyhow!("Failed acquire api status check")) + Err(anyhow::anyhow!("Failed to initialize storage properties.")) } else if !unsafe { capi::storage_properties_set_enable_multiscale(&mut out, value.enable_multiscale as u8) == 1 } { Err(anyhow::anyhow!("Failed acquire api status check")) + } else if !unsafe { + if value.acquisition_dimensions.is_empty() { + true + } else { + out.acquisition_dimensions.init = Some(init_storage_dimension_array); + out.acquisition_dimensions.destroy = Some(destroy_storage_dimension_array); + capi::storage_properties_dimensions_init( + &mut out, + value.acquisition_dimensions.len(), + ) == 1 + } + } { + Err(anyhow::anyhow!( + "Failed to initialize storage dimension array." + )) } else { + // initialize each dimension separately + for (i, pydim) in value.acquisition_dimensions.iter().enumerate() { + let dim = Python::with_gil(|py| -> PyResult<_> { + let storage_dim: StorageDimension = pydim.extract(py)?; + Ok(storage_dim) + })?; + + unsafe { + *out.acquisition_dimensions.data.add(i) = (&dim).try_into()?; + } + } + Ok(out) } } diff --git a/tests/test_basic.py b/tests/test_basic.py index 73ce75e..cc28799 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -19,10 +19,10 @@ def runtime(): def test_version(): - assert isinstance(acquire.__version__, str) - # this will fail if pip install -e . has not been run - # so feel free to remove this line if it's not what you want to test - assert acquire.__version__ != "uninstalled" + assert isinstance(acquire.__version__, str) + # this will fail if pip install -e . has not been run + # so feel free to remove this line if it's not what you want to test + assert acquire.__version__ != "uninstalled" def test_set(): diff --git a/tests/test_zarr.py b/tests/test_zarr.py index b14638b..b3aaa53 100644 --- a/tests/test_zarr.py +++ b/tests/test_zarr.py @@ -22,6 +22,90 @@ def runtime(): yield acquire.Runtime() +def test_set_acquisition_dimensions( + runtime: Runtime, request: pytest.FixtureRequest +): + dm = runtime.device_manager() + props = runtime.get_configuration() + props.video[0].camera.identifier = dm.select( + DeviceKind.Camera, ".*empty.*" + ) + props.video[0].camera.settings.shape = (64, 48) + + props.video[0].storage.identifier = dm.select(DeviceKind.Storage, "Zarr") + props.video[0].storage.settings.filename = f"{request.node.name}.zarr" + props.video[0].max_frame_count = 32 + + # configure storage dimensions + dimension_x = acquire.StorageDimension() + dimension_x.name = "x" + dimension_x.kind = acquire.DimensionType.Space + dimension_x.array_size_px = 64 + dimension_x.chunk_size_px = 64 + assert dimension_x.shard_size_chunks == 0 + + dimension_y = acquire.StorageDimension() + dimension_y.name = "y" + dimension_y.kind = acquire.DimensionType.Space + dimension_y.array_size_px = 48 + dimension_y.chunk_size_px = 48 + assert dimension_y.shard_size_chunks == 0 + + dimension_t = acquire.StorageDimension() + dimension_t.name = "t" + dimension_t.kind = acquire.DimensionType.Time + dimension_t.array_size_px = 32 + dimension_t.chunk_size_px = 32 + assert dimension_t.shard_size_chunks == 0 + + props.video[0].storage.settings.acquisition_dimensions = [ + dimension_x, + dimension_y, + dimension_t, + ] + assert len(props.video[0].storage.settings.acquisition_dimensions) == 3 + + # sleep(10) + + # set and test + props = runtime.set_configuration(props) + assert len(props.video[0].storage.settings.acquisition_dimensions) == 3 + + assert ( + props.video[0].storage.settings.acquisition_dimensions[0].name + == dimension_x.name + ) + assert ( + props.video[0].storage.settings.acquisition_dimensions[0].kind + == dimension_x.kind + ) + assert ( + props.video[0].storage.settings.acquisition_dimensions[0].array_size_px + == dimension_x.array_size_px + ) + assert ( + props.video[0].storage.settings.acquisition_dimensions[0].chunk_size_px + == dimension_x.chunk_size_px + ) + + assert ( + props.video[0].storage.settings.acquisition_dimensions[2].name + == dimension_t.name + ) + assert ( + props.video[0].storage.settings.acquisition_dimensions[2].kind + == dimension_t.kind + ) + assert ( + props.video[0].storage.settings.acquisition_dimensions[2].array_size_px + == dimension_t.array_size_px + ) + assert ( + props.video[0].storage.settings.acquisition_dimensions[2].chunk_size_px + == dimension_t.chunk_size_px + ) + + def test_write_external_metadata_to_zarr( runtime: Runtime, request: pytest.FixtureRequest ): @@ -37,9 +121,34 @@ def test_write_external_metadata_to_zarr( metadata = {"hello": "world"} p.video[0].storage.settings.external_metadata_json = json.dumps(metadata) p.video[0].storage.settings.pixel_scale_um = (0.5, 4) - p.video[0].storage.settings.chunk_dims_px.width = 33 - p.video[0].storage.settings.chunk_dims_px.height = 47 - p.video[0].storage.settings.chunk_dims_px.planes = 4 + + # configure storage dimensions + dimension_x = acquire.StorageDimension() + dimension_x.name = "x" + dimension_x.kind = acquire.DimensionType.Space + dimension_x.array_size_px = 33 + dimension_x.chunk_size_px = 33 + assert dimension_x.shard_size_chunks == 0 + + dimension_y = acquire.StorageDimension() + dimension_y.name = "y" + dimension_y.kind = acquire.DimensionType.Space + dimension_y.array_size_px = 47 + dimension_y.chunk_size_px = 47 + assert dimension_y.shard_size_chunks == 0 + + dimension_z = acquire.StorageDimension() + dimension_z.name = "z" + dimension_z.kind = acquire.DimensionType.Space + dimension_z.array_size_px = 0 + dimension_z.chunk_size_px = 4 + assert dimension_z.shard_size_chunks == 0 + + p.video[0].storage.settings.acquisition_dimensions = [ + dimension_x, + dimension_y, + dimension_z, + ] p = runtime.set_configuration(p) @@ -65,7 +174,6 @@ def test_write_external_metadata_to_zarr( image_data = multi_scale_image_node.data[0] assert image_data.shape == ( p.video[0].max_frame_count, - 1, p.video[0].camera.settings.shape[1], p.video[0].camera.settings.shape[0], ) @@ -74,13 +182,13 @@ def test_write_external_metadata_to_zarr( axes = multi_scale_image_metadata["axes"] axis_names = tuple(a["name"] for a in axes) - assert axis_names == ("t", "c", "y", "x") + assert axis_names == ("z", "y", "x") axis_types = tuple(a["type"] for a in axes) - assert axis_types == ("time", "channel", "space", "space") + assert axis_types == ("space", "space", "space") axis_units = tuple(a.get("unit") for a in axes) - assert axis_units == (None, None, "micrometer", "micrometer") + assert axis_units == (None, "micrometer", "micrometer") # We only have one multi-scale level and one transform. transform = multi_scale_image_metadata["coordinateTransformations"][0][0] @@ -123,6 +231,43 @@ def test_write_compressed_zarr( p.video[0].storage.settings.filename = filename metadata = {"foo": "bar"} p.video[0].storage.settings.external_metadata_json = json.dumps(metadata) + + # configure storage dimensions + dimension_x = acquire.StorageDimension() + dimension_x.name = "x" + dimension_x.kind = acquire.DimensionType.Space + dimension_x.array_size_px = 64 + dimension_x.chunk_size_px = 64 + assert dimension_x.shard_size_chunks == 0 + + dimension_y = acquire.StorageDimension() + dimension_y.name = "y" + dimension_y.kind = acquire.DimensionType.Space + dimension_y.array_size_px = 48 + dimension_y.chunk_size_px = 48 + assert dimension_y.shard_size_chunks == 0 + + dimension_c = acquire.StorageDimension() + dimension_c.name = "c" + dimension_c.kind = acquire.DimensionType.Channel + dimension_c.array_size_px = 1 + dimension_c.chunk_size_px = 1 + assert dimension_c.shard_size_chunks == 0 + + dimension_t = acquire.StorageDimension() + dimension_t.name = "t" + dimension_t.kind = acquire.DimensionType.Time + dimension_t.array_size_px = 0 + dimension_t.chunk_size_px = 70 + assert dimension_t.shard_size_chunks == 0 + + p.video[0].storage.settings.acquisition_dimensions = [ + dimension_x, + dimension_y, + dimension_c, + dimension_t, + ] + runtime.set_configuration(p) runtime.start() @@ -186,9 +331,33 @@ def test_write_zarr_with_chunking( p.video[0].storage.settings.filename = f"{request.node.name}.zarr" p.video[0].max_frame_count = number_of_frames - p.video[0].storage.settings.chunk_dims_px.width = 1920 // 2 - p.video[0].storage.settings.chunk_dims_px.height = 1080 // 2 - p.video[0].storage.settings.chunk_dims_px.planes = 64 + # configure storage dimensions + dimension_x = acquire.StorageDimension() + dimension_x.name = "x" + dimension_x.kind = acquire.DimensionType.Space + dimension_x.array_size_px = 1920 + dimension_x.chunk_size_px = 960 + assert dimension_x.shard_size_chunks == 0 + + dimension_y = acquire.StorageDimension() + dimension_y.name = "y" + dimension_y.kind = acquire.DimensionType.Space + dimension_y.array_size_px = 1080 + dimension_y.chunk_size_px = 540 + assert dimension_y.shard_size_chunks == 0 + + dimension_t = acquire.StorageDimension() + dimension_t.name = "t" + dimension_t.kind = acquire.DimensionType.Time + dimension_t.array_size_px = 0 + dimension_t.chunk_size_px = 64 + assert dimension_t.shard_size_chunks == 0 + + p.video[0].storage.settings.acquisition_dimensions = [ + dimension_x, + dimension_y, + dimension_t, + ] runtime.set_configuration(p) @@ -198,11 +367,10 @@ def test_write_zarr_with_chunking( group = zarr.open(p.video[0].storage.settings.filename) data = group["0"] - assert data.chunks == (64, 1, 1080 // 2, 1920 // 2) + assert data.chunks == (64, 540, 960) assert data.shape == ( number_of_frames, - 1, p.video[0].camera.settings.shape[1], p.video[0].camera.settings.shape[0], ) @@ -233,14 +401,33 @@ def test_write_zarr_multiscale( p.video[0].storage.settings.pixel_scale_um = (1, 1) p.video[0].max_frame_count = 100 - p.video[0].storage.settings.chunk_dims_px.width = ( - p.video[0].camera.settings.shape[0] // 3 - ) - p.video[0].storage.settings.chunk_dims_px.height = ( - p.video[0].camera.settings.shape[1] // 3 - ) - p.video[0].storage.settings.chunk_dims_px.planes = 64 - + # configure storage dimensions + dimension_x = acquire.StorageDimension() + dimension_x.name = "x" + dimension_x.kind = acquire.DimensionType.Space + dimension_x.array_size_px = 1920 + dimension_x.chunk_size_px = 640 + assert dimension_x.shard_size_chunks == 0 + + dimension_y = acquire.StorageDimension() + dimension_y.name = "y" + dimension_y.kind = acquire.DimensionType.Space + dimension_y.array_size_px = 1080 + dimension_y.chunk_size_px = 360 + assert dimension_y.shard_size_chunks == 0 + + dimension_t = acquire.StorageDimension() + dimension_t.name = "t" + dimension_t.kind = acquire.DimensionType.Time + dimension_t.array_size_px = 0 + dimension_t.chunk_size_px = 64 + assert dimension_t.shard_size_chunks == 0 + + p.video[0].storage.settings.acquisition_dimensions = [ + dimension_x, + dimension_y, + dimension_t, + ] p.video[0].storage.settings.enable_multiscale = True runtime.set_configuration(p) @@ -257,11 +444,11 @@ def test_write_zarr_multiscale( ] assert len(data) == 3 - image = data[0][0, 0, :, :].compute() # convert dask array to numpy array + image = data[0][0, :, :].compute() # convert dask array to numpy array for d in data: assert ( - np.linalg.norm(image - d[0, 0, :, :].compute()) == 0 + np.linalg.norm(image - d[0, :, :].compute()) == 0 ) # validate against the same method from scikit-image image = downscale_local_mean(image, (2, 2)).astype(np.uint8) @@ -299,9 +486,33 @@ def test_write_zarr_v3( p.video[0].storage.settings.filename = f"{request.node.name}.zarr" p.video[0].max_frame_count = number_of_frames - p.video[0].storage.settings.chunk_dims_px.width = 1920 // 2 - p.video[0].storage.settings.chunk_dims_px.height = 1080 // 2 - p.video[0].storage.settings.chunk_dims_px.planes = 64 + # configure storage dimensions + dimension_x = acquire.StorageDimension() + dimension_x.name = "x" + dimension_x.kind = acquire.DimensionType.Space + dimension_x.array_size_px = 1920 + dimension_x.chunk_size_px = 960 + dimension_x.shard_size_chunks = 2 + + dimension_y = acquire.StorageDimension() + dimension_y.name = "y" + dimension_y.kind = acquire.DimensionType.Space + dimension_y.array_size_px = 1080 + dimension_y.chunk_size_px = 540 + dimension_y.shard_size_chunks = 2 + + dimension_t = acquire.StorageDimension() + dimension_t.name = "t" + dimension_t.kind = acquire.DimensionType.Time + dimension_t.array_size_px = 0 + dimension_t.chunk_size_px = 64 + dimension_t.shard_size_chunks = 1 + + p.video[0].storage.settings.acquisition_dimensions = [ + dimension_x, + dimension_y, + dimension_t, + ] runtime.set_configuration(p) @@ -312,11 +523,10 @@ def test_write_zarr_v3( group = zarr.open(store=store, mode="r") data = group["0"] - assert data.chunks == (64, 1, 1080 // 2, 1920 // 2) + assert data.chunks == (64, 540, 960) assert data.shape == ( number_of_frames, - 1, p.video[0].camera.settings.shape[1], p.video[0].camera.settings.shape[0], ) @@ -342,6 +552,32 @@ def test_metadata_with_trailing_whitespace( p.video[0].storage.settings.external_metadata_json = ( json.dumps(metadata) + " " ) + + # configure storage dimensions + dimension_x = acquire.StorageDimension() + dimension_x.name = "x" + dimension_x.kind = acquire.DimensionType.Space + dimension_x.array_size_px = 64 + dimension_x.chunk_size_px = 64 + + dimension_y = acquire.StorageDimension() + dimension_y.name = "y" + dimension_y.kind = acquire.DimensionType.Space + dimension_y.array_size_px = 48 + dimension_y.chunk_size_px = 48 + + dimension_t = acquire.StorageDimension() + dimension_t.name = "t" + dimension_t.kind = acquire.DimensionType.Time + dimension_t.array_size_px = 0 + dimension_t.chunk_size_px = 64 + + p.video[0].storage.settings.acquisition_dimensions = [ + dimension_x, + dimension_y, + dimension_t, + ] + runtime.set_configuration(p) runtime.start()