def fuse(
images: list = None,
transform_key: str = None,
fusion_func: Callable = weighted_average_fusion,
fusion_func_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,
trim_overlap: bool = True,
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,
backend: str | None = None,
output_on_backend: bool = False,
sims: list | None = None,
):
"""
Fuse input images.
This function fuses all (Z)YX views ("fields") contained in the
input list of images, which can additionally contain C and T dimensions.
Inputs may be either all SpatialImages or all MultiscaleSpatialImages.
When fusing MultiscaleSpatialImages lazily, the returned object is also
multiscale and each output resolution is fused from a suitable input
resolution level.
Parameters
----------
images : list of SpatialImage or MultiscaleSpatialImage
Input images.
sims : list of SpatialImage or MultiscaleSpatialImage, optional
Deprecated alias for ``images``. Kept for backwards compatibility.
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_func_kwargs : dict, optional
weights_func : Callable, optional
Function to calculate fusion weights. This function always receives
transformed_views and may additionally receive params and
blending_weights when those parameters are declared in its
signature. 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 or a zarr-backed sim with stored chunk hints, its chunk grid
is used as the default. Otherwise, the default chunksize defined in
spatial_image_utils.py is used.
overlap_in_pixels : int or dict, optional
Pixel overlap added around each output chunk before fusion, by default
None.
trim_overlap : bool, optional
If True, trim fused chunks back to their requested output chunk size
after fusion. If False, keep the overlap in each fused chunk returned
by the lazy fusion graph. By default True.
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.
For MultiscaleSpatialImage inputs, the returned object remains multiscale:
OME-Zarr output returns the written MultiscaleSpatialImage, while plain Zarr output
returns a single-scale MultiscaleSpatialImage 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.
backend : str or None, optional
Compute backend to use for fusion. "numpy" (default) runs on the
CPU. "cupy" transfers each input chunk to the GPU via
cupy.asarray, runs resampling, blending-weight computation, and
the fusion function on the GPU, then returns a NumPy-backed result via
cupy.asnumpy. Requires CuPy to be installed; raises
ImportError if CuPy is not available. None is equivalent to
"numpy".
output_on_backend : bool, optional
If True, keep each fused chunk on the selected compute backend for
lazy, non-Zarr output. By default False, which converts CuPy-backed
results to NumPy before returning them.
Returns
-------
SpatialImage or MultiscaleSpatialImage
Fused image.
"""
if images is None:
if sims is None:
raise TypeError(
"fuse() missing 1 required positional argument: 'images'"
)
warnings.warn(
"The fuse(..., sims=...) parameter is deprecated and will be removed "
"in a future version. Use images=... instead.",
DeprecationWarning,
stacklevel=2,
)
images = sims
elif sims is not None:
raise TypeError(
"fuse() got both 'images' and deprecated 'sims'. Use only 'images'."
)
if not images:
raise ValueError("images must contain at least one image.")
input_is_msim = [msi_utils.is_msim(image) for image in images]
if any(input_is_msim) and not all(input_is_msim):
# Mixed inputs would make both scale selection and return type ambiguous.
raise ValueError(
"All input images must be of the same kind: either all SpatialImages "
"or all MultiscaleSpatialImages."
)
if all(input_is_msim):
msims = images
# Scale 0 defines the finest output geometry; coarser outputs derive from it.
scale0_sims = [
msi_utils.get_sim_from_msim(msim, scale="scale0") for msim in msims
]
scale0_output_stack_properties = process_output_stack_properties(
sims=scale0_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,
)
if output_zarr_url is not None:
# The Zarr path writes one sim, so choose the input level
# matching that output spacing.
selected_level_sims = [
msi_utils.get_sim_from_msim(
msim,
scale="scale%s"
% msi_utils.get_res_level_from_spacing(
msim, scale0_output_stack_properties["spacing"]
),
)
for msim in msims
]
fused = fuse(
images=selected_level_sims,
transform_key=transform_key,
fusion_func=fusion_func,
fusion_func_kwargs=fusion_func_kwargs,
weights_func=weights_func,
weights_func_kwargs=weights_func_kwargs,
output_stack_mode=output_stack_mode,
output_stack_properties=scale0_output_stack_properties,
output_chunksize=output_chunksize,
overlap_in_pixels=overlap_in_pixels,
trim_overlap=trim_overlap,
interpolation_order=interpolation_order,
blending_widths=blending_widths,
output_zarr_url=output_zarr_url,
zarr_options=zarr_options,
batch_options=batch_options,
backend=backend,
)
if (zarr_options or {}).get("ome_zarr", False):
return ngff_utils.read_msim_from_ome_zarr(
output_zarr_url,
transform_key=(
transform_key
if transform_key is not None
else si_utils.DEFAULT_TRANSFORM_KEY
),
# The fusion write path should keep returning a dask-backed
# msim even though the general OME-Zarr readers default to
# zarr-backed arrays.
array_backend="dask",
)
return msi_utils.get_msim_from_sim(fused, scale_factors=[])
res_shapes, _, res_abs_factors = msi_utils.calc_resolution_levels(
scale0_output_stack_properties["shape"],
)
fused_sims = []
for shape, abs_factors in zip(res_shapes, res_abs_factors):
# Match the center-of-pixel origin convention used by OME-Zarr
# output for downsampled levels.
curr_output_stack_properties = {
"shape": shape,
"spacing": {
dim: scale0_output_stack_properties["spacing"][dim]
* abs_factors[dim]
for dim in shape
},
"origin": {
dim: scale0_output_stack_properties["origin"][dim]
+ (abs_factors[dim] - 1)
* scale0_output_stack_properties["spacing"][dim]
/ 2
for dim in shape
},
}
# Fuse each output level from the coarsest input data that is
# still fine enough.
curr_sims = [
msi_utils.get_sim_from_msim(
msim,
scale="scale%s"
% msi_utils.get_res_level_from_spacing(
msim, curr_output_stack_properties["spacing"]
),
)
for msim in msims
]
fused_sims.append(
fuse(
images=curr_sims,
transform_key=transform_key,
fusion_func=fusion_func,
fusion_func_kwargs=fusion_func_kwargs,
weights_func=weights_func,
weights_func_kwargs=weights_func_kwargs,
output_stack_mode=output_stack_mode,
output_stack_properties=curr_output_stack_properties,
output_chunksize=output_chunksize,
overlap_in_pixels=overlap_in_pixels,
trim_overlap=trim_overlap,
interpolation_order=interpolation_order,
blending_widths=blending_widths,
backend=backend,
output_on_backend=output_on_backend,
)
)
# The levels have already been fused; assemble them without further
# downsampling.
return msi_utils.get_msim_from_sims(fused_sims)
sims = images
# 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)
# 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)
zarr_array_creation_kwargs = zarr_options.get(
"zarr_array_creation_kwargs", None
)
# 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(output_zarr_url):
shutil.rmtree(output_zarr_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 = {
"images": sims,
"transform_key": transform_key,
"fusion_func": fusion_func,
"fusion_func_kwargs": fusion_func_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,
"trim_overlap": trim_overlap,
"interpolation_order": interpolation_order,
"blending_widths": blending_widths,
"backend": backend,
# Direct Zarr writes require host arrays, so output_on_backend is
# intentionally not forwarded into these per-chunk fuse calls.
}
# 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,
zarr_array_creation_kwargs=zarr_array_creation_kwargs,
ngff_version=ngff_version,
)
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/fusion methods and user-supplied value
overlap_in_pixels = overlap_in_pixels or 0
# normalize to dict[str, int]
if not isinstance(overlap_in_pixels, dict):
overlap_in_pixels = {dim: overlap_in_pixels for dim in sdims}
shrink_distance = 0
for func, func_kwargs in [
(weights_func, weights_func_kwargs),
(fusion_func, fusion_func_kwargs),
]:
if func is not None and hasattr(func, "required_overlap"):
# Inject output_chunksize so the overlap can be clamped to it.
_kwargs_with_chunksize = dict(func_kwargs or {})
if (
has_keyword(func, "output_chunksize")
and output_chunksize is not None
):
_kwargs_with_chunksize.setdefault(
"output_chunksize", output_chunksize
)
curr_overlap = func.required_overlap(_kwargs_with_chunksize)
# normalize
if not isinstance(curr_overlap, dict):
curr_overlap = {dim: curr_overlap for dim in sdims}
overlap_in_pixels = {
dim: max(overlap_in_pixels[dim], curr_overlap[dim])
for dim in sdims
}
if func is not None and hasattr(func, "required_source_shrinkage"):
shrink_distance = func.required_source_shrinkage(func_kwargs)
# calculate output chunk bounding boxes
output_chunk_bbs, block_indices = mv_graph.get_chunk_bbs(
output_stack_properties, output_chunksize
)
# Chunk-grid shape reused to build the zarr-backed per-block param array.
nblocks_per_dim = tuple(int(x) for x in np.max(block_indices, 0) + 1)
# 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[dim]
* output_stack_properties["spacing"][dim]
for dim in sdims
}
}
| {
"shape": {
dim: output_chunk_bb["shape"][dim] + 2 * overlap_in_pixels[dim]
for dim in sdims
}
}
for output_chunk_bb in output_chunk_bbs
]
output_chunk_bbs_for_result = (
output_chunk_bbs if trim_overlap else output_chunk_bbs_with_overlap
)
views_bb = [si_utils.get_stack_properties_from_sim(sim) for sim in sims]
# All-zarr inputs use map_blocks (thin graph); dask inputs use per-chunk delayed.
input_is_zarr = all(si_utils.is_xarray_zarr_backed(sim) for sim in sims)
if input_is_zarr:
# Precompute lightweight tile representations once (outside the ns_coords loop).
# Spatial coordinate arrays are excluded — they are rebuilt cheaply from
# spacing + origin at compute time, avoiding large coord array serialisation.
# serialize_zarr_backed_sim also captures any dropped-dim selections so
# that sims pre-filtered with sim_sel_coords are handled correctly.
tile_series = [si_utils.serialize_zarr_backed_sim(sim) for sim in sims]
param_dependent_nsdims = [
dim for dim in nsdims if any(dim in param.dims for param in params)
]
spatial_plan_cache = {}
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 = {
dim: sim_coord_dict[dim] for dim in param_dependent_nsdims
}
spatial_plan_key = _get_spatial_plan_cache_key(
params_coord_dict,
param_dependent_nsdims,
)
if spatial_plan_key not in spatial_plan_cache:
spatial_plan_cache[spatial_plan_key] = _build_spatial_fusion_plan(
sparams=_select_params_for_coords(params, params_coord_dict),
views_bb=views_bb,
output_stack_properties=output_stack_properties,
output_chunksize=output_chunksize,
output_chunk_bbs=output_chunk_bbs,
output_chunk_bbs_with_overlap=output_chunk_bbs_with_overlap,
output_chunk_bbs_for_result=output_chunk_bbs_for_result,
block_indices=block_indices,
overlap_in_pixels=overlap_in_pixels,
trim_overlap=trim_overlap,
interpolation_order=interpolation_order,
sdims=sdims,
)
spatial_plan = spatial_plan_cache[spatial_plan_key]
sparams = spatial_plan["sparams"]
per_chunk_entries = spatial_plan["per_chunk_entries"]
if input_is_zarr:
if "zarr_chunk_params" not in spatial_plan:
spatial_plan["zarr_chunk_params"] = _build_zarr_chunk_params(
plan=spatial_plan,
tile_series=tile_series,
views_bb=views_bb,
block_indices=block_indices,
nblocks_per_dim=nblocks_per_dim,
)
chunk_params = spatial_plan["zarr_chunk_params"]
# === zarr path: thin dask graph via map_blocks ===
fused = da.map_blocks(
_fuse_block_zarr_backed,
chunk_params,
chunks=_chunks_from_chunk_bbs(
output_chunk_bbs_for_result, block_indices, sdims
),
dtype=sims[0].dtype,
output_dtype=sims[0].dtype,
sim_coord_dict=sim_coord_dict,
sdims=sdims,
fusion_func=fusion_func,
fusion_func_kwargs=fusion_func_kwargs,
weights_func=weights_func,
weights_func_kwargs=weights_func_kwargs,
overlap_in_pixels=overlap_in_pixels,
trim_overlap=trim_overlap,
interpolation_order=interpolation_order,
blending_widths=blending_widths,
shrink_distance=shrink_distance,
backend=backend,
output_on_backend=output_on_backend,
)
else:
# === dask path: per-chunk delayed tasks ===
fused_output_chunks = np.empty(
np.max(block_indices, 0) + 1, dtype=object
)
for block_index, entry in zip(block_indices, per_chunk_entries):
output_chunk_bb_with_overlap = entry["output_bb_overlap"]
output_chunk_bb_result = entry["output_bb_result"]
fuse_planewise = entry["fuse_planewise"]
chunk_views = entry["views"]
relevant_view_indices = [iview for iview, _ in chunk_views]
if not len(relevant_view_indices):
fused_output_chunks[tuple(block_index)] = da.zeros(
tuple(
[
output_chunk_bb_result["shape"][dim]
for dim in sdims
]
),
dtype=sims[0].dtype,
)
continue
tol = 1e-6
sims_slices = [
sims[iview].sel(
sim_coord_dict
| {
dim: slice(
tile_overlap_bb["origin"][dim] - tol,
tile_overlap_bb["origin"][dim]
+ (tile_overlap_bb["shape"][dim] - 1)
* tile_overlap_bb["spacing"][dim]
+ tol,
)
for dim in sdims
},
drop=True,
)
for iview, tile_overlap_bb in chunk_views
]
# determine whether to fuse plane 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 fuse_planewise:
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:
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_func_kwargs=fusion_func_kwargs,
weights_func=weights_func,
weights_func_kwargs=weights_func_kwargs,
# The overlap still defines the fusion target; this only
# controls the final crop inside fuse_np.
trim_overlap_in_pixels=(
overlap_in_pixels if trim_overlap else 0
),
interpolation_order=interpolation_order,
full_view_bbs=full_view_bbs,
blending_widths=blending_widths,
shrink_distance=shrink_distance,
backend=backend,
output_on_backend=output_on_backend,
)
fused_output_chunk = da.from_delayed(
fused_output_chunk,
shape=tuple(
[output_chunk_bb_result["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_utils.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