Skip to content

Commit 9ebcb95

Browse files
committed
ink inference: optional vc_reader backend for full-3D patch reads
--reader vc_reader reads each DataLoader batch of planned patches with one vc_reader read_regions call, so chunks shared by overlapping patches are fetched once and in parallel. The default zarr reader is unchanged.
1 parent 13a88a6 commit 9ebcb95

4 files changed

Lines changed: 257 additions & 9 deletions

File tree

‎vesuvius/src/vesuvius/ink_detection/inference/infer_full3d_tifxyz.py‎

Lines changed: 61 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from __future__ import annotations
44

5+
import importlib.util
56
import itertools
67
import logging
78
import math
@@ -46,8 +47,9 @@
4647
)
4748
from vesuvius.ink_detection.models.model import make_model
4849
from vesuvius.ink_detection.volume_io import (
50+
open_vc_reader,
4951
open_volume,
50-
read_bbox_with_padding,
52+
read_bboxes_with_padding,
5153
)
5254
from vesuvius.label_zarr import create_v2_array, open_v2_group
5355
from vesuvius.utils.cli import HyphenUnderscoreParser
@@ -149,13 +151,30 @@ def parse_args(argv: Sequence[str] | None = None):
149151
parser.add_argument("--max-target-chunks", type=int, default=None)
150152
parser.add_argument("--cache-dir", type=Path, default=None)
151153
parser.add_argument("--cache-max-gb", type=float, default=None)
154+
parser.add_argument("--reader", choices=("zarr", "vc_reader"), default="zarr")
155+
parser.add_argument("--reader-connections", type=int, default=32)
156+
parser.add_argument(
157+
"--reader-memory-cache-mb",
158+
type=int,
159+
default=512,
160+
help="vc_reader memory cache per DataLoader worker",
161+
)
152162
parser.add_argument("--log-level", default="INFO")
153163
args = parser.parse_args(argv)
154164
checks = (
155165
(args.batch_size > 0, "--batch-size must be positive"),
156166
(args.num_workers >= 0, "--num-workers must be >= 0"),
157167
(args.prefetch_factor > 0, "--prefetch-factor must be positive"),
158168
(args.downsample_workers > 0, "--downsample-workers must be positive"),
169+
(args.reader_connections > 0, "--reader-connections must be positive"),
170+
(
171+
args.reader_memory_cache_mb >= 0,
172+
"--reader-memory-cache-mb must be >= 0",
173+
),
174+
(
175+
args.reader != "vc_reader" or importlib.util.find_spec("vc_reader"),
176+
"--reader vc_reader needs the vc_reader package (villa/vc_reader)",
177+
),
159178
(
160179
args.tta_batch_size is None or args.tta_batch_size > 0,
161180
"--tta-batch-size must be positive",
@@ -849,7 +868,12 @@ def __init__(
849868
config,
850869
cache_dir=None,
851870
cache_max_gb=None,
871+
reader="zarr",
872+
reader_connections=32,
873+
reader_memory_cache_mb=512,
852874
) -> None:
875+
if reader not in {"zarr", "vc_reader"}:
876+
raise ValueError(f"reader must be 'zarr' or 'vc_reader', got {reader!r}")
853877
self.tifxyz_dir = Path(tifxyz_dir)
854878
self.volume_path = str(volume_path)
855879
self.resolution = str(resolution)
@@ -859,6 +883,9 @@ def __init__(
859883
self.normalization = config.data.normalization
860884
self.cache_dir = cache_dir
861885
self.cache_max_gb = cache_max_gb
886+
self.reader = reader
887+
self.reader_connections = int(reader_connections)
888+
self.reader_memory_cache_mb = int(reader_memory_cache_mb)
862889
(
863890
self.flat_grid_stride,
864891
self.native_coordinate_scale,
@@ -883,6 +910,15 @@ def __getstate__(self):
883910
return state
884911

885912
def _ensure_volume(self):
913+
if self._volume is None and self.reader == "vc_reader":
914+
self._volume = open_vc_reader(
915+
self.volume_path,
916+
self.resolution,
917+
cache_dir=self.cache_dir,
918+
cache_max_gb=self.cache_max_gb,
919+
memory_cache_mb=self.reader_memory_cache_mb,
920+
connections=self.reader_connections,
921+
)
886922
if self._volume is None:
887923
self._volume = _open_shared_volume(
888924
self.volume_path,
@@ -949,15 +985,28 @@ def _surface_mask(self, bbox_zyx) -> np.ndarray:
949985
)
950986

951987
def __getitem__(self, index: int):
952-
patch = self.patches[int(index)]
953-
starts = (int(patch.z0), int(patch.y0), int(patch.x0))
954-
stops = tuple(
955-
starts[axis] + self.patch_size_zyx[axis] for axis in range(3)
956-
)
957-
bbox = (*starts, *stops)
958-
crop, valid_slices = read_bbox_with_padding(
959-
self._ensure_volume(), bbox, fill_value=0
988+
return self.__getitems__([index])[0]
989+
990+
def __getitems__(self, indices):
991+
"""DataLoader batch fetch; a vc_reader volume reads the batch in one call."""
992+
bboxes = []
993+
for index in indices:
994+
patch = self.patches[int(index)]
995+
starts = (int(patch.z0), int(patch.y0), int(patch.x0))
996+
stops = tuple(
997+
starts[axis] + self.patch_size_zyx[axis] for axis in range(3)
998+
)
999+
bboxes.append((*starts, *stops))
1000+
reads = read_bboxes_with_padding(
1001+
self._ensure_volume(), bboxes, fill_value=0
9601002
)
1003+
return [
1004+
self._sample(bbox, crop, valid_slices)
1005+
for bbox, (crop, valid_slices) in zip(bboxes, reads)
1006+
]
1007+
1008+
def _sample(self, bbox, crop, valid_slices):
1009+
starts = bbox[:3]
9611010
crop = crop.astype(np.float32, copy=False)
9621011
if valid_slices is not None:
9631012
crop[valid_slices] = normalize_image(
@@ -1175,6 +1224,9 @@ def run_command(args) -> int:
11751224
config=config,
11761225
cache_dir=args.cache_dir,
11771226
cache_max_gb=args.cache_max_gb,
1227+
reader=args.reader,
1228+
reader_connections=args.reader_connections,
1229+
reader_memory_cache_mb=args.reader_memory_cache_mb,
11781230
)
11791231
run_native_inference(
11801232
args=args,

‎vesuvius/src/vesuvius/ink_detection/volume_io.py‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,3 +244,78 @@ def read_bbox_with_padding(
244244
)
245245
output[destination] = crop
246246
return output, destination
247+
248+
249+
def open_vc_reader(
250+
path: str | Path,
251+
resolution: int | str,
252+
*,
253+
cache_dir: str | Path | None = None,
254+
cache_max_gb: float | None = None,
255+
memory_cache_mb: int = 512,
256+
connections: int = 32,
257+
):
258+
"""Open one pyramid level with the optional compiled ``vc_reader`` package.
259+
260+
Each process (e.g. each DataLoader worker) gets its own memory cache and
261+
connection pool.
262+
"""
263+
264+
try:
265+
from vc_reader.reader import open_zarr
266+
except ImportError as exc:
267+
raise ImportError(
268+
"reader 'vc_reader' needs the vc_reader package (villa/vc_reader); "
269+
"install it or use the default zarr reader"
270+
) from exc
271+
return open_zarr(
272+
str(path),
273+
level=str(resolution),
274+
cache_dir=cache_dir,
275+
cache_max_bytes=None if cache_max_gb is None else int(cache_max_gb * 2**30),
276+
memory_cache=int(memory_cache_mb) << 20,
277+
connections=int(connections),
278+
)
279+
280+
281+
def read_bboxes_with_padding(
282+
volume: Any,
283+
bboxes_zyx: list[tuple[int, int, int, int, int, int]],
284+
*,
285+
fill_value: int | float = 0,
286+
) -> list[tuple[np.ndarray, tuple[slice, slice, slice] | None]]:
287+
"""Read several bboxes like read_bbox_with_padding.
288+
289+
A vc_reader volume reads them in one call, fetching shared chunks once.
290+
"""
291+
if not hasattr(volume, "read_regions"):
292+
return [
293+
read_bbox_with_padding(volume, bbox, fill_value=fill_value)
294+
for bbox in bboxes_zyx
295+
]
296+
bboxes = [tuple(int(value) for value in bbox) for bbox in bboxes_zyx]
297+
for bbox in bboxes:
298+
if any(bbox[axis + 3] <= bbox[axis] for axis in range(3)):
299+
raise ValueError(f"bbox must define a positive crop, got {bbox!r}")
300+
shape = tuple(int(value) for value in volume.shape[:3])
301+
sizes = [tuple(bbox[axis + 3] - bbox[axis] for axis in range(3)) for bbox in bboxes]
302+
crops = volume.read_regions(
303+
[(bbox[:3], size) for bbox, size in zip(bboxes, sizes)]
304+
)
305+
results = []
306+
for bbox, size, crop in zip(bboxes, sizes, crops):
307+
if crop.shape != size:
308+
raise ValueError(f"read_regions returned {crop.shape} for bbox {bbox!r}")
309+
# Copy out of vc_reader's recycled buffers and pad with fill_value.
310+
output = np.full(size, fill_value, dtype=np.dtype(volume.dtype))
311+
starts = [max(0, bbox[axis]) for axis in range(3)]
312+
stops = [min(shape[axis], bbox[axis + 3]) for axis in range(3)]
313+
destination = None
314+
if all(stop > start for start, stop in zip(starts, stops)):
315+
destination = tuple(
316+
slice(start - bbox[axis], stop - bbox[axis])
317+
for axis, (start, stop) in enumerate(zip(starts, stops))
318+
)
319+
output[destination] = crop[destination]
320+
results.append((output, destination))
321+
return results

‎vesuvius/tests/ink_detection/test_native_full3d_inference.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -551,6 +551,9 @@ def open_shared_volume_stub(path, resolution, *, cache_dir, cache_max_gb):
551551
prefetch_factor=2,
552552
cache_dir=None,
553553
cache_max_gb=None,
554+
reader="zarr",
555+
reader_connections=32,
556+
reader_memory_cache_mb=512,
554557
downsample_workers=1,
555558
)
556559
assert native.run_command(args) == 0
@@ -566,3 +569,53 @@ def open_shared_volume_stub(path, resolution, *, cache_dir, cache_max_gb):
566569
np.all(np.asarray(zarr.open_group(output, mode="r")[str(level)]) == 128)
567570
for level in range(6)
568571
)
572+
573+
574+
def test_native_dataset_vc_reader_batches_match_zarr_reads(monkeypatch):
575+
from torch.utils.data import DataLoader
576+
577+
from vesuvius.ink_detection.config import NormalizationConfig
578+
579+
array = np.random.default_rng(0).integers(0, 255, (6, 7, 8), dtype=np.uint8)
580+
581+
class RegionVolume:
582+
shape, dtype = array.shape, array.dtype
583+
584+
def read_regions(self, regions):
585+
calls.append(len(regions))
586+
padded = np.pad(array, 8, constant_values=7)
587+
return [
588+
padded[tuple(slice(s + 8, s + n + 8) for s, n in zip(start, size))]
589+
for start, size in regions
590+
]
591+
592+
calls = []
593+
monkeypatch.setattr(native, "_open_shared_volume", lambda *a, **k: array)
594+
monkeypatch.setattr(native, "open_vc_reader", lambda *a, **k: RegionVolume())
595+
config = SimpleNamespace(
596+
data=SimpleNamespace(
597+
mode="full_3d", normalization=NormalizationConfig.from_value("none")
598+
)
599+
)
600+
patches = [
601+
native.PatchSpec(z, y, x)
602+
for z, y, x in ((-2, 0, 1), (2, 3, 4), (5, 6, 7), (9, 0, 0))
603+
]
604+
batches = {}
605+
for reader in ("zarr", "vc_reader"):
606+
dataset = native.NativePatchDataset(
607+
tifxyz_dir=Path("unused"),
608+
volume_path="memory://native",
609+
resolution="0",
610+
patches=patches,
611+
patch_size_zyx=(3, 3, 3),
612+
config=config,
613+
reader=reader,
614+
)
615+
batches[reader] = list(DataLoader(dataset, batch_size=3))
616+
assert calls == [3, 1]
617+
for (image, starts), (want_image, want_starts) in zip(
618+
batches["vc_reader"], batches["zarr"]
619+
):
620+
assert torch.equal(image, want_image)
621+
assert torch.equal(starts, want_starts)

‎vesuvius/tests/ink_detection/test_volume_io.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
77
import multiprocessing as mp
88
import os
9+
import sys
910
from pathlib import Path
1011
import threading
1112
from types import SimpleNamespace
@@ -423,3 +424,70 @@ class CompatiblePathError(Exception):
423424
assert isinstance(
424425
volume_io._missing_node_error("missing"), CompatiblePathError
425426
)
427+
428+
429+
class _RegionVolume:
430+
"""vc_reader-like stub: read_regions fills outside the array with 7."""
431+
432+
def __init__(self, array):
433+
self.array = array
434+
self.shape = array.shape
435+
self.dtype = array.dtype
436+
self.calls = 0
437+
438+
def read_regions(self, regions):
439+
self.calls += 1
440+
padded = np.pad(self.array, 8, constant_values=7)
441+
crops = []
442+
for start, size in regions:
443+
index = tuple(slice(s + 8, s + n + 8) for s, n in zip(start, size))
444+
crops.append(padded[index]) # views, like vc_reader's recycled buffers
445+
return crops
446+
447+
448+
def test_read_bboxes_with_padding_matches_single_reads():
449+
array = np.arange(6 * 7 * 8, dtype=np.uint8).reshape(6, 7, 8)
450+
bboxes = [
451+
(-1, 1, 3, 2, 5, 7),
452+
(0, 0, 0, 6, 7, 8),
453+
(4, 5, 6, 9, 9, 9),
454+
(7, 0, 0, 9, 2, 2),
455+
]
456+
expected = [
457+
volume_io.read_bbox_with_padding(array, bbox, fill_value=0) for bbox in bboxes
458+
]
459+
for (crop, valid), (want, want_valid) in zip(
460+
volume_io.read_bboxes_with_padding(array, bboxes), expected
461+
):
462+
np.testing.assert_array_equal(crop, want)
463+
assert valid == want_valid
464+
stub = _RegionVolume(array)
465+
for (crop, valid), (want, want_valid) in zip(
466+
volume_io.read_bboxes_with_padding(stub, bboxes), expected
467+
):
468+
np.testing.assert_array_equal(crop, want)
469+
assert valid == want_valid
470+
assert stub.calls == 1
471+
472+
473+
def test_vc_reader_matches_zarr_reads(tmp_path):
474+
pytest.importorskip("vc_reader")
475+
array = np.random.default_rng(0).integers(0, 255, (9, 10, 11), dtype=np.uint8)
476+
_write_pyramid(tmp_path / "volume.zarr", array)
477+
bboxes = [(-2, 3, 4, 4, 9, 12), (5, 0, 0, 9, 10, 11), (8, 8, 8, 12, 12, 12)]
478+
zarr_reads = volume_io.read_bboxes_with_padding(
479+
open_volume(tmp_path / "volume.zarr", 0), bboxes
480+
)
481+
vc_reads = volume_io.read_bboxes_with_padding(
482+
volume_io.open_vc_reader(tmp_path / "volume.zarr", 0), bboxes
483+
)
484+
for (crop, valid), (want, want_valid) in zip(vc_reads, zarr_reads):
485+
np.testing.assert_array_equal(crop, want)
486+
assert valid == want_valid
487+
488+
489+
def test_open_vc_reader_without_package_says_how_to_fix(monkeypatch):
490+
monkeypatch.setitem(sys.modules, "vc_reader", None)
491+
monkeypatch.setitem(sys.modules, "vc_reader.reader", None)
492+
with pytest.raises(ImportError, match="needs the vc_reader package"):
493+
volume_io.open_vc_reader("volume.zarr", 0)

0 commit comments

Comments
 (0)