topo.stitch

Stitching utilities for tiled inference.

Split a volume into overlapping sub-crops, process each independently, then blend results back into the full volume using boundary-aware cosine weighting. The key insight is that cosine tapering should only be applied at internal subcrop seams (where two tiles overlap), not at volume boundaries where there is no neighbor to blend with.

Functions

compute_subcrop_slices(→ list[tuple[slice, Ellipsis]])

Compute sub-crop slices with overlap for each axis.

build_spatial_mask(→ numpy.ndarray)

Spatial mask: 1 in interior, 0 at boundary faces.

cosine_blend_weight(→ numpy.ndarray)

Cosine-ramp weight volume for blending overlapping sub-crops.

stitch_flows(→ numpy.ndarray)

Stitch per-subcrop flows into a full volume using cosine blending.

stitch_labels(→ numpy.ndarray)

Stitch per-subcrop label volumes using highest-weight-wins in overlaps.

stitch_volumes(→ numpy.ndarray)

Stitch per-subcrop multi-channel volumes (e.g. affinities) using cosine blending.

Module Contents

topo.stitch.compute_subcrop_slices(volume_shape: tuple[int, Ellipsis], crop_size: tuple[int, Ellipsis], overlap: tuple[int, Ellipsis]) → list[tuple[slice, Ellipsis]]

Compute sub-crop slices with overlap for each axis.

Parameters:

  • volume_shape – (D, H, W) of the full volume.

  • crop_size – (cd, ch, cw) size of each sub-crop.

  • overlap – (od, oh, ow) overlap between adjacent sub-crops.

Returns:

List of tuples of slices [(slice_d, slice_h, slice_w), …].

topo.stitch.build_spatial_mask(shape: tuple[int, Ellipsis]) → numpy.ndarray

Spatial mask: 1 in interior, 0 at boundary faces.

Used to mark subcrop boundaries so diffusion uses Neumann BC there (allowing flow to continue past the crop edge).

Note: Do NOT use this for the full volume if the volume boundary is real (not a crop edge). Pass spatial_mask=None to flow generation instead, so that Dirichlet BC keeps flow pointing inward.

topo.stitch.cosine_blend_weight(shape: tuple[int, Ellipsis], at_volume_boundary: list[bool] | None = None) → numpy.ndarray

Cosine-ramp weight volume for blending overlapping sub-crops.

At internal subcrop seams the weight tapers from 1 (center) to 0 (edge), enabling smooth blending with the neighboring tile. At volume boundaries (no neighbor), the weight stays at 1 so the flow is fully preserved.

Parameters:

  • shape – (d, h, w) of the subcrop.

  • at_volume_boundary – List of 6 bools — [z_lo, z_hi, y_lo, y_hi, x_lo, x_hi]. True means that face sits at the volume boundary (no overlap). If None, all faces are tapered (original behavior).

Returns:

[d, h, w] float32 weight volume.

topo.stitch.stitch_flows(volume_shape: tuple[int, int, int], subcrop_slices: list[tuple[slice, Ellipsis]], subcrop_flows: list[numpy.ndarray]) → numpy.ndarray

Stitch per-subcrop flows into a full volume using cosine blending.

Parameters:

  • volume_shape – (D, H, W) of the full volume.

  • subcrop_slices – List of (slice_d, slice_h, slice_w) tuples.

  • subcrop_flows – List of [3, cd, ch, cw] flow arrays.

Returns:

[3, D, H, W] float32 blended flow field.

Return type:

stitched

topo.stitch.stitch_labels(volume_shape: tuple[int, int, int], subcrop_slices: list[tuple[slice, Ellipsis]], subcrop_labels: list[numpy.ndarray]) → numpy.ndarray

Stitch per-subcrop label volumes using highest-weight-wins in overlaps.

Each voxel gets the label from whichever subcrop has the highest cosine blend weight at that position (i.e. the subcrop whose center is closest). Labels are relabeled to avoid collisions across sub-crops.

Parameters:

  • volume_shape – (D, H, W) of the full volume.

  • subcrop_slices – List of (slice_d, slice_h, slice_w) tuples.

  • subcrop_labels – List of [cd, ch, cw] int label arrays.

Returns:

[D, H, W] int32 stitched labels.

topo.stitch.stitch_volumes(volume_shape: tuple[int, int, int], subcrop_slices: list[tuple[slice, Ellipsis]], subcrop_volumes: list[numpy.ndarray], n_channels: int = 1) → numpy.ndarray

Stitch per-subcrop multi-channel volumes (e.g. affinities) using cosine blending.

Parameters:

  • volume_shape – (D, H, W) of the full volume.

  • subcrop_slices – List of (slice_d, slice_h, slice_w) tuples.

  • subcrop_volumes – List of [C, cd, ch, cw] arrays (or [cd, ch, cw] if n_channels=1).

  • n_channels – Number of channels.

Returns:

[C, D, H, W] float64 blended volume (or [D, H, W] if n_channels=1).