def fuse(
sims: list,
transform_key: str = None,
fusion_func: Callable = weighted_average_fusion,
fusion_method_kwargs: dict = None,
weights_func: Callable = None,
weights_func_kwargs: dict = None,
output_spacing: dict[str, float] = None,
output_stack_mode: str = "union",
output_origin: dict[str, float] = None,
output_shape: dict[str, int] = None,
output_stack_properties: BoundingBox = None,
output_chunksize: Union[int, dict[str, int]] = None,
overlap_in_pixels: int = None,
interpolation_order: int = 1,
blending_widths: dict[str, float] = None,
output_zarr_url: str | None = None,
zarr_options: dict | None = None,
batch_options: dict | None = None,
):
"""
Fuse input views.
This function fuses all (Z)YX views ("fields") contained in the
input list of images, which can additionally contain C and T dimensions.
Parameters
----------
sims : list of SpatialImage
Input views.
transform_key : str, optional
Which (extrinsic coordinate system) to use as transformation parameters.
By default None (intrinsic coordinate system).
fusion_func : Callable, optional
Fusion function to be applied. This function receives the following
inputs (as arrays if applicable): transformed_views, blending_weights, fusion_weights, params.
By default weighted_average_fusion
fusion_method_kwargs : dict, optional
weights_func : Callable, optional
Function to calculate fusion weights. This function receives the
following inputs: transformed_views (as spatial images), params.
It returns (non-normalized) fusion weights for each view.
By default None.
weights_func_kwargs : dict, optional
output_spacing : dict, optional
Spacing of the fused image for each spatial dimension, by default None
output_stack_mode : str, optional
Mode to determine output stack properties. Can be one of
"union", "intersection", "sample". By default "union"
output_origin : dict, optional
Origin of the fused image for each spatial dimension, by default None
output_shape : dict, optional
Shape of the fused image for each spatial dimension, by default None
output_stack_properties : dict, optional
Dictionary describing the output stack with keys
'spacing', 'origin', 'shape'. Other output_* are ignored
if this argument is present.
output_chunksize : int or dict, optional
Chunksize of the dask data array of the fused image. If the first tile is a chunked dask array,
its chunksize is used as the default. If the first tile is not a chunked dask array,
the default chunksize defined in spatial_image_utils.py is used.
output_zarr_url : str or None, optional
If not None, fuse directly into a Zarr store at this location and do so in batches of chunks,
with each chunk being processed independently. This allows for efficient memory usage and
works well for very large datasets (successfully tested ~0.5PB on a macbook).
When provided, fuse() performs eager fusion and returns a SpatialImage backed by the written store.
zarr_options: dict, optional
Additional (dict of) options to pass when creating the Zarr store. Keys:
- ome_zarr : bool, optional
If True and output_zarr_url is provided, write a NGFF/OME-Zarr multiscale image under
"<output_zarr_url>/". Otherwise, the fused array is written directly under output_zarr_url.
- ngff_version : str, optional
NGFF version used when ome_zarr=True. Default "0.4".
- zarr_array_creation_kwargs: dict = None, optional
Additional keyword arguments to pass when creating the Zarr array.
- overwrite: bool, by default True
batch_options : dict, optional
Options for chunked fusion when output_zarr_url is provided. Keys:
- batch_func: Callable, optional
Function to process each batch of fused chunks. Inputs:
1) a list of block_id(s)
2) function that performs fusion when passed a given block_id
By default None, in which case the each block is processed sequentially.
- n_batch: int
Number of blocks to process in each batch
(n_batch>1 only compatible with a provided batch_func). By default 1.
- batch_func_kwargs: dict, optional
Additional keyword arguments passed to batch_func.
Returns
-------
SpatialImage
Fused image.
"""
# If writing directly to Zarr/OME-Zarr, run chunked fusion path and return eagerly.
if output_zarr_url is not None:
# Collect batch options with defaults
batch_options = batch_options or {}
batch_func = batch_options.get("batch_func", None)
n_batch = batch_options.get("n_batch", 1)
batch_func_kwargs = batch_options.get("batch_func_kwargs", None)
zarr_array_creation_kwargs = batch_options.get("zarr_array_creation_kwargs", None)
# Collect zarr options with defaults
zarr_options = zarr_options or {}
ome_zarr = zarr_options.get("ome_zarr", False)
ngff_version = zarr_options.get("ngff_version", "0.4")
overwrite = zarr_options.get("overwrite", True)
# Resolve store path for data (OME-Zarr stores scale 0 under "<root>/0")
store_url = os.path.join(output_zarr_url, "0") if ome_zarr else output_zarr_url
if overwrite and os.path.exists(store_url):
shutil.rmtree(store_url)
if ome_zarr:
# Ensure creation kwargs reflect NGFF version when writing OME-Zarr
zarr_array_creation_kwargs = ngff_utils.update_zarr_array_creation_kwargs_for_ngff_version(
ngff_version, zarr_array_creation_kwargs
)
# Build kwargs for per-chunk fuse() calls (exclude zarr-specific args to avoid recursion)
per_chunk_fuse_kwargs = {
"sims": sims,
"transform_key": transform_key,
"fusion_func": fusion_func,
"fusion_method_kwargs": fusion_method_kwargs,
"weights_func": weights_func,
"weights_func_kwargs": weights_func_kwargs,
"output_spacing": output_spacing,
"output_stack_mode": output_stack_mode,
"output_origin": output_origin,
"output_shape": output_shape,
"output_stack_properties": output_stack_properties,
"output_chunksize": output_chunksize,
"overlap_in_pixels": overlap_in_pixels,
"interpolation_order": interpolation_order,
"blending_widths": blending_widths,
}
# Prepare block fusion and process in batches
block_fusion_info = prepare_block_fusion(
store_url,
fuse_kwargs=per_chunk_fuse_kwargs,
zarr_array_creation_kwargs=zarr_array_creation_kwargs,
)
fuse_chunk = block_fusion_info["func"]
nblocks = block_fusion_info["nblocks"]
osp = block_fusion_info["output_stack_properties"]
osp["shape"] = {dim: int(v) for dim, v in osp["shape"].items()}
print(f'Fusing {np.prod(nblocks)} blocks in batches of {n_batch}...')
for batch in tqdm(
misc_utils.ndindex_batches(nblocks, n_batch),
total=int(np.ceil(np.prod(nblocks) / n_batch)),
):
if batch_func is None:
for block_id in batch:
fuse_chunk(block_id)
else:
batch_func(fuse_chunk, batch, **(batch_func_kwargs or {}))
# Build SpatialImage from zarr array
fusion_transform_key = transform_key
fused = si_utils.get_sim_from_array(
array=da.from_zarr(store_url),
dims=list(sims[0].dims),
transform_key=fusion_transform_key,
scale=osp["spacing"],
translation=osp["origin"],
c_coords=sims[0].coords["c"].values,
t_coords=sims[0].coords["t"].values,
)
# If requested, write OME-Zarr metadata
# and multiscale pyramid
if ome_zarr:
ngff_utils.write_sim_to_ome_zarr(
fused,
output_zarr_url=output_zarr_url,
overwrite=False,
batch_options=batch_options,
)
return fused
# Default in-memory fusion path (unchanged)
output_chunksize = process_output_chunksize(sims, output_chunksize)
output_stack_properties = process_output_stack_properties(
sims=sims,
output_spacing=output_spacing,
output_origin=output_origin,
output_shape=output_shape,
output_stack_properties=output_stack_properties,
output_stack_mode=output_stack_mode,
transform_key=transform_key,
)
sdims = si_utils.get_spatial_dims_from_sim(sims[0])
nsdims = si_utils.get_nonspatial_dims_from_sim(sims[0])
params = [
si_utils.get_affine_from_sim(sim, transform_key=transform_key)
for sim in sims
]
# determine overlap from weights method
# (soon: fusion methods will also require overlap)
overlap_in_pixels = 0
if weights_func is not None:
overlap_in_pixels = np.max(
[
overlap_in_pixels,
weights.calculate_required_overlap(
weights_func, weights_func_kwargs
),
]
)
# calculate output chunk bounding boxes
output_chunk_bbs, block_indices = mv_graph.get_chunk_bbs(
output_stack_properties, output_chunksize
)
# add overlap to output chunk bounding boxes
output_chunk_bbs_with_overlap = [
output_chunk_bb
| {
"origin": {
dim: output_chunk_bb["origin"][dim]
- overlap_in_pixels * output_stack_properties["spacing"][dim]
for dim in sdims
}
}
| {
"shape": {
dim: output_chunk_bb["shape"][dim] + 2 * overlap_in_pixels
for dim in sdims
}
}
for output_chunk_bb in output_chunk_bbs
]
views_bb = [si_utils.get_stack_properties_from_sim(sim) for sim in sims]
merges = []
for ns_coords in itertools.product(
*tuple([sims[0].coords[nsdim] for nsdim in nsdims])
):
sim_coord_dict = {
ndsim: ns_coords[i] for i, ndsim in enumerate(nsdims)
}
params_coord_dict = {
ndsim: ns_coords[i]
for i, ndsim in enumerate(nsdims)
if ndsim in params[0].dims
}
# ssims = [sim.sel(sim_coord_dict) for sim in sims]
sparams = [param.sel(params_coord_dict) for param in params]
# should this be done within the loop over output chunks?
fix_dims = []
for dim in sdims:
other_dims = [odim for odim in sdims if odim != dim]
if (
any((param.sel(x_in=dim, x_out=dim) - 1) for param in sparams)
or any(
any(param.sel(x_in=dim, x_out=other_dims))
for param in sparams
)
or any(
any(param.sel(x_in=other_dims, x_out=dim))
for param in sparams
)
or any(
output_stack_properties["spacing"][dim]
- views_bb[iview]["spacing"][dim]
for iview in range(len(sims))
)
or any(
float(
output_stack_properties["origin"][dim]
- param.sel(x_in=dim, x_out="1")
)
% output_stack_properties["spacing"][dim]
for param in sparams
)
):
continue
fix_dims.append(dim)
fused_output_chunks = np.empty(
np.max(block_indices, 0) + 1, dtype=object
)
for output_chunk_bb, output_chunk_bb_with_overlap, block_index in zip(
output_chunk_bbs, output_chunk_bbs_with_overlap, block_indices
):
# calculate relevant slices for each output chunk
# this is specific to each non spatial coordinate
views_overlap_bb = [
mv_graph.get_overlap_for_bbs(
target_bb=output_chunk_bb_with_overlap,
query_bbs=[view_bb],
param=sparams[iview],
additional_extent_in_pixels={
dim: 0 if dim in fix_dims else int(interpolation_order)
for dim in sdims
},
)[0]
for iview, view_bb in enumerate(views_bb)
]
# append to output
relevant_view_indices = np.where(
[
view_overlap_bb is not None
for view_overlap_bb in views_overlap_bb
]
)[0]
if not len(relevant_view_indices):
fused_output_chunks[tuple(block_index)] = da.zeros(
tuple([output_chunk_bb["shape"][dim] for dim in sdims]),
dtype=sims[0].dtype,
)
continue
tol = 1e-6
sims_slices = [
sims[iview].sel(
sim_coord_dict
| {
dim: slice(
views_overlap_bb[iview]["origin"][dim] - tol,
views_overlap_bb[iview]["origin"][dim]
+ (views_overlap_bb[iview]["shape"][dim] - 1)
* views_overlap_bb[iview]["spacing"][dim]
+ tol,
)
for dim in sdims
},
drop=True,
)
for iview in relevant_view_indices
]
# determine whether to fuse plany by plane
# to avoid weighting edge artifacts
# fuse planewise if:
# - z dimension is present
# - params don't affect z dimension
# - shape in z dimension is 1 (i.e. only one plane)
# (the last criterium above could be dropped if we find a way
# (to propagate metadata through xr.apply_ufunc)
if (
"z" in fix_dims
and output_chunk_bb_with_overlap["shape"]["z"] == 1
):
fuse_planewise = True
sims_slices = [sim.isel(z=0) for sim in sims_slices]
tmp_params = [
sparams[iview].sel(
x_in=["y", "x", "1"],
x_out=["y", "x", "1"],
)
for iview in relevant_view_indices
]
output_chunk_bb_with_overlap = mv_graph.project_bb_along_dim(
output_chunk_bb_with_overlap, dim="z"
)
full_view_bbs = [
mv_graph.project_bb_along_dim(views_bb[iview], dim="z")
for iview in relevant_view_indices
]
else:
fuse_planewise = False
tmp_params = [
sparams[iview] for iview in relevant_view_indices
]
full_view_bbs = [
views_bb[iview] for iview in relevant_view_indices
]
fused_output_chunk = delayed(
lambda append_leading_axis, **kwargs: fuse_np(**kwargs)[
np.newaxis
]
if append_leading_axis
else fuse_np(**kwargs),
)(
append_leading_axis=fuse_planewise,
sims=sims_slices,
params=tmp_params,
output_properties=output_chunk_bb_with_overlap,
fusion_func=fusion_func,
fusion_method_kwargs=fusion_method_kwargs,
weights_func=weights_func,
weights_func_kwargs=weights_func_kwargs,
trim_overlap_in_pixels=overlap_in_pixels,
interpolation_order=1,
full_view_bbs=full_view_bbs,
blending_widths=blending_widths,
)
fused_output_chunk = da.from_delayed(
fused_output_chunk,
shape=tuple([output_chunk_bb["shape"][dim] for dim in sdims]),
dtype=sims[0].dtype,
)
fused_output_chunks[tuple(block_index)] = fused_output_chunk
fused = da.block(fused_output_chunks.tolist())
merge = si.to_spatial_image(
fused,
dims=sdims,
scale=output_stack_properties["spacing"],
translation=output_stack_properties["origin"],
)
merge = merge.expand_dims(nsdims)
merge = merge.assign_coords(
{ns_coord.name: [ns_coord.values] for ns_coord in ns_coords}
)
merges.append(merge)
if len(merges) > 1:
# suppress pandas future warning occuring within xarray.concat
with warnings.catch_warnings():
warnings.simplefilter(action="ignore", category=FutureWarning)
# if sims are named, combine_by_coord returns a dataset
res = xr.combine_by_coords([m.rename(None) for m in merges])
else:
res = merge
res = si_utils.get_sim_from_xim(res)
si_utils.set_sim_affine(
res,
param_utils.identity_transform(len(sdims)),
transform_key,
)
# order channels in the same way as first input sim
# (combine_by_coords may change coordinate order)
if "c" in res.dims:
res = res.sel({"c": sims[0].coords["c"].values})
return res