Skip to content

Commit 4c877dc

Browse files
committed
fix #8298
Signed-off-by: YunLiu <[email protected]>
1 parent df1ba5d commit 4c877dc

File tree

2 files changed

+24
-7
lines changed

2 files changed

+24
-7
lines changed

monai/inferers/merger.py

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,13 @@
1515
from abc import ABC, abstractmethod
1616
from collections.abc import Sequence
1717
from contextlib import nullcontext
18+
from tempfile import TemporaryDirectory
1819
from typing import TYPE_CHECKING, Any
1920

2021
import numpy as np
2122
import torch
2223

23-
from monai.utils import ensure_tuple_size, optional_import, require_pkg
24+
from monai.utils import ensure_tuple_size, get_package_version, optional_import, require_pkg, version_geq
2425

2526
if TYPE_CHECKING:
2627
import zarr
@@ -233,7 +234,7 @@ def __init__(
233234
store: zarr.storage.Store | str = "merged.zarr",
234235
value_store: zarr.storage.Store | str | None = None,
235236
count_store: zarr.storage.Store | str | None = None,
236-
compressor: str = "default",
237+
compressor: str | None = None,
237238
value_compressor: str | None = None,
238239
count_compressor: str | None = None,
239240
chunks: Sequence[int] | bool = True,
@@ -246,8 +247,20 @@ def __init__(
246247
self.value_dtype = value_dtype
247248
self.count_dtype = count_dtype
248249
self.store = store
249-
self.value_store = zarr.storage.TempStore() if value_store is None else value_store
250-
self.count_store = zarr.storage.TempStore() if count_store is None else count_store
250+
if version_geq(get_package_version("zarr"), "3.0.0"):
251+
if value_store is None:
252+
with TemporaryDirectory() as tmpdir:
253+
self.value_store = zarr.storage.LocalStore(tmpdir)
254+
else:
255+
self.value_store = value_store
256+
if count_store is None:
257+
with TemporaryDirectory() as tmpdir:
258+
self.count_store = zarr.storage.LocalStore(tmpdir)
259+
else:
260+
self.count_store = count_store
261+
else:
262+
self.value_store = zarr.storage.TempStore() if value_store is None else value_store
263+
self.count_store = zarr.storage.TempStore() if count_store is None else count_store
251264
self.chunks = chunks
252265
self.compressor = compressor
253266
self.value_compressor = value_compressor

tests/test_zarr_avg_merger.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -289,13 +289,17 @@ class ZarrAvgMergerTests(unittest.TestCase):
289289
def test_zarr_avg_merger_patches(self, arguments, patch_locations, expected):
290290
if "compressor" in arguments:
291291
if arguments["compressor"] != "default":
292-
arguments["compressor"] = zarr.codec_registry[arguments["compressor"].lower()]()
292+
arguments["compressor"] = numcodecs.registry.codec_registry[arguments["compressor"].lower()]()
293293
if "value_compressor" in arguments:
294294
if arguments["value_compressor"] != "default":
295-
arguments["value_compressor"] = zarr.codec_registry[arguments["value_compressor"].lower()]()
295+
arguments["value_compressor"] = numcodecs.registry.codec_registry[
296+
arguments["value_compressor"].lower()
297+
]()
296298
if "count_compressor" in arguments:
297299
if arguments["count_compressor"] != "default":
298-
arguments["count_compressor"] = zarr.codec_registry[arguments["count_compressor"].lower()]()
300+
arguments["count_compressor"] = numcodecs.registry.codec_registry[
301+
arguments["count_compressor"].lower()
302+
]()
299303
merger = ZarrAvgMerger(**arguments)
300304
for pl in patch_locations:
301305
merger.aggregate(pl[0], pl[1])

0 commit comments

Comments
 (0)