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,
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,
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,
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.
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.
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,
interpolation_order=interpolation_order,
blending_widths=blending_widths,
output_zarr_url=output_zarr_url,
zarr_options=zarr_options,
batch_options=batch_options,
)
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
),
)
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,
interpolation_order=interpolation_order,
blending_widths=blending_widths,
)
)
# 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,
"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,
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
)
# 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
]
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)
# Pre-compute the inverse of each tile's affine transform once.
# This avoids recomputing np.linalg.inv inside get_overlap_for_bbs for
# every (tile, chunk) pair.
inv_sparams = [
xr.DataArray(
np.linalg.inv(sp.data),
dims=sp.dims,
coords=sp.coords,
)
for sp in sparams
]
# Build a chunk_index -> [tile_indices] mapping by iterating over tiles
# and using the regular output chunk grid to find overlapping chunks via
# simple index arithmetic — O(N_tiles * ndim) instead of O(N_tiles * N_chunks).
#
# Output chunks form a regular grid: all blocks have the same pixel size
# per dimension except possibly the last block (which may be smaller).
# The first (uniform) block size is sufficient to map any physical
# coordinate to a chunk index via floor division; we clamp to the valid
# range so the last partial block is always included when needed.
_normalized_chunks = normalize_chunks(
[output_chunksize[dim] for dim in sdims],
[output_stack_properties["shape"][dim] for dim in sdims],
)
_n_blocks_per_dim = [len(c) for c in _normalized_chunks]
_uniform_cs_per_dim = [c[0] for c in _normalized_chunks]
_osp_origin = np.array(
[output_stack_properties["origin"][dim] for dim in sdims]
)
_osp_spacing = np.array(
[output_stack_properties["spacing"][dim] for dim in sdims]
)
# Physical padding: interpolation extent + chunk overlap added to output chunks
_additional_extent_pixels = np.array(
[0.0 if dim in fix_dims else float(interpolation_order) for dim in sdims]
)
_padding_phys = (
_additional_extent_pixels * _osp_spacing
+ np.array([overlap_in_pixels[dim] for dim in sdims]) * _osp_spacing
)
_chunk_to_tiles: dict = {}
for iview in range(len(sims)):
# Forward-project tile corners through the affine to get its AABB
# in output (world) space.
tile_corners_output = transformation.transform_pts(
mv_graph.get_vertices_from_stack_props(views_bb[iview]),
sparams[iview].data,
)
aabb_min = np.min(tile_corners_output, axis=0) - _padding_phys
aabb_max = np.max(tile_corners_output, axis=0) + _padding_phys
# Map the padded AABB to a range of chunk indices per dimension.
idx_ranges = []
skip = False
for idim in range(len(sdims)):
cs_phys = _uniform_cs_per_dim[idim] * _osp_spacing[idim]
i_first = max(
0,
int(np.floor((aabb_min[idim] - _osp_origin[idim]) / cs_phys)),
)
i_last = min(
_n_blocks_per_dim[idim] - 1,
int(np.floor((aabb_max[idim] - _osp_origin[idim]) / cs_phys)),
)
if i_first > i_last:
skip = True
break
idx_ranges.append(range(i_first, i_last + 1))
if skip:
continue
for chunk_idx in product(*idx_ranges):
_chunk_to_tiles.setdefault(chunk_idx, []).append(iview)
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
):
# Look up candidate tiles via the pre-built mapping (O(1) per chunk).
candidate_view_indices = _chunk_to_tiles.get(tuple(block_index), [])
views_overlap_bb = [None] * len(sims)
for iview in candidate_view_indices:
views_overlap_bb[iview] = mv_graph.get_overlap_for_bbs(
target_bb=output_chunk_bb_with_overlap,
query_bbs=[views_bb[iview]],
param=inv_sparams[iview],
additional_extent_in_pixels={
dim: 0 if dim in fix_dims else int(interpolation_order)
for dim in sdims
},
param_is_inverse=True,
)[0]
# 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_func_kwargs=fusion_func_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,
shrink_distance=shrink_distance,
)
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