From 18841bf83842f51a9e248002e4ac496d4ef1eeaf Mon Sep 17 00:00:00 2001 From: uermel Date: Tue, 18 Aug 2026 10:26:52 -0700 Subject: [PATCH 1/3] feat: delegate feature persistence to copick --- src/copick_utils/features/skimage.py | 37 ++++++++------- tests/test_delegated_writers.py | 68 ++++++++++++++++++++++++++++ tests/test_zarr_migration.py | 54 +++++++++++++--------- 3 files changed, 119 insertions(+), 40 deletions(-) diff --git a/src/copick_utils/features/skimage.py b/src/copick_utils/features/skimage.py index 8ef5d5b..23b0fa4 100644 --- a/src/copick_utils/features/skimage.py +++ b/src/copick_utils/features/skimage.py @@ -1,6 +1,4 @@ import numpy as np -import zarr -from numcodecs import Blosc from skimage.feature import multiscale_basic_features from copick_utils.io.zarr import get_level_array @@ -16,6 +14,9 @@ def compute_skimage_features( sigma_min=0.5, sigma_max=16.0, feature_chunk_size=None, + *, + chunks=None, + shards=None, ): """ Processes the tomogram chunkwise and computes the multiscale basic features. @@ -42,25 +43,11 @@ def compute_skimage_features( ) num_features = test_features.shape[-1] - # Prepare output Zarr array directly in the tomogram store + # Preserve the existing entity-creation timing, but defer all persistence + # until the complete feature tensor has been assembled. print(f"Creating new feature store with {num_features} features...") copick_features = tomogram.new_features(feature_type) - feature_store = copick_features.zarr() - - # Use the provided feature chunk size if available, otherwise default to the input chunk size - if feature_chunk_size is None: - feature_chunk_size = (num_features, *chunk_size) - else: - feature_chunk_size = (num_features, *feature_chunk_size) - - out_array = zarr.create( - shape=(num_features, *image.shape), - chunks=feature_chunk_size, - dtype="float32", - compressor=Blosc(cname="zstd", clevel=3, shuffle=2), - store=feature_store, - overwrite=True, - ) + out_array = np.empty((num_features, *image.shape), dtype=np.float32) # Process each chunk for z in range(0, image.shape[0], chunk_size[0]): @@ -98,6 +85,18 @@ def compute_skimage_features( x : x + chunk_size[2], ] = contiguous_chunk + storage_chunks = chunks + if storage_chunks is None and feature_chunk_size is not None: + storage_chunks = feature_chunk_size + + copick_features.from_numpy( + out_array, + chunks=storage_chunks, + shards=shards, + dtype=np.float32, + overwrite=True, + ) + print(f"Features saved under feature type '{feature_type}'") return copick_features diff --git a/tests/test_delegated_writers.py b/tests/test_delegated_writers.py index a4a96c5..c84bb4b 100644 --- a/tests/test_delegated_writers.py +++ b/tests/test_delegated_writers.py @@ -8,6 +8,7 @@ from copick.impl.filesystem import CopickConfigFSSpec, CopickRootFSSpec from copick.util.ome import get_level_path from copick_utils.converters import lazy_converter +from copick_utils.features.skimage import compute_skimage_features from copick_utils.io import writers from copick_utils.process.rescale import rescale_segmentation @@ -67,6 +68,73 @@ def test_processor_reads_and_writes_through_copick_entities(run): _assert_canonical_volume(derived, labels) +@pytest.mark.parametrize( + ("feature_chunk_size", "chunks", "shards", "expected_chunks", "expected_shards"), + [ + (None, None, None, (1, 128, 128, 128), (1, 128, 128, 128)), + ((3, 4, 5), None, None, (1, 3, 4, 5), (1, 6, 8, 10)), + ((3, 4, 5), (1, 2, 3, 4), (1, 6, 6, 8), (1, 2, 3, 4), (1, 6, 6, 8)), + ], +) +def test_feature_writer_uses_core_layout_policy( + run, + feature_chunk_size, + chunks, + shards, + expected_chunks, + expected_shards, +): + volume = np.arange(5 * 6 * 7, dtype=np.float32).reshape(5, 6, 7) + writers.tomogram(run, volume, voxel_size=10, algorithm="wbp") + tomogram = run.get_voxel_spacing(10).get_tomogram("wbp") + + features = compute_skimage_features( + tomogram, + "skimage", + run.root, + intensity=True, + edges=False, + texture=False, + sigma_min=0.5, + sigma_max=0.5, + feature_chunk_size=feature_chunk_size, + chunks=chunks, + shards=shards, + ) + + group = zarr.open_group(store=features.zarr(), mode="r") + level = group[get_level_path(group, 0)] + axes = tuple(axis["name"] for axis in group.attrs["ome"]["multiscales"][0]["axes"]) + assert group.metadata.zarr_format == 3 + assert axes == ("feature", "z", "y", "x") + assert level.metadata.dimension_names == axes + assert level.shape == (1, 5, 6, 7) + assert level.chunks == expected_chunks + assert level.shards == expected_shards + np.testing.assert_array_equal(features.numpy(), level[:]) + + +def test_invalid_explicit_feature_layout_is_rejected_by_core(run): + volume = np.arange(5 * 6 * 7, dtype=np.float32).reshape(5, 6, 7) + writers.tomogram(run, volume, voxel_size=10, algorithm="wbp") + tomogram = run.get_voxel_spacing(10).get_tomogram("wbp") + + with pytest.raises(ValueError, match="shard"): + compute_skimage_features( + tomogram, + "invalid-layout", + run.root, + intensity=True, + edges=False, + texture=False, + sigma_min=0.5, + sigma_max=0.5, + feature_chunk_size=(3, 4, 5), + chunks=(1, 2, 3, 4), + shards=(1, 3, 6, 8), + ) + + def test_lazy_worker_preserves_core_writer_failure(monkeypatch): task = { "segmentation": SimpleNamespace(session_id="input"), diff --git a/tests/test_zarr_migration.py b/tests/test_zarr_migration.py index 23e0b63..0e863db 100644 --- a/tests/test_zarr_migration.py +++ b/tests/test_zarr_migration.py @@ -5,6 +5,7 @@ """ import hashlib +import inspect import numpy as np import pytest @@ -28,10 +29,15 @@ def _tomogram_store(path="0", zarr_format=3): class _Features: def __init__(self): - self.store = _memory_store() + self.data = None + self.write_calls = [] - def zarr(self): - return self.store + def from_numpy(self, data, **kwargs): + self.data = np.array(data, copy=True) + self.write_calls.append(kwargs) + + def numpy(self): + return np.array(self.data, copy=True) class _Tomogram: @@ -67,11 +73,6 @@ def test_level_array_rejects_out_of_range_levels(level): get_level_array(_Tomogram(store), level) -@pytest.mark.xfail( - raises=ValueError, - strict=True, - reason="The retained feature writer is migrated in U3", -) def test_pre_migration_feature_result_is_frozen(): """Protect the existing chunk subdivision and boundary behavior.""" store, _ = _tomogram_store() @@ -83,7 +84,7 @@ def test_pre_migration_feature_result_is_frozen(): sigma_max=0.5, feature_chunk_size=(3, 4, 5), ) - result = zarr.open(features.zarr(), mode="r")[:] + result = features.numpy() assert result.shape == (5, 5, 6, 7) assert result.dtype == np.float32 @@ -91,12 +92,7 @@ def test_pre_migration_feature_result_is_frozen(): assert rounded_digest == "8364181d58811d79fe86847872316a97370ed2737aeee6411a38753124305312" -@pytest.mark.xfail( - raises=ValueError, - strict=True, - reason="The retained feature writer is migrated in U3", -) -def test_pre_migration_feature_store_documents_reader_incompatibility(): +def test_feature_writer_delegates_one_final_float32_write(): store, _ = _tomogram_store() features = compute_skimage_features( _Tomogram(store), @@ -110,9 +106,25 @@ def test_pre_migration_feature_store_documents_reader_incompatibility(): feature_chunk_size=(3, 4, 5), ) - # The old implementation writes an array at the store root. CopickFeatures - # expects an OME group and therefore cannot resolve a metadata-defined level. - root = zarr.open(features.zarr(), mode="r") - assert isinstance(root, zarr.Array) - with pytest.raises((AttributeError, TypeError, zarr.errors.ContainsArrayError)): - zarr.open_group(store=features.zarr(), mode="r") + assert features.data.shape == (1, 5, 6, 7) + assert features.data.dtype == np.float32 + assert features.write_calls == [ + { + "chunks": (3, 4, 5), + "shards": None, + "dtype": np.float32, + "overwrite": True, + }, + ] + + +def test_feature_layout_controls_are_optional_keyword_only(): + signature = inspect.signature(compute_skimage_features) + + # All pre-migration parameters still bind positionally in their original order. + signature.bind(object(), "features", object(), True, True, True, 0.5, 16.0, (32, 32, 32)) + + assert signature.parameters["chunks"].kind is inspect.Parameter.KEYWORD_ONLY + assert signature.parameters["chunks"].default is None + assert signature.parameters["shards"].kind is inspect.Parameter.KEYWORD_ONLY + assert signature.parameters["shards"].default is None From fc4f7a15212b4d225de35d32559811a10a0c4680 Mon Sep 17 00:00:00 2001 From: uermel Date: Tue, 18 Aug 2026 14:09:33 -0700 Subject: [PATCH 2/3] fix: bound feature staging memory and overlap slices --- src/copick_utils/features/skimage.py | 112 ++++++++++++++----------- tests/test_zarr_migration.py | 121 +++++++++++++++++++++++++-- 2 files changed, 178 insertions(+), 55 deletions(-) diff --git a/src/copick_utils/features/skimage.py b/src/copick_utils/features/skimage.py index 23b0fa4..1cd4948 100644 --- a/src/copick_utils/features/skimage.py +++ b/src/copick_utils/features/skimage.py @@ -1,9 +1,25 @@ +import tempfile + import numpy as np from skimage.feature import multiscale_basic_features from copick_utils.io.zarr import get_level_array +def _axis_slices(origin, chunk_size, overlap, image_size): + target_end = min(origin + chunk_size, image_size) + read_start = max(origin - overlap, 0) + read_end = min(target_end + overlap, image_size) + crop_start = origin - read_start + crop_end = target_end - read_start + + return ( + slice(read_start, read_end), + slice(crop_start, crop_end), + slice(origin, target_end), + ) + + def compute_skimage_features( tomogram, feature_type, @@ -44,58 +60,54 @@ def compute_skimage_features( num_features = test_features.shape[-1] # Preserve the existing entity-creation timing, but defer all persistence - # until the complete feature tensor has been assembled. + # until the complete feature tensor has been assembled on disk. print(f"Creating new feature store with {num_features} features...") copick_features = tomogram.new_features(feature_type) - out_array = np.empty((num_features, *image.shape), dtype=np.float32) - - # Process each chunk - for z in range(0, image.shape[0], chunk_size[0]): - for y in range(0, image.shape[1], chunk_size[1]): - for x in range(0, image.shape[2], chunk_size[2]): - z_start = max(z - overlap, 0) - z_end = min(z + chunk_size[0] + overlap, image.shape[0]) - y_start = max(y - overlap, 0) - y_end = min(y + chunk_size[1] + overlap, image.shape[1]) - x_start = max(x - overlap, 0) - x_end = min(x + chunk_size[2] + overlap, image.shape[2]) - - chunk = image[z_start:z_end, y_start:y_end, x_start:x_end] - chunk_features = multiscale_basic_features( - chunk, - intensity=intensity, - edges=edges, - texture=texture, - sigma_min=sigma_min, - sigma_max=sigma_max, - ) - - # Adjust indices for overlap - z_slice = slice(overlap if z_start > 0 else 0, None if z_end == image.shape[0] else -overlap) - y_slice = slice(overlap if y_start > 0 else 0, None if y_end == image.shape[1] else -overlap) - x_slice = slice(overlap if x_start > 0 else 0, None if x_end == image.shape[2] else -overlap) - - # Ensure contiguous array and correct slicing - contiguous_chunk = np.ascontiguousarray(chunk_features[z_slice, y_slice, x_slice].transpose(3, 0, 1, 2)) - - out_array[ - 0:num_features, - z : z + chunk_size[0], - y : y + chunk_size[1], - x : x + chunk_size[2], - ] = contiguous_chunk - - storage_chunks = chunks - if storage_chunks is None and feature_chunk_size is not None: - storage_chunks = feature_chunk_size - - copick_features.from_numpy( - out_array, - chunks=storage_chunks, - shards=shards, - dtype=np.float32, - overwrite=True, - ) + with tempfile.TemporaryDirectory(prefix="copick-utils-skimage-") as directory: + out_array = np.memmap( + f"{directory}/features.dat", + dtype=np.float32, + mode="w+", + shape=(num_features, *image.shape), + ) + try: + # Process each chunk + for z in range(0, image.shape[0], chunk_size[0]): + for y in range(0, image.shape[1], chunk_size[1]): + for x in range(0, image.shape[2], chunk_size[2]): + z_read, z_crop, z_output = _axis_slices(z, chunk_size[0], overlap, image.shape[0]) + y_read, y_crop, y_output = _axis_slices(y, chunk_size[1], overlap, image.shape[1]) + x_read, x_crop, x_output = _axis_slices(x, chunk_size[2], overlap, image.shape[2]) + + chunk = image[z_read, y_read, x_read] + chunk_features = multiscale_basic_features( + chunk, + intensity=intensity, + edges=edges, + texture=texture, + sigma_min=sigma_min, + sigma_max=sigma_max, + ) + + contiguous_chunk = np.ascontiguousarray( + chunk_features[z_crop, y_crop, x_crop].transpose(3, 0, 1, 2), + ) + out_array[:, z_output, y_output, x_output] = contiguous_chunk + + storage_chunks = chunks + if storage_chunks is None and feature_chunk_size is not None: + storage_chunks = feature_chunk_size + + out_array.flush() + copick_features.from_numpy( + out_array, + chunks=storage_chunks, + shards=shards, + dtype=np.float32, + overwrite=True, + ) + finally: + del out_array print(f"Features saved under feature type '{feature_type}'") return copick_features diff --git a/tests/test_zarr_migration.py b/tests/test_zarr_migration.py index 0e863db..0e3bf00 100644 --- a/tests/test_zarr_migration.py +++ b/tests/test_zarr_migration.py @@ -6,10 +6,13 @@ import hashlib import inspect +import tracemalloc +from pathlib import Path import numpy as np import pytest import zarr + from copick_utils.features.skimage import compute_skimage_features from copick_utils.io.zarr import get_level_array @@ -18,21 +21,35 @@ def _memory_store(): return zarr.storage.MemoryStore() -def _tomogram_store(path="0", zarr_format=3): +def _tomogram_store(path="0", zarr_format=3, shape=(5, 6, 7), chunks=(3, 4, 5)): store = _memory_store() group = zarr.open_group(store=store, mode="w", zarr_format=zarr_format) - data = ((np.indices((5, 6, 7)) * np.array([11, 5, 2])[:, None, None, None]).sum(0) % 17).astype(np.float32) - group.create_array(path, data=data, chunks=(3, 4, 5)) + data = ((np.indices(shape) * np.array([11, 5, 2])[:, None, None, None]).sum(0) % 17).astype(np.float32) + group.create_array(path, data=data, chunks=chunks) group.attrs["multiscales"] = [{"datasets": [{"path": path}]}] return store, data +def _empty_tomogram_store(shape, chunks): + store = _memory_store() + group = zarr.open_group(store=store, mode="w", zarr_format=3) + group.create_array("0", shape=shape, chunks=chunks, dtype=np.float32, fill_value=0) + group.attrs["multiscales"] = [{"datasets": [{"path": "0"}]}] + return store + + class _Features: def __init__(self): self.data = None + self.input_is_memmap = False + self.staging_path = None + self.staging_existed_during_write = False self.write_calls = [] def from_numpy(self, data, **kwargs): + self.input_is_memmap = isinstance(data, np.memmap) + self.staging_path = Path(data.filename) + self.staging_existed_during_write = self.staging_path.exists() self.data = np.array(data, copy=True) self.write_calls.append(kwargs) @@ -40,17 +57,33 @@ def numpy(self): return np.array(self.data, copy=True) +class _NoCopyFeatures(_Features): + def from_numpy(self, data, **kwargs): + self.input_is_memmap = isinstance(data, np.memmap) + self.staging_path = Path(data.filename) + self.staging_existed_during_write = self.staging_path.exists() + self.data = (data.shape, data.dtype) + self.write_calls.append(kwargs) + + +class _FailingFeatures(_NoCopyFeatures): + def from_numpy(self, data, **kwargs): + super().from_numpy(data, **kwargs) + raise RuntimeError("feature write failed") + + class _Tomogram: - def __init__(self, store): + def __init__(self, store, features_factory=_Features): self.store = store self.features = None + self.features_factory = features_factory def zarr(self): return self.store def new_features(self, feature_type): assert feature_type == "golden" - self.features = _Features() + self.features = self.features_factory() return self.features @@ -108,6 +141,9 @@ def test_feature_writer_delegates_one_final_float32_write(): assert features.data.shape == (1, 5, 6, 7) assert features.data.dtype == np.float32 + assert features.input_is_memmap + assert features.staging_existed_during_write + assert not features.staging_path.exists() assert features.write_calls == [ { "chunks": (3, 4, 5), @@ -118,6 +154,81 @@ def test_feature_writer_delegates_one_final_float32_write(): ] +@pytest.mark.parametrize("shape", [(13, 8, 8), (8, 13, 8), (8, 8, 13)]) +def test_feature_overlap_trimming_handles_image_end_inside_overlap(monkeypatch, shape): + def constant_features(chunk, **kwargs): + return np.full((*chunk.shape, 1), 7, dtype=np.float32) + + monkeypatch.setattr("copick_utils.features.skimage.multiscale_basic_features", constant_features) + store, _ = _tomogram_store(shape=shape, chunks=(4, 4, 4)) + + features = compute_skimage_features( + _Tomogram(store), + "golden", + None, + intensity=True, + edges=False, + texture=False, + feature_chunk_size=(4, 4, 4), + ) + + assert features.data.shape == (1, *shape) + np.testing.assert_array_equal(features.data, np.full((1, *shape), 7, dtype=np.float32)) + + +def test_feature_staging_is_removed_when_final_write_fails(): + store, _ = _tomogram_store() + tomogram = _Tomogram(store, _FailingFeatures) + + with pytest.raises(RuntimeError, match="feature write failed"): + compute_skimage_features( + tomogram, + "golden", + None, + intensity=True, + edges=False, + texture=False, + sigma_min=0.5, + sigma_max=0.5, + feature_chunk_size=(3, 4, 5), + ) + + assert tomogram.features.input_is_memmap + assert tomogram.features.staging_existed_during_write + assert not tomogram.features.staging_path.exists() + + +def test_feature_staging_has_bounded_python_memory(monkeypatch): + shape = (96, 96, 96) + feature_count = 30 + + def constant_features(chunk, **kwargs): + return np.full((*chunk.shape, feature_count), 7, dtype=np.float32) + + monkeypatch.setattr("copick_utils.features.skimage.multiscale_basic_features", constant_features) + tomogram = _Tomogram(_empty_tomogram_store(shape, (16, 16, 16)), _NoCopyFeatures) + + tracemalloc.start() + features = compute_skimage_features( + tomogram, + "golden", + None, + intensity=True, + edges=False, + texture=False, + feature_chunk_size=(16, 16, 16), + ) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + logical_size = feature_count * np.prod(shape) * np.dtype(np.float32).itemsize + assert peak < 64 * 1024 * 1024 + assert peak < logical_size * 0.75 + assert features.data == ((feature_count, *shape), np.dtype(np.float32)) + assert features.input_is_memmap + assert not features.staging_path.exists() + + def test_feature_layout_controls_are_optional_keyword_only(): signature = inspect.signature(compute_skimage_features) From a91bb84f5eb6bdab410af3684cd809d2af9ed5c9 Mon Sep 17 00:00:00 2001 From: uermel Date: Tue, 18 Aug 2026 14:11:51 -0700 Subject: [PATCH 3/3] style: satisfy pinned import ordering --- tests/test_zarr_migration.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_zarr_migration.py b/tests/test_zarr_migration.py index 0e3bf00..c25057c 100644 --- a/tests/test_zarr_migration.py +++ b/tests/test_zarr_migration.py @@ -12,7 +12,6 @@ import numpy as np import pytest import zarr - from copick_utils.features.skimage import compute_skimage_features from copick_utils.io.zarr import get_level_array