| """Export a rendered surface volume (zarr) to numbered layer TIFFs: 00.tif, 01.tif, ... |
| |
| villa's `vc_render_tifxyz --zarr-output` and the team's published surface volumes are |
| zarr stores of shape (layers, height, width), uint8, either a bare array or an OME-Zarr |
| group whose full-resolution level is "0". v8in (YoussefMoNader/ink-8um-v8in) reads a |
| directory of numbered layer images instead, and uses the central 24 when there are more. |
| |
| Chunks of these stores usually hold every layer (w045: 28 x 128 x 128), so reading one |
| layer decompresses all depths in each touched chunk. The export therefore reads all requested |
| layers in row bands into memory when they fit under max_bytes (w045: about 1.4 GB), |
| and falls back to one layer at a time otherwise. Depth order is kept as |
| stored; reverse it at inference time (v8in `--reverse`), not here, so one export serves both |
| directions. Needs numpy, tifffile and zarr (all in villa's environment). |
| |
| The depth-shuffle control (after aviad12g's vesuvius-depth-order-control) is the exception: |
| `shuffle` writes the layers in a fixed random order, never the stored or the reversed one. A |
| model that reads ink loses it there; on PHerc0841 every released reader we tried fell to |
| about 0.5 AUC on shuffled layers, a cleaner null than the reversed volume, which keeps a |
| depth structure. `kit shuffle` writes a zarr (villa's input); `kit layers --shuffle` writes |
| v8in's layer folder in the same order. |
| """ |
|
|
| from pathlib import Path |
|
|
| from .verify import VerifyError, _numpy, missing |
|
|
|
|
| def _zarr(): |
| try: |
| import zarr |
| except ImportError as exc: |
| raise missing("zarr") from exc |
| return zarr |
|
|
|
|
| def open_volume(path, level="0"): |
| """Return the (layers, height, width) array of a surface-volume zarr.""" |
| zarr = _zarr() |
| try: |
| node = zarr.open(str(path), mode="r") |
| except Exception as exc: |
| raise VerifyError(f"cannot open {path} as zarr: {exc}") from exc |
| if not hasattr(node, "shape"): |
| if level not in node: |
| raise VerifyError(f"{path}: zarr group has no level {level!r}; levels: {sorted(node.keys())}") |
| node = node[level] |
| if len(node.shape) != 3: |
| raise VerifyError(f"{path}: expected (layers, height, width), got shape {node.shape}") |
| return node |
|
|
|
|
| def depth_permutation(n, seed): |
| """A fixed random order of n layers that is neither the stored nor the reversed order.""" |
| np = _numpy() |
| if n < 3: |
| raise VerifyError(f"a depth shuffle needs at least 3 layers, got {n}") |
| rng = np.random.default_rng(seed) |
| ident = list(range(n)) |
| while True: |
| perm = [int(i) for i in rng.permutation(n)] |
| if perm != ident and perm != ident[::-1]: |
| return perm |
|
|
|
|
| DEFAULT_MAX_BYTES = 4 * 1024 ** 3 |
| BAND_ROWS = 1024 |
|
|
|
|
| def export_layers(volume, out_dir, start=0, count=None, compression="zlib", crop=None, |
| max_bytes=DEFAULT_MAX_BYTES, order=None): |
| """Write layers start .. start+count-1 of volume as 00.tif, 01.tif, ... in out_dir, which must |
| not already hold numbered layer TIFFs. |
| |
| crop is (y0, y1, x0, x1) in full-resolution pixels, or None for the whole layer. order, a |
| permutation of range(count), writes layer start+order[i] as file i (the shuffle control).""" |
| np = _numpy() |
| import tifffile |
|
|
| total = volume.shape[0] |
| if count is None: |
| count = total - start |
| if start < 0 or count < 1 or start + count > total: |
| raise VerifyError(f"layers {start}..{start + count - 1} are outside the volume's {total} layers") |
| h, w = volume.shape[1:] |
| y0, y1, x0, x1 = 0, h, 0, w |
| if crop is not None: |
| y0, y1, x0, x1 = crop |
| if not (0 <= y0 < y1 <= h and 0 <= x0 < x1 <= w): |
| raise VerifyError(f"crop {list(crop)} is outside the {h} x {w} layer") |
| region = (slice(y0, y1), slice(x0, x1)) |
| if volume.dtype != np.uint8: |
| raise VerifyError(f"the volume is {volume.dtype}; v8in expects uint8 layers") |
| out = Path(out_dir) |
| if out.is_dir() and any(p.suffix == ".tif" and p.stem.isdigit() for p in out.iterdir()): |
| raise VerifyError(f"{out} already holds layer TIFFs from an earlier export, which v8in would read too; " |
| "remove them or give an empty folder") |
| stack = None |
| if count * (y1 - y0) * (x1 - x0) <= max_bytes: |
| stack = np.empty((count, y1 - y0, x1 - x0), dtype=np.uint8) |
| for band in range(y0, y1, BAND_ROWS): |
| end = min(band + BAND_ROWS, y1) |
| stack[:, band - y0:end - y0] = volume[start:start + count, band:end, x0:x1] |
| if order is None: |
| order = list(range(count)) |
| elif sorted(order) != list(range(count)): |
| raise VerifyError(f"order must be a permutation of 0..{count - 1}") |
| out.mkdir(parents=True, exist_ok=True) |
| width = max(2, len(str(count - 1))) |
| written = [] |
| for i, src in enumerate(order): |
| layer = stack[src] if stack is not None else np.asarray(volume[(start + src, *region)]) |
| path = out / f"{i:0{width}d}.tif" |
| tifffile.imwrite(path, layer, compression=compression) |
| written.append(path) |
| return written |
|
|
|
|
| def export_file(src, out_dir, start=0, count=None, level="0", crop=None, max_bytes=DEFAULT_MAX_BYTES, |
| shuffle_seed=None): |
| volume = open_volume(src, level) |
| n = (volume.shape[0] - start) if count is None else count |
| order = None if shuffle_seed is None else depth_permutation(n, shuffle_seed) |
| written = export_layers(volume, out_dir, start, count, crop=crop, max_bytes=max_bytes, order=order) |
| return {"source": str(src), "shape": list(volume.shape), "start": start, "crop": crop, |
| "layers": len(written), "out_dir": str(out_dir), "order": order} |
|
|
|
|
| def shuffle_file(src, dst, seed, level="0", crop=None): |
| """Write src's layers in a fixed random order to a new zarr (array "0"), for villa's reader.""" |
| np = _numpy() |
| zarr = _zarr() |
| if Path(dst).exists(): |
| raise VerifyError(f"{dst} already exists; remove it or give a new path for the shuffled zarr") |
| volume = open_volume(src, level) |
| n, h, w = volume.shape |
| y0, y1, x0, x1 = crop if crop is not None else (0, h, 0, w) |
| if not (0 <= y0 < y1 <= h and 0 <= x0 < x1 <= w): |
| raise VerifyError(f"crop {list(crop)} is outside the {h} x {w} layer") |
| order = depth_permutation(n, seed) |
| data = np.empty((n, y1 - y0, x1 - x0), dtype=volume.dtype) |
| for band in range(y0, y1, BAND_ROWS): |
| end = min(band + BAND_ROWS, y1) |
| data[:, band - y0:end - y0] = volume[:, band:end, x0:x1][order] |
| group = zarr.open_group(str(dst), mode="w-") |
| group.create_array("0", data=data, chunks=(n, min(128, y1 - y0), min(128, x1 - x0))) |
| group.attrs["depth_shuffle"] = {"source": str(src), "seed": seed, "order": order, "crop": crop} |
| return {"source": str(src), "out": str(dst), "shape": list(data.shape), "seed": seed, "order": order} |
|
|
