"""Write one MCAM dataset to storage a block at a time.
``save``/``export`` need the whole dataset resident before they write anything.
For a plate sized z stack that is hundreds of gigabytes, so the operations that
produce one -- an XY scan, a stack reduction -- cannot use them without holding
their own output in RAM. The writers here invert that: the caller hands over one
*block* of the dataset at a time, each block is written as it arrives and then
released, and only the metadata (which carries no pixels) stays resident.
Every storage layout ``owl.mcam_data`` can write is available:
* a directory of per-tile ``.tif``, ``.png``, ``.jpeg``, ``.bmp`` or
``.j2c`` files with a metadata file beside them, as ``export`` writes
(:class:`StreamedExportWriter`),
* the same directory layout with the images in one ``.mp4`` container per
tile (``'multi_mp4'``); ``'tiledmp4'`` composites the whole dataset into
a single video and so cannot be written a block at a time,
* one ``.nc`` file with the images as a chunked HDF5 variable, as ``save``
writes (:class:`StreamedSaveWriter`).
Pick between them the way you pick between ``save`` and ``export``: what the
data is for decides the layout, not the path it is going to.
"""
from functools import partial
from pathlib import Path
from threading import Event
import numpy as np
import xarray as xr
from ..util import make_timestamp, nvidia
from ..util.io_utils import check_disk_space
from ._color import bayer_dataset_to_rgb
from ._core import add_software_timestamp
from ._export import (
_add_imagename_cache_to_metadata,
_add_tiff_refs_to_metadata,
_adjust_images_dims_for_z_stack,
_exporters,
_get_imagej_axes,
_write_metadata,
save_metadata,
split_imagename_format,
validate_export_layout,
)
from ._export_constants import (
_default_imagename_format,
_default_imagename_format_wells,
_default_imagename_format_wells_field_id,
_default_tile_dims,
_expected_compression_factor,
)
from ._netcdf import _get_chunksizes
from ._properties import get_chroma, get_stack_dims, is_well_dataset
from ._ramona_backend import RamonaH5NetCDFStore
from ._util import remove_invalid_images
# A format that writes one file for the whole dataset cannot be streamed: each
# block would reopen and overwrite it. 'tiledmp4' composites every tile into a
# single video, so it is one of those.
_unstreamable_image_export_formats = ('tiledmp4',)
__all__ = [
'StreamedExportWriter',
'StreamedSaveWriter',
'default_imagename_format_for',
]
def _unwritten_images(shape, dtype):
"""A full-shape stand-in for images that have not been produced yet.
The writers describe the whole dataset to the storage layer before any
block exists, which needs an array of the final shape and dtype but none of
its bytes. A zero-strided broadcast is exactly that: it reports the right
shape and dtype while owning one element.
"""
return np.broadcast_to(np.zeros((1,) * len(shape), dtype=dtype), shape)
def _as_indexer(indices):
"""A slice when the indices step evenly, an index array when they do not.
Two things turn on this. ``Dataset.isel`` hands back a *view* for a slice
and a *copy* for an index array, and the metadata merge writes through what
it is given -- so a regular block has to stay a slice or the merge writes
into nothing. And a slice is what the exporters and HDF5 handle best. Most
blocks are regular: one XY scan position steps by the number of positions,
one well is a contiguous run. The irregular ones still work, just through a
longer path.
"""
indices = list(indices)
if len(indices) == 1:
return slice(indices[0], indices[0] + 1)
step = indices[1] - indices[0]
if step > 0 and all(
second - first == step for first, second in zip(indices[:-1], indices[1:])
):
return slice(indices[0], indices[-1] + 1, step)
return np.array(indices)
def _indexer_positions(indexer, size):
"""The positions an indexer selects, as a list."""
if isinstance(indexer, slice):
return list(range(*indexer.indices(size)))
return list(indexer)
class _BlockGrid:
"""The tiles of a dataset, divided into the blocks a writer fills it with.
A block is a set of tile indices, in the shape
:func:`~owl.mcam_data.get_groupings` returns: a list of ``(image_y,
image_x)`` tuples. That is the general form, and it is deliberately the
only one -- a well whose fields are staggered across the plate, or whose
corner tile was never acquired, is not a rectangle, and a model that only
speaks in rectangles either cannot express it or has to round it up to one.
Selecting an arbitrary set of tiles is another matter. Indexing a dataset
pointwise collapses ``image_y`` and ``image_x`` into a single dimension,
and every projection, exporter and metadata path downstream reads the tile
grid by name. So a block is *selected* as the outer product of the indices
it uses along each dimension -- which keeps the grid two dimensional, and
is exactly the block when the block is a rectangle -- and the tiles in that
rectangle the block does not own are masked out of what gets written.
:meth:`is_whole` says which case a block is, so the common one pays
nothing.
"""
__slots__ = ('_blocks', '_dims', '_indexers', '_sizes', '_owned')
def __init__(self, blocks, dims):
self._dims = tuple(dims)
self._blocks = [
[tuple(int(index) for index in tile) for tile in block]
for block in blocks
]
for number, block in enumerate(self._blocks):
if not block:
raise ValueError(f"Block {number} holds no tiles.")
for tile in block:
if len(tile) != len(self._dims):
raise ValueError(
f"Block {number} holds the tile index {tile}, which does "
f"not name a position in {self._dims}."
)
seen = {}
for number, block in enumerate(self._blocks):
for tile in block:
if tile in seen:
raise ValueError(
f"Tile {dict(zip(self._dims, tile))} is in block "
f"{seen[tile]} and block {number}. A tile written twice "
"is a file written twice, so the blocks have to be "
"disjoint."
)
seen[tile] = number
self._sizes = tuple(
max((tile[axis] for tile in seen), default=-1) + 1
for axis in range(len(self._dims))
)
# Worked out once: every write needs them, and a block is small.
self._indexers = [
{
dim: _as_indexer(sorted({tile[axis] for tile in block}))
for axis, dim in enumerate(self._dims)
}
for block in self._blocks
]
self._owned = [None] * len(self._blocks)
def __len__(self):
return len(self._blocks)
@property
def dims(self):
"""The dimensions a tile index names."""
return self._dims
def size(self, dim):
"""How far the whole grid reaches along ``dim``."""
return self._sizes[self._dims.index(dim)]
def indexers(self, index):
"""The outer-product indexers that select block ``index``'s rectangle."""
return self._indexers[index]
def start(self, index, dim):
"""The first position block ``index`` occupies along ``dim``."""
indexer = self._indexers[index][dim]
return indexer.start if isinstance(indexer, slice) else int(indexer[0])
def is_whole(self, index):
"""Whether the block is the whole rectangle its indices span."""
return self.owned(index).all()
def owned(self, index):
"""Which tiles of block ``index``'s rectangle the block actually holds."""
if self._owned[index] is None:
indexers = self._indexers[index]
positions = [
{
position: axis
for axis, position in enumerate(_indexer_positions(
indexers[dim], self._sizes[axis_number],
))
}
for axis_number, dim in enumerate(self._dims)
]
owned = np.zeros(
tuple(
len(_indexer_positions(indexers[dim], self._sizes[axis_number]))
for axis_number, dim in enumerate(self._dims)
),
dtype=bool,
)
for tile in self._blocks[index]:
owned[tuple(
positions[axis][position]
for axis, position in enumerate(tile)
)] = True
self._owned[index] = owned
return self._owned[index]
def tiles(self, index):
"""The tile indices block ``index`` holds."""
return self._blocks[index]
class _StreamedWriter:
"""Accumulate one MCAM dataset from blocks and write it as they arrive.
``blocks`` says which tiles make up each block, in the shape
:func:`~owl.mcam_data.get_groupings` returns: a list of ``(image_y,
image_x)`` tuples per block. Any set of tiles is a block -- the cameras of
one XY scan position, laced across the plate; the fields of one well,
staggered or not; a single tile. The blocks must be disjoint, since a tile
written twice is a file written twice, and together they should cover the
grid. Dimensions the tile indices do not name are not divided at all, and
every block must span them whole.
The full metadata is not declared up front. The first block written builds
the skeleton -- every variable it carries, grown to the full grid along the
expanded dimensions -- and each later block is copied into its place. A
block that carries a dimension longer than the skeleton's (a later position
with more frames, say) grows the skeleton rather than being truncated.
This class holds everything the storage layout does not decide: what the
blocks mean, how their metadata is folded together and when it is written.
Subclasses supply the two operations that do depend on the layout, writing
one block's images and writing the finished metadata.
It is internal, and never instantiated on its own -- it writes nowhere,
since a destination is a thing a layout decides. Use
:class:`StreamedExportWriter` or :class:`StreamedSaveWriter`, which is also
where the behaviour described here is tested from.
Parameters
----------
blocks : sequence or callable, optional
The blocks the dataset is divided into, each a sequence of tile index
tuples as :func:`~owl.mcam_data.get_groupings` gives them. Defaults to
no division, which makes one block covering the whole grid.
A caller that cannot know the tiles up front -- an acquisition learns
how many cameras a position holds only when the first one arrives --
may pass a callable instead. It is called once, with a mapping of each
of ``block_dims`` to the first block's size along it, and returns the
blocks.
block_dims : tuple of str
The dimensions a tile index names, in order. Defaults to
``('image_y', 'image_x')``.
extra_metadata : xarray.Dataset or dict, optional
Metadata folded into the skeleton the first block builds. Use it for
plate-wide information the blocks themselves do not carry;
:meth:`add_extra_metadata` indexes it back down to a single block. This
is not :attr:`metadata`, which is what the blocks have contributed so
far.
stop_event : threading.Event, optional
Checked between the stages of a write. A set event abandons the block
in progress. The metadata is only written by :meth:`close`, so a
stopped stream leaves what it wrote and no metadata to describe it.
tqdm : callable, optional
Progress bar factory. ``export`` gives one bar for the dataset it
writes, so this gives one bar for the blocks that make it up rather
than one per block. It is a plain attribute, so a caller that only
builds its progress bar once the stream is running can assign it later.
"""
def __init__(
self,
*,
blocks=None,
block_dims=('image_y', 'image_x'),
extra_metadata=None,
stop_event=None,
tqdm=None,
):
metadata = extra_metadata
if metadata is not None and not isinstance(metadata, xr.Dataset):
# Each block indexes into these by dimension name, so a plain dict
# of ``(dims, data)`` tuples has to be promoted first. A Dataset is
# already in that form, and xarray refuses to build one out of one.
metadata = xr.Dataset(metadata)
if stop_event is None:
stop_event = Event()
self.stop_event = stop_event
self.export_attrs = {}
self._blocks = blocks
self._block_dims = tuple(block_dims)
# Explicit blocks are settled now, so block_indexers works before the
# first write -- which is how a caller selects the block to write.
self._grid = (
None if blocks is None or callable(blocks)
else _BlockGrid(blocks, self._block_dims)
)
self._extra_metadata = metadata
self._metadata = None
self.tqdm = tqdm
self._closed = False
self._blocks_written = set()
self._progress = None
@property
def metadata(self):
"""The metadata accumulated from the blocks written so far.
``None`` until the first block is written. Assigning to it replaces what
:meth:`close` will write, which is how a caller folds in something no
single block knows.
"""
return self._metadata
@metadata.setter
def metadata(self, value):
self._metadata = value
@property
def closed(self):
"""Whether :meth:`close` has already run."""
return self._closed
def __enter__(self):
return self
def __exit__(self, exception_type, exception_value, traceback):
# A dataset without metadata is not loadable, so only a clean run gets
# one: an exception or a stop leaves the written blocks behind as
# evidence rather than passing them off as a finished dataset.
if exception_type is None and not self.stop_event.is_set():
self.close()
else:
self._release_storage()
return False
def block_indexers(self, index):
"""The indexers that select block ``index``'s place in the full dataset.
Parameters
----------
index : int
Which block.
Returns
-------
indexers : dict
Suitable for ``Dataset.isel``. For a block that is not the whole
rectangle its tiles span, this selects the rectangle; ``owned``
says which of it belongs to the block.
"""
if self._grid is None:
return {}
return self._grid.indexers(index)
def add_extra_metadata(self, block, index=0):
"""Return ``block`` with the writer's extra ``metadata`` folded in.
Variables spanning an expanded dimension are indexed down to this
block; variables the block already carries are left alone. :meth:`write`
does not do this itself, so that a caller can run its own per-block work
-- an analysis, say -- against a block that already carries the
plate-wide metadata.
"""
if self._extra_metadata is None:
return block
# This runs before write(), so it may be the first sight of a block.
self._resolve_grid(block)
indexers = self.block_indexers(index)
block = block.copy()
for key, value in self._extra_metadata.items():
if key in block:
continue
block[key] = value.isel({
dim: indexer
for dim, indexer in indexers.items()
if dim in value.dims
})
return block
def write(self, block, index=0):
"""Write one block of the dataset.
The block's images go to storage immediately and are not referenced
again; its metadata is copied into the full metadata, which
:meth:`close` writes out once at the end.
Parameters
----------
block : xarray.Dataset
An MCAM dataset covering the rectangle block ``index`` spans. It
must span every dimension the tile indices do not name. When the
block does not fill that rectangle, the tiles it does not own are
left alone rather than written.
index : int
Which of the declared blocks this is.
"""
if self._closed:
raise RuntimeError(f"Cannot write to a closed {type(self).__name__}.")
block = self._prepare_block(block)
if self.stop_event.is_set():
return
self._resolve_grid(block)
self._resolve_export_attrs(block, index)
# The attributes have to be resolved against the prepared block, so the
# skeleton is built here rather than in __init__.
if self._metadata is None:
self._check_disk_space(block)
self._metadata = self._create_metadata_skeleton(block, index)
self._open_storage(block)
self.merge_block_metadata(block, index)
# Hand the writer the block's own slice of the full metadata rather
# than the block itself, so a variable the block does not carry is still
# described wherever the images end up.
block_to_write = self._metadata.isel(self.block_indexers(index))
# Assigning a DataArray aligns it against the coordinates already
# there, which would quietly crop a block that does not fit its slot
# instead of writing all of it. Say so rather than losing the pixels.
mismatched = {
dim: (block.sizes[dim], block_to_write.sizes[dim])
for dim in block.images.dims
if dim in block_to_write.sizes and block.sizes[dim] != block_to_write.sizes[dim]
}
if mismatched:
sizes = ', '.join(
f"{dim} is {block_size} where the dataset has {slot_size}"
for dim, (block_size, slot_size) in sorted(mismatched.items())
)
raise ValueError(
f"Block {index} does not fit its slot in the dataset: {sizes}."
)
block_to_write['images'] = block.images
block_to_write = self._mask_unowned_tiles(block_to_write, index)
self._write_block(block_to_write, index)
self._blocks_written.add(index)
if self.tqdm is not None:
if self._progress is None:
self._progress = self.tqdm(total=self._total_blocks())
self._progress.update()
def close(self, *, before_save=None, allow_incomplete=False):
"""Write the metadata and finish the dataset.
Parameters
----------
allow_incomplete : bool
Finish even though some blocks were never written. Only a caller
that knows a block was deliberately skipped -- an acquisition told
which positions hold no sample, say -- should pass this; otherwise
an unwritten block reads back as zeros with nothing to say it was
never there.
before_save : callable, optional
Called as ``before_save(metadata)`` with the consolidated metadata,
and must return the metadata to write. This is the last chance to
fold in something no block carries, or to derive a companion output
from the finished dataset.
Returns
-------
path : pathlib.Path
Where the dataset was written.
"""
if self._closed:
raise RuntimeError(f"This {type(self).__name__} is already closed.")
self._closed = True
if self._progress is not None:
self._progress.close()
self._progress = None
metadata = self._metadata
if metadata is None:
# Nothing was ever written, so there is nothing to describe.
return self.path
expected = self._total_blocks()
if not allow_incomplete and len(self._blocks_written) != expected:
# The metadata describes the whole grid, so writing it now would
# hand off a dataset whose missing tiles read back as zeros with
# nothing to say they were never written.
raise RuntimeError(
f"Only {len(self._blocks_written)} of {expected} blocks were "
"written, so this dataset is incomplete. Write the rest before "
"closing, pass allow_incomplete=True if the missing blocks were "
"skipped deliberately, or abandon the writer to leave the blocks "
"on disk without metadata describing them as finished."
)
metadata = self._consolidate_metadata(metadata)
if before_save is not None:
metadata = before_save(metadata)
self._metadata = metadata
return self._write_metadata(metadata)
# Layout specific hooks. ``_write_block`` and ``_write_metadata`` are the
# two a subclass must supply; the rest have defaults that suit most.
@property
def path(self):
"""Where this writer is writing."""
raise NotImplementedError
def _prepare_block(self, block):
"""Convert one block into the form this layout stores."""
return block
def _resolve_grid(self, block):
"""Settle the blocks, now that one of them can say how big a block is."""
if self._grid is not None or self._blocks is None:
return
blocks = self._blocks
if callable(blocks):
# Why ``blocks`` may be a callable, and what would let it stop
# being one.
#
# Every caller but one knows its tiles up front and passes a plain
# sequence. ``run_xy_scan`` cannot: a block is the cameras of one
# XY position, and how many cameras that is depends on the
# acquisition function it was handed. ``mcam.acquire_high_speed_video``
# takes a ``selection_slice`` and returns a subset of the array, and
# that argument lives inside the opaque ``func_kwargs`` the scan
# forwards without inspecting. So the scan can only learn its own
# block shape by acquiring one, which is after the writer is built.
# tests/analysis/test_scan_functions.py::test_video_subset pins the
# case: a 2 x 2 selection out of a larger array.
#
# To remove this: give the scan the camera selection as a parameter
# of its own rather than something buried in func_kwargs, or have
# MCAM report the shape an acquisition will return before it runs.
# Either lets run_xy_scan build its blocks in __init__ like everyone
# else, and then this branch, the ``_resolve_grid`` calls in
# ``write`` and ``add_extra_metadata``, and the deferred ``_grid``
# can all go, leaving ``blocks`` a plain sequence.
blocks = blocks({
dim: block.sizes[dim]
for dim in self._block_dims
if dim in block.sizes
})
self._grid = _BlockGrid(blocks, self._block_dims)
def _total_blocks(self):
"""How many blocks make up the whole dataset."""
return 1 if self._grid is None else len(self._grid)
def _owned_tiles(self, index):
"""Which tiles of block ``index``'s rectangle it owns, or None for all."""
if self._grid is None or self._grid.is_whole(index):
return None
return self._grid.owned(index)
def _mask_unowned_tiles(self, block, index):
"""Hide the tiles of the rectangle that belong to another block.
The exporters skip a tile that ``valid_data`` marks invalid, which is
how a block that is not a rectangle writes only its own files.
"""
owned = self._owned_tiles(index)
if owned is None:
return block
block = block.copy()
if 'valid_data' in block:
valid = np.asarray(block['valid_data'])
block['valid_data'] = (block['valid_data'].dims, valid & owned)
else:
block['valid_data'] = (self._grid.dims, owned.copy())
return block
def _check_disk_space(self, block):
"""Refuse the write if the whole dataset will not fit.
``save`` and ``export`` measure the dataset they were handed. A streamed
writer has only ever seen one block, so it scales that block by how many
are coming.
"""
def _open_storage(self, block):
"""Create the storage now that the first block has fixed the shape."""
def _release_storage(self):
"""Let go of the storage without finishing the dataset."""
def _write_block(self, block, index):
raise NotImplementedError
def _consolidate_metadata(self, metadata):
return metadata
def _write_metadata(self, metadata):
raise NotImplementedError
def _resolve_export_attrs(self, block, index=0):
"""Fill in the attributes that only the first block can supply."""
if 'images_dims' not in self.export_attrs:
self.export_attrs['images_dims'] = block.images.dims
def _create_metadata_skeleton(self, block, index, *, include_extra=True):
"""Grow one block's metadata into the full grid, filled with blanks.
``include_extra`` is False when this is called to grow a handful of
variables that outgrew the skeleton: the attributes and the extra
metadata already sit on the metadata those get merged back into.
"""
# owl.calibration reaches back into owl.mcam_data, so this cannot be
# imported while this module is.
from owl.calibration._keys import (
all_base_keys_constructors as all_base_calibration_keys_constructors,
)
block_metadata = block.drop_vars(['images'], errors='ignore')
expanded_dims = () if self._grid is None else tuple(
dim
for dim in self._grid.dims
if dim in block_metadata.dims
)
metadata = block_metadata.drop_dims(expanded_dims)
new_coords = {}
for dim in expanded_dims:
# The block's own coordinates are already the global ones, and its
# smallest entry sits exactly where the grid says the block starts,
# so that pins the origin of the full grid.
origin = int(block_metadata[dim].min()) - self._grid.start(index, dim)
new_coords[dim] = np.arange(self._grid.size(dim)) + origin
metadata = metadata.assign_coords(new_coords)
for key, value in block_metadata.items():
if key in metadata:
continue
dims = value.dims
shape = tuple(
# Prefer the expanded size, and fall back to the block's for a
# dimension the expansion does not touch.
metadata.sizes.get(dim, block_metadata.sizes.get(dim))
for dim in dims
)
dtype = value.dtype
if key == 'software_timestamp':
metadata = add_software_timestamp(metadata, dims=dims)
continue
if np.issubdtype(dtype, np.floating) and value.attrs.get(
'__owl_settings_type__'
) in (
'analysis_output_positioned_label',
'analysis_output_bounding_box',
'analysis_output_point',
):
# A missing detection is absent, not at the origin, so these
# start as nan rather than as the zeros everything else gets.
data = np.full(shape=shape, fill_value=np.nan, dtype=dtype)
else:
data = all_base_calibration_keys_constructors.get(key, np.zeros)(
shape=shape, dtype=dtype
)
metadata[key] = (dims, data)
metadata[key].attrs = value.attrs
metadata[key].encoding = value.encoding
if 'valid_data' in metadata:
# Start invalid everywhere: a stream abandoned half way through then
# describes only what it actually wrote.
metadata['valid_data'][...] = False
if include_extra:
for key, value in self.export_attrs.items():
metadata.attrs[key] = value
if self._extra_metadata is not None:
for key, value in self._extra_metadata.items():
metadata[key] = value
return metadata
def merge_block_metadata(self, block_metadata, index=0):
"""Record one block's metadata without writing any of its images.
:meth:`write` calls this for every block it writes. Call it directly to
describe a block whose images are not being written -- a position the
scan skipped, say -- so the metadata still spans the whole grid.
"""
block_metadata = block_metadata.drop_vars(['images'], errors='ignore')
if self._metadata is None:
self._metadata = self._create_metadata_skeleton(block_metadata, index)
metadata = self._metadata
keys_to_update = []
for key in block_metadata:
if key not in metadata:
keys_to_update.append(key)
continue
for dim in block_metadata[key].dims:
if self._grid is not None and dim in self._grid.dims:
continue
if block_metadata.sizes[dim] > metadata.sizes[dim]:
keys_to_update.append(key)
break
if keys_to_update:
# This block is longer along some dimension than the skeleton the
# first one built. Rebuild those variables at the larger size and
# copy what has already been recorded into the corner.
updated_keys = self._create_metadata_skeleton(
block_metadata[keys_to_update], index, include_extra=False,
)
keys_to_drop = []
for key in updated_keys:
if key not in metadata:
continue
keys_to_drop.append(key)
updated_keys[key].attrs.update(metadata[key].attrs)
updated_keys[key].data[
tuple(slice(None, metadata.sizes[dim]) for dim in metadata[key].dims)
] = metadata[key].data
metadata = metadata.drop_vars(keys_to_drop)
# The variables being replaced were dropped just above, so there
# is nothing to reconcile; state the compat rather than inherit
# whichever one xarray defaults to.
metadata = xr.merge([metadata, updated_keys], compat='no_conflicts')
metadata_selection = metadata.isel(self.block_indexers(index))
owned = self._owned_tiles(index)
for key in metadata_selection:
if key not in block_metadata:
continue
new_axis = tuple(
axis
for axis, dim in enumerate(metadata_selection[key].dims)
if dim not in block_metadata[key].dims
)
if len(metadata_selection[key].dims) > 0:
slice_tuple = tuple(
slice(None, block_metadata.sizes[dim])
for dim in block_metadata[key].dims
)
else:
slice_tuple = Ellipsis
# The ellipsis is used to ensure that the data
# is copied to the correct location
# Consider the following code
# >>> a = np.zeros((2))
# >>> np.copyto(a[0], 2)
# will not work -- TypeError: copyto() argument 1 must be numpy.ndarray, not numpy.float64 # noqa
# instead
# >>> np.copyto(a[0][...], 2)
# will work
destination = metadata_selection[key].data[slice_tuple]
source = np.expand_dims(block_metadata[key].data, new_axis)
if owned is None:
np.copyto(destination, source)
else:
# Only this block's own tiles: the rest of the rectangle
# belongs to another block, which will write its own values.
np.copyto(
destination, source,
where=self._broadcast_ownership(
owned, metadata_selection[key].dims, destination.shape,
),
)
self._metadata = metadata
def _broadcast_ownership(self, owned, dims, shape):
"""Line the tile ownership mask up with a variable's dimensions."""
axes = tuple(dims.index(dim) for dim in self._grid.dims if dim in dims)
mask = owned
if len(axes) != len(self._grid.dims):
# The variable does not span every grid dimension, so a tile it
# does cover is owned if any tile behind it is.
keep = tuple(
axis for axis, dim in enumerate(self._grid.dims) if dim in dims
)
mask = owned.any(axis=tuple(
axis for axis in range(owned.ndim) if axis not in keep
))
expanded = np.expand_dims(
mask, tuple(axis for axis in range(len(shape)) if axis not in axes),
)
return np.broadcast_to(expanded, shape)
def _full_images_shape(self, block):
"""The shape the whole images variable will have once every block lands."""
if self._grid is None:
return tuple(block.images.shape)
return tuple(
self._grid.size(dim) if dim in self._grid.dims else size
for dim, size in zip(block.images.dims, block.images.shape)
)
def default_imagename_format_for(dataset):
"""The per-format default tile names ``export`` would choose for a dataset.
This is a property of the whole dataset, not of any one block: a block that
happens to hold only field 0 of each well would otherwise be named as
though the plate had one field per well, and every field of a well would
collide on the same filename. A caller that has the whole dataset in hand
should resolve it once, up front, and pass the result to
:class:`StreamedExportWriter` as ``default_imagename_format``.
"""
if not is_well_dataset(dataset):
return _default_imagename_format
if 'field_id' not in dataset or dataset.field_id.max() == 0:
return _default_imagename_format_wells
return _default_imagename_format_wells_field_id
[docs]
class StreamedExportWriter(_StreamedWriter):
"""Stream a dataset into a directory of image files, as ``export`` writes.
Each block's tiles become their own image files as soon as the block
arrives, and a single metadata file describing all of them is written at
:meth:`close`. Every format ``export`` supports is available, chosen the
same way: ``image_export_format`` picks the layout and the extension of
``imagename_format`` picks the file type -- ``.tif``, ``.png``, ``.jpeg``,
``.bmp``, ``.j2c`` for the ``'images'`` and ``'flat_image_stacks'``
layouts, ``.mp4`` for ``'multi_mp4'`` and ``'tiledmp4'``.
Parameters
----------
directory : path-like
The directory the dataset is written to. It is created on the first
write.
save_mode_options : optional
An assay ``SaveOptionsModel``. Its fields supply the defaults for
``image_export_format``, ``imagename_format``, ``metadata_filename``,
``save_filename_separator``, ``compression``, ``compression_args`` and
``pyramid_levels``; anything passed explicitly wins.
default_imagename_format : dict, optional
The per-format default image names to fall back on, for callers that
name their tiles by well rather than by camera.
include_timestamp : bool
Append a timestamp to ``directory``. The resulting path is available as
:attr:`path` before the first block is written.
convert_bayer_to_rgb_metadata_only : bool, optional
What to do with a bayer block. ``None``, the default, leaves it alone,
which is what ``export`` does -- bayer is a chroma the image layouts
write as it is. ``True`` relabels the block as RGB without demosaicing
it, for a pipeline that demosaics as it encodes. ``False`` demosaics up
front, which costs a copy three times the size of the block.
mode : {'w', 'x'}
Passed to the exporter for every image file, as ``export`` passes it.
``'x'`` refuses to overwrite, which suits a stream into a directory
that is meant to be new; not every image format honours it.
remove_invalid : bool
Drop rows and columns of the tile grid that no block reported as valid.
An acquisition wants this, because a position it skipped has no files
behind it. A dataset being rewritten does not: ``export`` keeps invalid
tiles in the metadata, and matching it keeps the grid the same shape
through the round trip. Defaults to False, as ``export`` behaves.
Examples
--------
Reduce a plate sized z stack one tile at a time, keeping one tile of source
and one tile of output resident rather than the whole stack:
>>> blocks = [[(iy, ix)] for iy, ix in np.ndindex(
... dataset.sizes['image_y'], dataset.sizes['image_x'])]
>>> writer = StreamedExportWriter(directory, blocks=blocks)
>>> for index, tiles in enumerate(blocks):
... tile = dataset.isel(image_y=[tiles[0][0]], image_x=[tiles[0][1]])
... writer.write(project(tile), index)
>>> directory = writer.close()
A plate reduced well by well is the same call with the groups
:func:`~owl.mcam_data.get_groupings` found:
>>> blocks = list(get_groupings(dataset, 'well_id').values())
"""
def __init__(
self,
directory,
*,
blocks=None,
block_dims=('image_y', 'image_x'),
extra_metadata=None,
save_mode_options=None,
image_export_format=None,
imagename_format=None,
metadata_filename=None,
save_filename_separator=None,
default_imagename_format=None,
tile_dims=None,
compression=None,
compression_args=None,
pyramid_levels=None,
save_tiff_references=None,
include_timestamp=True,
disk_space_tolerance=0.01E9,
convert_bayer_to_rgb_metadata_only=None,
mode='w',
remove_invalid=False,
tqdm=None,
stop_event=None,
**kwargs,
):
super().__init__(
blocks=blocks,
block_dims=block_dims,
extra_metadata=extra_metadata,
stop_event=stop_event,
tqdm=tqdm,
)
if save_mode_options is not None:
if image_export_format is None:
image_export_format = save_mode_options.image_export_format
if imagename_format is None:
imagename_format = save_mode_options.imagename_format
if metadata_filename is None:
metadata_filename = save_mode_options.metadata_filename
if save_filename_separator is None:
save_filename_separator = save_mode_options.save_filename_separator
if compression is None:
compression = save_mode_options.compression
if compression_args is None:
compression_args = save_mode_options.compression_args
if pyramid_levels is None:
pyramid_levels = save_mode_options.pyramid_levels
if image_export_format is None:
image_export_format = 'images'
if image_export_format not in _exporters:
raise ValueError(
f"Unknown export format {image_export_format}. "
f"Supported formats are {sorted(_exporters)}."
)
if image_export_format in _unstreamable_image_export_formats:
raise ValueError(
f"The {image_export_format} export format writes one file for the "
"whole dataset, so it cannot be written a block at a time -- every "
"block would overwrite the last. Use 'multi_mp4', which writes one "
"file per tile, or owl.mcam_data.export for a dataset that is "
"already in memory."
)
if metadata_filename is None:
metadata_filename = 'metadata.nc'
if imagename_format is None and default_imagename_format is not None:
imagename_format = default_imagename_format[image_export_format]
if tile_dims is None:
# A format that does not declare one works it out from the first
# block's image dimensions instead.
tile_dims = _default_tile_dims.get(image_export_format)
directory = Path(directory)
if include_timestamp:
directory = directory.parent / (directory.name + '_' + make_timestamp())
imagename_prefix = ''
if save_filename_separator is not None:
imagename_prefix = directory.name + save_filename_separator
directory = directory.parent
metadata_filename = imagename_prefix + metadata_filename
if imagename_format is not None:
imagename_format = imagename_prefix + imagename_format
export_func = _exporters[image_export_format]
if image_export_format in ('images', 'flat_image_stacks'):
export_func = partial(
export_func,
compression=compression,
compression_args=compression_args,
pyramid_levels=pyramid_levels,
)
elif image_export_format == 'multi_mp4':
gpu_capabilities = nvidia.get_gpu_capabilities()
if len(gpu_capabilities) == 0:
N_simultaneous_streams = 1
else:
N_simultaneous_streams = min(2, gpu_capabilities[0]['number_of_nvenc'])
export_func = partial(
export_func,
N_simultaneous_streams=N_simultaneous_streams,
)
if kwargs:
# ``export`` forwards whatever it does not recognise to the
# exporter -- max_workers, predictor, codec, rowsperstrip -- and so
# does this.
export_func = partial(export_func, **kwargs)
self.directory = directory
# The attributes every block is stamped with. ``images_dims`` and the
# ImageJ axes are only knowable once a block has been seen, so the first
# write fills them in and every later one reuses them.
self.export_attrs = {
'image_export_format': image_export_format,
'image_export_tile_dims': tile_dims,
'preapplied_exif_transform': 1,
'imagename_format': imagename_format,
}
self._metadata_filename = metadata_filename
self._default_imagename_format = default_imagename_format
self._save_tiff_references = save_tiff_references
self._disk_space_tolerance = disk_space_tolerance
self._imagename_prefix = imagename_prefix
self._image_export_format = image_export_format
self._export_func = export_func
self._convert_bayer_to_rgb_metadata_only = convert_bayer_to_rgb_metadata_only
self._mode = mode
self._remove_invalid = remove_invalid
self._relabelled_bayer = False
self._validated = False
@property
def path(self):
return self.directory
def _prepare_block(self, block):
# We do not want to ever export bayer data. This could be done by the
# caller, but we force it here to ensure it is done.
if (
get_chroma(block) != 'bayer'
or self._convert_bayer_to_rgb_metadata_only is None
):
# ``export`` writes bayer as bayer, so by default so does this.
return block
if self._convert_bayer_to_rgb_metadata_only:
# Only the labelling changes here: the images stay bayer and the
# export pipeline demosaics them as it writes, so the images gain
# an ``rgb`` dimension that the block itself does not show yet.
self._relabelled_bayer = True
return block.assign_coords(rgb=['r', 'g', 'b'])
return bayer_dataset_to_rgb(block)
def _resolve_imagename_format(self, block):
"""Name the tiles the way ``export`` would name this dataset's."""
default_imagename_format = (
self._default_imagename_format
if self._default_imagename_format is not None
else default_imagename_format_for(block)
)
imagename_format = default_imagename_format[self._image_export_format]
if self._image_export_format == 'flat_image_stacks':
stack_dims = get_stack_dims(block)
if stack_dims:
stem, extension = split_imagename_format(
default_imagename_format['images']
)
stack_parts = '_'.join(f'{{{dim}}}' for dim in stack_dims)
imagename_format = f'{stem}_{stack_parts}{extension}'
return self._imagename_prefix + imagename_format
def _resolve_export_attrs(self, block, index=0):
if self.export_attrs['imagename_format'] is None:
self.export_attrs['imagename_format'] = self._resolve_imagename_format(block)
if (images_dims := self.export_attrs.get('images_dims')) is None:
images_dims = block.images.dims
if self._relabelled_bayer:
images_dims = images_dims + ('rgb',)
self.export_attrs['images_dims'] = images_dims
if (tile_dims := self.export_attrs.get('image_export_tile_dims')) is None:
if self._image_export_format == 'flat_image_stacks':
# A flat stack puts every frame in its own file, so the stack
# dimensions are flattened into the tile grid rather than kept
# inside one image.
tile_dims = tuple(
dim
for dim in images_dims
if dim in get_stack_dims(block) or dim in ('image_y', 'image_x')
)
else:
tile_dims = tuple(
dim
for dim in images_dims
if dim not in ('y', 'x', 'rgb', 'rgba')
)
divided_but_not_tiled = () if self._grid is None else tuple(
dim
for dim in self._grid.dims
if dim in block.images.dims and dim not in tile_dims
)
if divided_but_not_tiled:
raise ValueError(
f"Cannot divide the dataset along {sorted(divided_but_not_tiled)}: "
f"the {self._image_export_format} layout keeps those dimensions "
"inside a single file, so each block would rewrite the file the "
"last one wrote rather than adding to it. Divide along the tile "
f"dimensions {sorted(tile_dims)} instead, or pass tile_dims so "
"that each block writes its own files."
)
# Every check export makes against the dataset, made against the first
# block. A block is a whole dataset in every respect the checks care
# about -- chroma, image dimensions, stack dimensions -- so what holds
# for it holds for the rest. This also settles the order of tile_dims,
# which is persisted into the metadata.
if not self._validated:
tile_dims = validate_export_layout(
block,
image_export_format=self._image_export_format,
imagename_format=self.export_attrs['imagename_format'],
tile_dims=tile_dims,
)
self._validated = True
self.export_attrs['image_export_tile_dims'] = tile_dims
if Path(self.export_attrs['imagename_format']).suffix.lower() in ('.tif', '.tiff'):
# This logic must match the logic in _save_tiff_with_metadata
per_tiff_dims = tuple(
dim
for dim in images_dims
if dim not in tile_dims
)
per_tiff_dims = _adjust_images_dims_for_z_stack(per_tiff_dims, block)
imagej_axes, axes_transposer = _get_imagej_axes(per_tiff_dims)
self.export_attrs['image_export_imagej_axes'] = imagej_axes
self.export_attrs['image_export_imagej_axes_transposer'] = axes_transposer
def _write_block(self, block, index):
self._export_func(
block,
directory=self.directory,
metadata_filename=None,
mode=self._mode,
tile_dims=self.export_attrs['image_export_tile_dims'],
stop_event=self.stop_event,
)
def _check_disk_space(self, block):
expected_compression = _expected_compression_factor[self._image_export_format]
if 'valid_data' in block:
fraction_valid = float(block.valid_data.mean())
else:
fraction_valid = 1.
check_disk_space(
self.directory, 1,
block.nbytes * self._total_blocks() * expected_compression * fraction_valid,
absolute_tolerance=self._disk_space_tolerance,
)
def _consolidate_metadata(self, metadata):
if self._remove_invalid:
# A row or column no block ever wrote has no files behind it, so it
# is dropped rather than described as empty.
metadata = remove_invalid_images(metadata)
metadata = _add_imagename_cache_to_metadata(metadata)
extension = Path(self.export_attrs['imagename_format']).suffix.lower()
save_tiff_references = self._save_tiff_references
if save_tiff_references is None:
save_tiff_references = extension in ('.tif', '.tiff')
if save_tiff_references and extension not in ('.tif', '.tiff'):
raise RuntimeError(
'save_tiff_references is only compatible with .tif files. '
f'not {extension} files')
if save_tiff_references:
metadata = _add_tiff_refs_to_metadata(metadata, self.directory)
return metadata
def _write_metadata(self, metadata):
return save_metadata(
metadata,
self.directory,
metadata_filename=self._metadata_filename,
include_timestamp=False,
mode=self._mode,
)
[docs]
class StreamedSaveWriter(_StreamedWriter):
"""Stream a dataset into one ``.nc`` file, as ``save`` writes.
The images become a single chunked HDF5 variable, and each block is written
into its own region of it. The variable has to be created before any of it
can be written, so unlike :class:`StreamedExportWriter` the first block
fixes the shape of the images for good: a later block that outgrows it
along an unexpanded dimension is an error rather than a resize. The
metadata, which is written at :meth:`close`, is still free to grow.
Blocks are written into the file in whatever order they arrive, so a block
whose region does not line up with the HDF5 chunks costs a
read-modify-write of the chunks it straddles. Prefer blocks that partition
the leading dimensions -- one tile, one row of tiles, one channel -- over
interleaved blocks, whose stride touches every chunk.
Parameters
----------
filename : path-like
The file to write. A ``.nc`` extension is added when there is none.
include_timestamp : bool
Append a timestamp to the filename, before the extension.
mode : {'w', 'x'}
As ``save`` takes it. ``'x'`` refuses to overwrite an existing file.
chunksizes : tuple of int, optional
The HDF5 chunk shape for the images. Defaults to what ``save_video``
would choose for a dataset of this shape.
"""
def __init__(
self,
filename,
*,
blocks=None,
block_dims=('image_y', 'image_x'),
extra_metadata=None,
include_timestamp=True,
engine='ramona',
disk_space_tolerance=0.1E9,
mode='w',
chunksizes=None,
tqdm=None,
stop_event=None,
):
super().__init__(
blocks=blocks,
block_dims=block_dims,
extra_metadata=extra_metadata,
stop_event=stop_event,
tqdm=tqdm,
)
if mode not in ('x', 'w'):
raise ValueError(
f"Invalid file creation mode ({mode}). Valid options are 'w' or 'x'"
)
filename = Path(filename)
# '.nc' is the only extension this writer produces, so it is the only
# one to strip. Path.stem would treat a dot inside the dataset name,
# such as the '.02step' in '..._11z_.02step_xyzc_stack_...', as an
# extension and truncate the name at it.
stem = filename.name.removesuffix('.nc')
if include_timestamp:
stem = stem + '_' + make_timestamp()
self.filename = filename.parent / (stem + '.nc')
self._mode = mode
self._engine = engine
self._disk_space_tolerance = disk_space_tolerance
self._chunksizes = chunksizes
self._store = None
self._images_h5ds = None
self._images_chunksizes = None
self._images_shape = None
@property
def path(self):
return self.filename
def _check_disk_space(self, block):
check_disk_space(
self.filename, 1, block.nbytes * self._total_blocks(),
absolute_tolerance=self._disk_space_tolerance,
)
def _open_storage(self, block):
"""Create the file and the full-size images variable it will be filled into."""
images_dims = self.export_attrs['images_dims']
images_shape = self._full_images_shape(block)
images = xr.Variable(
images_dims,
_unwritten_images(images_shape, block.images.dtype),
attrs=dict(block.images.attrs),
)
# ``original_shape`` describes the file the block was read out of, and
# ``dtype`` is xarray's record of an on-disk encoding neither of which
# applies to the file being created here.
images.encoding = {
key: value
for key, value in block.images.encoding.items()
if key not in ('original_shape', 'dtype')
}
chunksizes = self._chunksizes
if chunksizes is None:
# save_video keeps a chunk layout the dataset already carries and
# only works one out when there is none. The block's layout describes
# one block, so it is only meaningful when it fits the whole.
chunksizes = images.encoding.get('chunksizes')
if chunksizes is not None and len(chunksizes) != len(images_shape):
chunksizes = None
if chunksizes is None:
skeleton = self._metadata.assign(images=images)
chunksizes = _get_chunksizes(skeleton)
chunksizes = tuple(
min(chunk, size) for chunk, size in zip(chunksizes, images_shape)
)
images.encoding['chunksizes'] = chunksizes
self._images_chunksizes = chunksizes
self._images_shape = images_shape
# Write the coordinates the images variable depends on first, so the
# file has the dimensions to hang it off. The rest of the metadata is
# written at close, once the blocks have finished contributing to it.
keep = set(images_dims) | {
'__owl_version__', '__sys_version__', '__owl_sys_info__',
}
coords_metadata = self._metadata.drop_vars(
self._metadata.variables.keys() - keep
)
coords_metadata.to_netcdf(
self.filename, format='NETCDF4', engine=self._engine, mode=self._mode,
)
self._store = RamonaH5NetCDFStore.open(self.filename, mode='a')
target, _data = self._store.prepare_variable('images', images)
self._images_h5ds = target.get_array()._h5ds
def _write_block(self, block, index):
indexers = self.block_indexers(index)
images = np.asarray(block.images)
images_dims = self.export_attrs['images_dims']
region = tuple(
indexers.get(dim, slice(None)) for dim in images_dims
)
region_shape = tuple(
len(_indexer_positions(indexer, size))
for indexer, size in zip(region, self._images_h5ds.shape)
)
if images.shape != region_shape:
raise ValueError(
f"Block of shape {images.shape} does not fit the region "
f"{region_shape} of the images variable. The first block "
"written fixes the shape of a netCDF stream."
)
owned = self._owned_tiles(index)
if owned is None and all(isinstance(part, slice) for part in region):
# One hyperslab. h5py takes a tuple of slices directly, and this is
# what a block of evenly stepped tiles is.
self._images_h5ds[region] = images
return
# Either the block is not the whole rectangle it spans, or its tiles do
# not step evenly. Place them one at a time rather than overwriting a
# neighbour's, or asking HDF5 for a selection it cannot express.
grid_axes = tuple(images_dims.index(dim) for dim in self._grid.dims)
positions = {
axis: _indexer_positions(region[axis], self._images_h5ds.shape[axis])
for axis in grid_axes
}
if owned is None:
owned = np.ones(
tuple(len(positions[axis]) for axis in grid_axes), dtype=bool,
)
for offsets in np.argwhere(owned):
source = [slice(None)] * images.ndim
destination = [slice(None)] * images.ndim
for axis, offset in zip(grid_axes, offsets):
source[axis] = int(offset)
destination[axis] = positions[axis][int(offset)]
self._images_h5ds[tuple(destination)] = images[tuple(source)]
def _release_storage(self):
if self._store is not None:
self._store.close()
self._store = None
self._images_h5ds = None
def _plan_metadata_chunks(self, metadata):
"""Chunk the mask variables the way ``save_video`` chunks them.
``save_video`` plans a chunk layout for ``images`` and for anything
shaped like it -- an ``analysis_output_mask`` -- rather than letting
HDF5 pick one. The metadata is written through ``save_metadata`` here,
which does not plan, so the plan is stamped on first.
"""
if self._images_chunksizes is None:
return metadata
for name in metadata.data_vars:
variable = metadata[name]
if variable.attrs.get('__owl_settings_type__') != 'analysis_output_mask':
continue
if variable.shape != self._images_shape:
continue
variable.encoding['chunksizes'] = self._images_chunksizes
return metadata
def _write_metadata(self, metadata):
# The metadata is written through xarray, which opens the file itself.
self._release_storage()
metadata = self._plan_metadata_chunks(metadata)
# ``images`` is already in the file; writing it again would mean
# holding the whole stack in memory, which is what this class exists
# to avoid.
_write_metadata(
metadata.drop_vars('images', errors='ignore'),
self.filename,
mode='a',
)
return self.filename
[docs]
def close(self, *, before_save=None, allow_incomplete=False):
"""Write the metadata, close the file and finish the dataset."""
try:
return super().close(
before_save=before_save, allow_incomplete=allow_incomplete,
)
finally:
self._release_storage()