torch_em.transform.label

  1from typing import Optional
  2
  3import numpy as np
  4import skimage.measure
  5import skimage.segmentation
  6import vigra
  7
  8from ..util import ensure_array, ensure_spatial_array
  9
 10try:
 11    from affogato.affinities import compute_affinities
 12except ImportError:
 13    compute_affinities = None
 14
 15
 16def connected_components(labels, ndim=None, ensure_zero=False):
 17    labels = ensure_array(labels) if ndim is None else ensure_spatial_array(labels, ndim)
 18    labels = skimage.measure.label(labels)
 19    if ensure_zero and 0 not in labels:
 20        labels -= 1
 21    return labels
 22
 23
 24def labels_to_binary(labels, background_label=0):
 25    return (labels != background_label).astype(labels.dtype)
 26
 27
 28def label_consecutive(labels, with_background=True):
 29    if with_background:
 30        seg = skimage.segmentation.relabel_sequential(labels)[0]
 31    else:
 32        if 0 in labels:
 33            labels += 1
 34        seg = skimage.segmentation.relabel_sequential(labels)[0]
 35        assert seg.min() == 1
 36        seg -= 1
 37    return seg
 38
 39
 40# TODO smoothing
 41class BoundaryTransform:
 42    def __init__(self, mode="thick", add_binary_target=False, ndim=None):
 43        self.mode = mode
 44        self.add_binary_target = add_binary_target
 45        self.ndim = ndim
 46
 47    def __call__(self, labels):
 48        labels = ensure_array(labels) if self.ndim is None else ensure_spatial_array(labels, self.ndim)
 49        boundaries = skimage.segmentation.find_boundaries(labels, mode=self.mode)[None]
 50        if self.add_binary_target:
 51            binary = labels_to_binary(labels)[None].astype(boundaries.dtype)
 52            target = np.concatenate([binary, boundaries], axis=0)
 53        else:
 54            target = boundaries
 55        return target
 56
 57
 58# TODO smoothing
 59class NoToBackgroundBoundaryTransform:
 60    def __init__(self, bg_label=0, mask_label=-1, mode="thick", add_binary_target=False, ndim=None):
 61        self.bg_label = bg_label
 62        self.mask_label = mask_label
 63        self.mode = mode
 64        self.ndim = ndim
 65        self.add_binary_target = add_binary_target
 66
 67    def __call__(self, labels):
 68        labels = ensure_array(labels) if self.ndim is None else ensure_spatial_array(labels, self.ndim)
 69        # calc normal boundaries
 70        boundaries = skimage.segmentation.find_boundaries(labels, mode=self.mode)[None]
 71
 72        # make label image binary and calculate to-background-boundaries
 73        labels_binary = (labels != self.bg_label)
 74        to_bg_boundaries = skimage.segmentation.find_boundaries(labels_binary, mode=self.mode)[None]
 75
 76        # mask the to-background-boundaries
 77        boundaries = boundaries.astype(np.int8)
 78        boundaries[to_bg_boundaries] = self.mask_label
 79
 80        if self.add_binary_target:
 81            binary = labels_to_binary(labels, self.bg_label).astype(boundaries.dtype)
 82            binary[labels == self.mask_label] = self.mask_label
 83            target = np.concatenate([binary[None], boundaries], axis=0)
 84        else:
 85            target = boundaries
 86
 87        return target
 88
 89
 90# TODO smoothing
 91class BoundaryTransformWithIgnoreLabel:
 92    def __init__(self, ignore_label=-1, mode="thick", add_binary_target=False, ndim=None):
 93        self.ignore_label = ignore_label
 94        self.mode = mode
 95        self.ndim = ndim
 96        self.add_binary_target = add_binary_target
 97
 98    def __call__(self, labels):
 99        labels = ensure_array(labels) if self.ndim is None else ensure_spatial_array(labels, self.ndim)
100        # calculate the normal boundaries
101        boundaries = skimage.segmentation.find_boundaries(labels, mode=self.mode)[None]
102
103        # calculate the boundaries for the ignore label
104        labels_ignore = (labels == self.ignore_label)
105        to_ignore_boundaries = skimage.segmentation.find_boundaries(labels_ignore, mode=self.mode)[None]
106
107        # mask the to-background-boundaries
108        boundaries = boundaries.astype(np.int8)
109        boundaries[to_ignore_boundaries] = self.ignore_label
110
111        if self.add_binary_target:
112            binary = labels_to_binary(labels).astype(boundaries.dtype)
113            binary[labels == self.ignore_label] = self.ignore_label
114            target = np.concatenate([binary[None], boundaries], axis=0)
115        else:
116            target = boundaries
117
118        return target
119
120
121# TODO affinity smoothing
122class AffinityTransform:
123    def __init__(self, offsets,
124                 ignore_label=None,
125                 add_binary_target=False,
126                 add_mask=False,
127                 include_ignore_transitions=False):
128        assert compute_affinities is not None
129        self.offsets = offsets
130        self.ndim = len(self.offsets[0])
131        assert self.ndim in (2, 3)
132
133        self.ignore_label = ignore_label
134        self.add_binary_target = add_binary_target
135        self.add_mask = add_mask
136        self.include_ignore_transitions = include_ignore_transitions
137
138    def add_ignore_transitions(self, affs, mask, labels):
139        ignore_seg = (labels == self.ignore_label).astype(labels.dtype)
140        ignore_transitions, invalid_mask = compute_affinities(ignore_seg, self.offsets)
141        invalid_mask = np.logical_not(invalid_mask)
142        # NOTE affinity convention returned by affogato: transitions are marked by 0
143        ignore_transitions = ignore_transitions == 0
144        ignore_transitions[invalid_mask] = 0
145        affs[ignore_transitions] = 1
146        mask[ignore_transitions] = 1
147        return affs, mask
148
149    def __call__(self, labels):
150        dtype = "uint64"
151        if np.dtype(labels.dtype) in (np.dtype("int16"), np.dtype("int32"), np.dtype("int64")):
152            dtype = "int64"
153        labels = ensure_spatial_array(labels, self.ndim, dtype=dtype)
154        affs, mask = compute_affinities(labels, self.offsets,
155                                        have_ignore_label=self.ignore_label is not None,
156                                        ignore_label=0 if self.ignore_label is None else self.ignore_label)
157        # we use the "disaffinity" convention for training; i.e. 1 means repulsive, 0 attractive
158        affs = 1. - affs
159
160        # remove transitions to the ignore label from the mask
161        if self.ignore_label is not None and self.include_ignore_transitions:
162            affs, mask = self.add_ignore_transitions(affs, mask, labels)
163
164        if self.add_binary_target:
165            binary = labels_to_binary(labels)[None].astype(affs.dtype)
166            assert binary.ndim == affs.ndim
167            affs = np.concatenate([binary, affs], axis=0)
168
169        if self.add_mask:
170            if self.add_binary_target:
171                if self.ignore_label is None:
172                    mask_for_bin = np.ones((1,) + labels.shape, dtype=mask.dtype)
173                else:
174                    mask_for_bin = (labels != self.ignore_label)[None].astype(mask.dtype)
175                assert mask.ndim == mask_for_bin.ndim
176                mask = np.concatenate([mask_for_bin, mask], axis=0)
177            assert affs.shape == mask.shape
178            affs = np.concatenate([affs, mask.astype(affs.dtype)], axis=0)
179
180        return affs
181
182
183class OneHotTransform:
184    def __init__(self, class_ids=None):
185        self.class_ids = list(range(class_ids)) if isinstance(class_ids, int) else class_ids
186
187    def __call__(self, labels):
188        class_ids = np.unique(labels).tolist() if self.class_ids is None else self.class_ids
189        n_classes = len(class_ids)
190        one_hot = np.zeros((n_classes,) + labels.shape, dtype="float32")
191        for i, class_id in enumerate(class_ids):
192            one_hot[i][labels == class_id] = 1.0
193        return one_hot
194
195
196class DistanceTransform:
197    """Compute distances to foreground in the labels.
198
199    Args:
200        distances: Whether to compute the absolute distances.
201        directed_distances: Whether to compute the directed distances (vector distances).
202        normalize: Whether to normalize the computed distances.
203        max_distance: Maximal distance at which to threshold the distances.
204        foreground_id: Label id to which the distance is compute.
205        invert Whether to invert the distances:
206        func: Normalization function for the distances.
207    """
208    eps = 1e-7
209
210    def __init__(
211        self,
212        distances: bool = True,
213        directed_distances: bool = False,
214        normalize: bool = True,
215        max_distance: Optional[float] = None,
216        foreground_id=1,
217        invert=False,
218        func=None
219    ):
220        if sum((distances, directed_distances)) == 0:
221            raise ValueError("At least one of 'distances' or 'directed_distances' must be set to 'True'")
222        self.directed_distances = directed_distances
223        self.distances = distances
224        self.normalize = normalize
225        self.max_distance = max_distance
226        self.foreground_id = foreground_id
227        self.invert = invert
228        self.func = func
229
230    def _compute_distances(self, directed_distances):
231        distances = np.linalg.norm(directed_distances, axis=0)
232        if self.max_distance is not None:
233            distances = np.clip(distances, 0, self.max_distance)
234        if self.normalize:
235            distances /= (distances.max() + self.eps)
236        if self.invert:
237            distances = distances.max() - distances
238        if self.func is not None:
239            distances = self.func(distances)
240        return distances
241
242    def _compute_directed_distances(self, directed_distances):
243        if self.max_distance is not None:
244            directed_distances = np.clip(directed_distances, -self.max_distance, self.max_distance)
245        if self.normalize:
246            directed_distances /= (np.abs(directed_distances).max(axis=(1, 2), keepdims=True) + self.eps)
247        if self.invert:
248            directed_distances = directed_distances.max(axis=(1, 2), keepdims=True) - directed_distances
249        if self.func is not None:
250            directed_distances = self.func(directed_distances)
251        return directed_distances
252
253    def _get_distances_for_empty_labels(self, labels):
254        shape = labels.shape
255        fill_value = 0.0 if self.invert else np.sqrt(np.linalg.norm(list(shape)) ** 2 / 2)
256        data = np.full((labels.ndim,) + shape, fill_value)
257        return data
258
259    def __call__(self, labels):
260        distance_mask = (labels == self.foreground_id).astype("uint32")
261        # the distances are not computed corrected if they are all zero
262        # so this case needs to be handled separately
263        if distance_mask.sum() == 0:
264            directed_distances = self._get_distances_for_empty_labels(labels)
265        else:
266            ndim = distance_mask.ndim
267            to_channel_first = (ndim,) + tuple(range(ndim))
268            directed_distances = vigra.filters.vectorDistanceTransform(distance_mask).transpose(to_channel_first)
269
270        if self.distances:
271            distances = self._compute_distances(directed_distances)
272
273        if self.directed_distances:
274            directed_distances = self._compute_directed_distances(directed_distances)
275
276        if self.distances and self.directed_distances:
277            return np.concatenate((distances[None], directed_distances), axis=0)
278        if self.distances:
279            return distances
280        if self.directed_distances:
281            return directed_distances
282
283
284class PerObjectDistanceTransform:
285    """Compute normalized distances per object in a segmentation.
286
287    Args:
288        distances: Whether to compute the undirected distances.
289        boundary_distances: Whether to compute the distances to the object boundaries.
290        directed_distances: Whether to compute the directed distances (vector distances).
291        foreground: Whether to return a foreground channel.
292        apply_label: Whether to apply connected components to the labels before computing distances.
293        correct_centers: Whether to correct centers that are not in the objects.
294        min_size: Minimal size of objects for distance calculdation.
295        distance_fill_value: Fill value for the distances outside of objects.
296    """
297    eps = 1e-7
298
299    def __init__(
300        self,
301        distances=True,
302        boundary_distances=True,
303        directed_distances=False,
304        foreground=True,
305        instances=False,
306        apply_label=True,
307        correct_centers=True,
308        min_size=0,
309        distance_fill_value=1.0,
310    ):
311        if sum([distances, directed_distances, boundary_distances]) == 0:
312            raise ValueError("At least one of distances or directed distances has to be passed.")
313        self.distances = distances
314        self.boundary_distances = boundary_distances
315        self.directed_distances = directed_distances
316        self.foreground = foreground
317        self.instances = instances
318
319        self.apply_label = apply_label
320        self.correct_centers = correct_centers
321        self.min_size = min_size
322        self.distance_fill_value = distance_fill_value
323
324    def compute_normalized_object_distances(self, mask, boundaries, bb, center, distances):
325        # Crop the mask and generate array with the correct center.
326        cropped_mask = mask[bb]
327        cropped_center = tuple(ce - b.start for ce, b in zip(center, bb))
328
329        # The centroid might not be inside of the object.
330        # In this case we correct the center by taking the maximum of the distance to the boundary.
331        # Note: the centroid is still the best estimate for the center, as long as it's in the object.
332        correct_center = not cropped_mask[cropped_center]
333
334        # Compute the boundary distances if necessary.
335        # (Either if we need to correct the center, or compute the boundary distances anyways.)
336        if correct_center or self.boundary_distances:
337            # Crop the boundary mask and compute the boundary distances.
338            cropped_boundary_mask = boundaries[bb]
339            boundary_distances = vigra.filters.distanceTransform(cropped_boundary_mask)
340            boundary_distances[~cropped_mask] = 0
341            max_dist_point = np.unravel_index(np.argmax(boundary_distances), boundary_distances.shape)
342
343        # Set the crop center to the max dist point
344        if correct_center:
345            # Find the center (= maximal distance from the boundaries).
346            cropped_center = max_dist_point
347
348        cropped_center_mask = np.zeros_like(cropped_mask, dtype="uint32")
349        cropped_center_mask[cropped_center] = 1
350
351        # Compute the directed distances,
352        if self.distances or self.directed_distances:
353            this_distances = vigra.filters.vectorDistanceTransform(cropped_center_mask)
354        else:
355            this_distances = None
356
357        # Keep only the specified distances:
358        if self.distances and self.directed_distances:  # all distances
359            # Compute the undirected ditacnes from directed distances and concatenate,
360            undir = np.linalg.norm(this_distances, axis=-1, keepdims=True)
361            this_distances = np.concatenate([undir, this_distances], axis=-1)
362
363        elif self.distances:  # only undirected distances
364            # Compute the undirected distances from directed distances and keep only them.
365            this_distances = np.linalg.norm(this_distances, axis=-1, keepdims=True)
366
367        elif self.directed_distances:  # only directed distances
368            pass  # We don't have to do anything becasue the directed distances are already computed.
369
370        # Add an extra channel for the boundary distances if specified.
371        if self.boundary_distances:
372            boundary_distances = (boundary_distances[max_dist_point] - boundary_distances)[..., None]
373            if this_distances is None:
374                this_distances = boundary_distances
375            else:
376                this_distances = np.concatenate([this_distances, boundary_distances], axis=-1)
377
378        # Set distances outside of the mask to zero.
379        this_distances[~cropped_mask] = 0
380
381        # Normalize the distances.
382        spatial_axes = tuple(range(mask.ndim))
383        this_distances /= (np.abs(this_distances).max(axis=spatial_axes, keepdims=True) + self.eps)
384
385        # Set the distance values in the global result.
386        distances[bb][cropped_mask] = this_distances[cropped_mask]
387
388        return distances
389
390    def __call__(self, labels):
391        # Apply label (connected components) if specified.
392        if self.apply_label:
393            labels = skimage.measure.label(labels).astype("uint32")
394        else:  # Otherwise just relabel the segmentation.
395            labels = vigra.analysis.relabelConsecutive(labels)[0].astype("uint32")
396
397        # Filter out small objects if min_size is specified.
398        if self.min_size > 0:
399            ids, sizes = np.unique(labels, return_counts=True)
400            discard_ids = ids[sizes < self.min_size]
401            labels[np.isin(labels, discard_ids)] = 0
402            labels = vigra.analysis.relabelConsecutive(labels)[0].astype("uint32")
403
404        # Compute the boundaries. They will be used to determine the most central point,
405        # and if 'self.boundary_distances is True' to add the boundary distances.
406        boundaries = skimage.segmentation.find_boundaries(labels, mode="inner").astype("uint32")
407
408        # Compute region properties to derive bounding boxes and centers.
409        ndim = labels.ndim
410        props = skimage.measure.regionprops(labels)
411        bounding_boxes = {
412            prop.label: tuple(slice(prop.bbox[i], prop.bbox[i + ndim]) for i in range(ndim))
413            for prop in props
414        }
415
416        # Compute the object centers from centroids.
417        centers = {prop.label: np.round(prop.centroid).astype("int") for prop in props}
418
419        # Compute how many distance channels we have.
420        n_channels = 0
421        if self.distances:  # We need one channel for the overall distances.
422            n_channels += 1
423        if self.boundary_distances:  # We need one channel for the boundary distances.
424            n_channels += 1
425        if self.directed_distances:  # And ndim channels for directed distances.
426            n_channels += ndim
427
428        # Compute the per object distances.
429        distances = np.full(labels.shape + (n_channels,), self.distance_fill_value, dtype="float32")
430        for prop in props:
431            label_id = prop.label
432            mask = labels == label_id
433            distances = self.compute_normalized_object_distances(
434                mask, boundaries, bounding_boxes[label_id], centers[label_id], distances
435            )
436
437        # Bring the distance channel to the first dimension.
438        to_channel_first = (ndim,) + tuple(range(ndim))
439        distances = distances.transpose(to_channel_first)
440
441        # Add the foreground mask as first channel if specified.
442        if self.foreground:
443            binary_labels = (labels > 0).astype("float32")
444            distances = np.concatenate([binary_labels[None], distances], axis=0)
445
446        if self.instances:
447            distances = np.concatenate([labels[None], distances], axis=0)
448
449        return distances