@@ -18,6 +18,10 @@ def mutex_watershed(
1818 affinities : np .ndarray ,
1919 offsets : Sequence [Sequence [int ]],
2020 number_of_attractive_channels : int ,
21+ * ,
22+ strides : Sequence [int ] | None = None ,
23+ randomized_strides : bool = False ,
24+ mask : np .ndarray | None = None ,
2125) -> np .ndarray :
2226 """Run mutex watershed on a 2D or 3D image-derived grid graph.
2327
@@ -33,6 +37,17 @@ def mutex_watershed(
3337 number_of_attractive_channels:
3438 The first this many affinity channels are attractive merge edges. The
3539 remaining channels are mutex edges.
40+ strides:
41+ Optional spatial sub-sampling strides for mutex edges. Attractive
42+ channels are always kept. If given, it must have one positive integer
43+ per spatial dimension.
44+ randomized_strides:
45+ If ``True``, sub-sample mutex edges randomly with probability
46+ ``1 / np.prod(strides)`` instead of on a regular grid.
47+ mask:
48+ Optional boolean foreground mask with shape ``affinities.shape[1:]``.
49+ Edges touching ``False`` pixels are ignored and ``False`` pixels are
50+ labeled as background ``0`` in the output.
3651
3752 Returns
3853 -------
@@ -73,5 +88,140 @@ def mutex_watershed(
7388 if number_of_attractive_channels > array .shape [0 ]:
7489 raise ValueError ("number_of_attractive_channels must be <= number of channels" )
7590
91+ normalized_strides = _normalize_strides (strides , spatial_ndim , randomized_strides )
92+ valid_edges = _compute_valid_edges (
93+ array .shape ,
94+ normalized_offsets ,
95+ number_of_attractive_channels ,
96+ normalized_strides ,
97+ randomized_strides ,
98+ mask ,
99+ )
100+
76101 contiguous = np .ascontiguousarray (array )
77- return run (contiguous , normalized_offsets , number_of_attractive_channels )
102+ labels = run (
103+ contiguous ,
104+ valid_edges ,
105+ normalized_offsets ,
106+ number_of_attractive_channels ,
107+ )
108+ if mask is not None :
109+ labels [np .logical_not (np .asarray (mask ))] = 0
110+ return labels
111+
112+
113+ def _normalize_strides (
114+ strides : Sequence [int ] | None ,
115+ spatial_ndim : int ,
116+ randomized_strides : bool ,
117+ ) -> tuple [int , ...] | None :
118+ if strides is None :
119+ if randomized_strides :
120+ raise ValueError ("randomized_strides requires strides" )
121+ return None
122+
123+ normalized = tuple (int (stride ) for stride in strides )
124+ if len (normalized ) != spatial_ndim :
125+ raise ValueError (
126+ "strides length must match the spatial ndim, got "
127+ f"strides length={ len (normalized )} { spatial_ndim }
128+ )
129+ if any (stride <= 0 for stride in normalized ):
130+ raise ValueError ("strides must contain only positive integers" )
131+ return normalized
132+
133+
134+ def _valid_source_slices (
135+ image_shape : tuple [int , ...],
136+ offset : tuple [int , ...],
137+ ) -> tuple [slice , ...] | None :
138+ slices = []
139+ for axis_size , step in zip (image_shape , offset , strict = True ):
140+ if step > 0 :
141+ if step >= axis_size :
142+ return None
143+ slices .append (slice (0 , axis_size - step ))
144+ elif step < 0 :
145+ if - step >= axis_size :
146+ return None
147+ slices .append (slice (- step , axis_size ))
148+ else :
149+ slices .append (slice (None ))
150+ return tuple (slices )
151+
152+
153+ def _neighbor_slices (
154+ image_shape : tuple [int , ...],
155+ offset : tuple [int , ...],
156+ ) -> tuple [slice , ...] | None :
157+ slices = []
158+ for axis_size , step in zip (image_shape , offset , strict = True ):
159+ if step > 0 :
160+ if step >= axis_size :
161+ return None
162+ slices .append (slice (step , axis_size ))
163+ elif step < 0 :
164+ if - step >= axis_size :
165+ return None
166+ slices .append (slice (0 , axis_size + step ))
167+ else :
168+ slices .append (slice (None ))
169+ return tuple (slices )
170+
171+
172+ def _compute_valid_edges (
173+ affinity_shape : tuple [int , ...],
174+ offsets : Sequence [tuple [int , ...]],
175+ number_of_attractive_channels : int ,
176+ strides : tuple [int , ...] | None ,
177+ randomized_strides : bool ,
178+ mask : np .ndarray | None ,
179+ ) -> np .ndarray :
180+ image_shape = tuple (int (size ) for size in affinity_shape [1 :])
181+ valid_edges = np .zeros (affinity_shape , dtype = bool )
182+
183+ for channel , offset in enumerate (offsets ):
184+ source_slices = _valid_source_slices (image_shape , offset )
185+ if source_slices is not None :
186+ valid_edges [(channel ,) + source_slices ] = True
187+
188+ if strides is not None :
189+ stride_edges = np .zeros_like (valid_edges , dtype = bool )
190+ stride_edges [:number_of_attractive_channels ] = True
191+ if randomized_strides :
192+ stride_factor = 1.0 / np .prod (strides )
193+ stride_edges [number_of_attractive_channels :] = (
194+ np .random .random (
195+ valid_edges [number_of_attractive_channels :].shape
196+ )
197+ < stride_factor
198+ )
199+ else :
200+ valid_slice = (slice (number_of_attractive_channels , None ),) + tuple (
201+ slice (None , None , stride ) for stride in strides
202+ )
203+ stride_edges [valid_slice ] = True
204+ valid_edges &= stride_edges
205+
206+ if mask is not None :
207+ mask_array = np .asarray (mask )
208+ if mask_array .shape != image_shape :
209+ raise ValueError (
210+ "mask shape must match affinities spatial shape, got "
211+ f"mask shape={ mask_array .shape } { image_shape }
212+ )
213+ if mask_array .dtype != np .dtype ("bool" ):
214+ raise TypeError (f"mask must have dtype bool, got dtype={ mask_array .dtype } )
215+
216+ invalid = np .logical_not (mask_array )
217+ for channel , offset in enumerate (offsets ):
218+ source_slices = _valid_source_slices (image_shape , offset )
219+ neighbor_slices = _neighbor_slices (image_shape , offset )
220+ if source_slices is None or neighbor_slices is None :
221+ continue
222+ touches_invalid = invalid [source_slices ] | invalid [neighbor_slices ]
223+ channel_valid = valid_edges [(channel ,) + source_slices ]
224+ channel_valid [touches_invalid ] = False
225+ valid_edges [(channel ,) + source_slices ] = channel_valid
226+
227+ return np .ascontiguousarray (valid_edges , dtype = np .uint8 )
0 commit comments