torch_em.transform.label

  1from typing import Callable, List, Optional, Sequence, Union
  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: np.ndarray, ndim: Optional[int] = None, ensure_zero: bool = False) -> np.ndarray:
 17    """Apply connected components to a segmentation.
 18
 19    Args:
 20        labels: The input segmentation.
 21        ndim: The expected dimensionality of the data.
 22        ensure_zero: Whether to ensure that the data has a zero label.
 23
 24    Returns:
 25        The segmentation after connected components.
 26    """
 27    labels = ensure_array(labels) if ndim is None else ensure_spatial_array(labels, ndim)
 28    labels = skimage.measure.label(labels)
 29    if ensure_zero and 0 not in labels:
 30        labels -= 1
 31    return labels
 32
 33
 34def labels_to_binary(labels: np.ndarray, background_label: int = 0) -> np.ndarray:
 35    """Transform a segmentation to binary labels.
 36
 37    Args:
 38        labels: The input segmentation.
 39        background_label: The id of the background label.
 40
 41    Returns:
 42        The binary segmentation.
 43    """
 44    return (labels != background_label).astype(labels.dtype)
 45
 46
 47def label_consecutive(labels: np.ndarray, with_background: bool = True) -> np.ndarray:
 48    """Ensure that the input segmentation is labeled consecutively.
 49
 50    Args:
 51        labels: The input segmentation.
 52        with_background: Whether this segmentation has a background label.
 53
 54    Returns:
 55        The consecutively labeled segmentation.
 56    """
 57    if with_background:
 58        seg = skimage.segmentation.relabel_sequential(labels)[0]
 59    else:
 60        if 0 in labels:
 61            labels += 1
 62        seg = skimage.segmentation.relabel_sequential(labels)[0]
 63        assert seg.min() == 1
 64        seg -= 1
 65    return seg
 66
 67
 68class MinSizeLabelTransform:
 69    """Transformation to filter out objects smaller than a minimal size from the segmentation.
 70
 71    Args:
 72        min_size: The minimal object size of the segmentation.
 73        ndim: The dimensionality of the segmentation.
 74        ensure_zero: Ensure that the segmentation contains the id zero.
 75    """
 76    def __init__(self, min_size: Optional[int] = None, ndim: Optional[int] = None, ensure_zero: bool = False):
 77        self.min_size = min_size
 78        self.ndim = ndim
 79        self.ensure_zero = ensure_zero
 80
 81    def __call__(self, labels: np.ndarray) -> np.ndarray:
 82        """Filter out small objects from segmentation.
 83
 84        Args:
 85            labels: The input segmentation.
 86
 87        Returns:
 88            The size filtered segmentation.
 89        """
 90        components = connected_components(labels, ndim=self.ndim, ensure_zero=self.ensure_zero)
 91        if self.min_size is not None:
 92            ids, sizes = np.unique(components, return_counts=True)
 93            filter_ids = ids[sizes < self.min_size]
 94            components[np.isin(components, filter_ids)] = 0
 95            components, _, _ = skimage.segmentation.relabel_sequential(components)
 96        return components
 97
 98
 99# TODO smoothing
100class BoundaryTransform:
101    """Transformation to convert an instance segmentation into boundaries.
102
103    Args:
104        mode: The mode for converting the segmentation to boundaries.
105        add_binary_target: Whether to add a binary mask channel to the transformation output.
106        ndim: The expected dimensionality of the data.
107    """
108    def __init__(self, mode: str = "thick", add_binary_target: bool = False, ndim: Optional[int] = None):
109        self.mode = mode
110        self.add_binary_target = add_binary_target
111        self.ndim = ndim
112
113    def __call__(self, labels: np.ndarray) -> np.ndarray:
114        """Apply the boundary transformation to an input segmentation.
115
116        Args:
117            labels: The input segmentation.
118
119        Returns:
120            The boundaries.
121        """
122        labels = ensure_array(labels) if self.ndim is None else ensure_spatial_array(labels, self.ndim)
123        boundaries = skimage.segmentation.find_boundaries(labels, mode=self.mode)[None]
124        if self.add_binary_target:
125            binary = labels_to_binary(labels)[None].astype(boundaries.dtype)
126            target = np.concatenate([binary, boundaries], axis=0)
127        else:
128            target = boundaries
129        return target
130
131
132# TODO smoothing
133class NoToBackgroundBoundaryTransform:
134    """Transformation to convert an instance segmentation into boundaries.
135
136    This transformation sets boundaries with the ignore label to the ignore label
137    in the output of the transformation.
138
139    Args:
140        bg_label: The background label.
141        mask_label: The mask label.
142        mode: The mode for converting the segmentation to boundaries.
143        add_binary_target: Whether to add a binary mask channel to the transformation output.
144        ndim: The expected dimensionality of the data.
145    """
146    def __init__(
147        self,
148        bg_label: int = 0,
149        mask_label: int = -1,
150        mode: str = "thick",
151        add_binary_target: bool = False,
152        ndim: Optional[int] = None,
153    ):
154        self.bg_label = bg_label
155        self.mask_label = mask_label
156        self.mode = mode
157        self.ndim = ndim
158        self.add_binary_target = add_binary_target
159
160    def __call__(self, labels: np.ndarray) -> np.ndarray:
161        """Apply the boundary transformation to an input segmentation.
162
163        Args:
164            labels: The input segmentation.
165
166        Returns:
167            The boundaries.
168        """
169        labels = ensure_array(labels) if self.ndim is None else ensure_spatial_array(labels, self.ndim)
170        # calc normal boundaries
171        boundaries = skimage.segmentation.find_boundaries(labels, mode=self.mode)[None]
172
173        # make label image binary and calculate to-background-boundaries
174        labels_binary = (labels != self.bg_label)
175        to_bg_boundaries = skimage.segmentation.find_boundaries(labels_binary, mode=self.mode)[None]
176
177        # mask the to-background-boundaries
178        boundaries = boundaries.astype(np.int8)
179        boundaries[to_bg_boundaries] = self.mask_label
180
181        if self.add_binary_target:
182            binary = labels_to_binary(labels, self.bg_label).astype(boundaries.dtype)
183            binary[labels == self.mask_label] = self.mask_label
184            target = np.concatenate([binary[None], boundaries], axis=0)
185        else:
186            target = boundaries
187
188        return target
189
190
191# TODO smoothing
192class BoundaryTransformWithIgnoreLabel:
193    """Transformation to convert an instance segmentation into boundaries.
194
195    This transformation sets boundaries with the ignore label to the ignore label
196    in the output of the transformation.
197
198    Args:
199        ignore_label: The ignore label.
200        mode: The mode for converting the segmentation to boundaries.
201        add_binary_target: Whether to add a binary mask channel to the transformation output.
202        ndim: The expected dimensionality of the data.
203    """
204    def __init__(
205        self,
206        ignore_label: int = -1,
207        mode: str = "thick",
208        add_binary_target: bool = False,
209        ndim: Optional[int] = None,
210    ):
211        self.ignore_label = ignore_label
212        self.mode = mode
213        self.ndim = ndim
214        self.add_binary_target = add_binary_target
215
216    def __call__(self, labels: np.ndarray) -> np.ndarray:
217        """Apply the boundary transformation to an input segmentation.
218
219        Args:
220            labels: The input segmentation.
221
222        Returns:
223            The boundaries.
224        """
225        labels = ensure_array(labels) if self.ndim is None else ensure_spatial_array(labels, self.ndim)
226        # calculate the normal boundaries
227        boundaries = skimage.segmentation.find_boundaries(labels, mode=self.mode)[None]
228
229        # calculate the boundaries for the ignore label
230        labels_ignore = (labels == self.ignore_label)
231        to_ignore_boundaries = skimage.segmentation.find_boundaries(labels_ignore, mode=self.mode)[None]
232
233        # mask the to-background-boundaries
234        boundaries = boundaries.astype(np.int8)
235        boundaries[to_ignore_boundaries] = self.ignore_label
236
237        if self.add_binary_target:
238            binary = labels_to_binary(labels).astype(boundaries.dtype)
239            binary[labels == self.ignore_label] = self.ignore_label
240            target = np.concatenate([binary[None], boundaries], axis=0)
241        else:
242            target = boundaries
243
244        return target
245
246
247# TODO affinity smoothing
248class AffinityTransform:
249    """Transformation to compute affinities from a segmentation.
250
251    Args:
252        offsets: The offsets for computing affinities.
253        ignore_label: The ignore label to use for computing the ignore mask.
254        add_binary_target: Whether to add a binary channel to the affinities.
255        add_mask: Whether to add the ignore mask as extra output channels.
256        include_ignore_transitions: Whether transitions to the ignore label
257            should be positive in the ignore mask or negative in it.
258    """
259    def __init__(
260        self,
261        offsets: List[List[int]],
262        ignore_label: Optional[bool] = None,
263        add_binary_target: bool = False,
264        add_mask: bool = False,
265        include_ignore_transitions: bool = False,
266    ):
267        assert compute_affinities is not None
268        self.offsets = offsets
269        self.ndim = len(self.offsets[0])
270        assert self.ndim in (2, 3)
271
272        self.ignore_label = ignore_label
273        self.add_binary_target = add_binary_target
274        self.add_mask = add_mask
275        self.include_ignore_transitions = include_ignore_transitions
276
277    def add_ignore_transitions(self, affs, mask, labels):
278        """@private
279        """
280        ignore_seg = (labels == self.ignore_label).astype(labels.dtype)
281        ignore_transitions, invalid_mask = compute_affinities(ignore_seg, self.offsets)
282        invalid_mask = np.logical_not(invalid_mask)
283        # NOTE affinity convention returned by affogato: transitions are marked by 0
284        ignore_transitions = ignore_transitions == 0
285        ignore_transitions[invalid_mask] = 0
286        affs[ignore_transitions] = 1
287        mask[ignore_transitions] = 1
288        return affs, mask
289
290    def __call__(self, labels: np.ndarray) -> np.ndarray:
291        """Compute the affinities.
292
293        Args:
294            labels: The segmentation.
295
296        Returns:
297            The affinities.
298        """
299        dtype = "uint64"
300        if np.dtype(labels.dtype) in (np.dtype("int16"), np.dtype("int32"), np.dtype("int64")):
301            dtype = "int64"
302        labels = ensure_spatial_array(labels, self.ndim, dtype=dtype)
303        affs, mask = compute_affinities(labels, self.offsets,
304                                        have_ignore_label=self.ignore_label is not None,
305                                        ignore_label=0 if self.ignore_label is None else self.ignore_label)
306        # we use the "disaffinity" convention for training; i.e. 1 means repulsive, 0 attractive
307        affs = 1. - affs
308
309        # remove transitions to the ignore label from the mask
310        if self.ignore_label is not None and self.include_ignore_transitions:
311            affs, mask = self.add_ignore_transitions(affs, mask, labels)
312
313        if self.add_binary_target:
314            binary = labels_to_binary(labels)[None].astype(affs.dtype)
315            assert binary.ndim == affs.ndim
316            affs = np.concatenate([binary, affs], axis=0)
317
318        if self.add_mask:
319            if self.add_binary_target:
320                if self.ignore_label is None:
321                    mask_for_bin = np.ones((1,) + labels.shape, dtype=mask.dtype)
322                else:
323                    mask_for_bin = (labels != self.ignore_label)[None].astype(mask.dtype)
324                assert mask.ndim == mask_for_bin.ndim
325                mask = np.concatenate([mask_for_bin, mask], axis=0)
326            assert affs.shape == mask.shape
327            affs = np.concatenate([affs, mask.astype(affs.dtype)], axis=0)
328
329        return affs
330
331
332class OneHotTransform:
333    """Transformations to compute one-hot labels from a semantic segmentation.
334
335    Args:
336        class_ids: The class ids to convert to one-hot labels.
337    """
338    def __init__(self, class_ids: Optional[Union[int, Sequence[int]]] = None):
339        self.class_ids = list(range(class_ids)) if isinstance(class_ids, int) else class_ids
340
341    def __call__(self, labels: np.ndarray) -> np.ndarray:
342        """Compute the one hot transformation.
343
344        Args:
345            labels: The segmentation.
346
347        Returns:
348            The one-hot transformation.
349        """
350        class_ids = np.unique(labels).tolist() if self.class_ids is None else self.class_ids
351        n_classes = len(class_ids)
352        one_hot = np.zeros((n_classes,) + labels.shape, dtype="float32")
353        for i, class_id in enumerate(class_ids):
354            one_hot[i][labels == class_id] = 1.0
355        return one_hot
356
357
358class DistanceTransform:
359    """Transformation to compute distances to foreground in the labels.
360
361    Args:
362        distances: Whether to compute the absolute distances.
363        directed_distances: Whether to compute the directed distances (vector distances).
364        normalize: Whether to normalize the computed distances.
365        max_distance: Maximal distance at which to threshold the distances.
366        foreground_id: Label id to which the distance is compute.
367        invert Whether to invert the distances:
368        func: Normalization function for the distances.
369    """
370    eps = 1e-7
371
372    def __init__(
373        self,
374        distances: bool = True,
375        directed_distances: bool = False,
376        normalize: bool = True,
377        max_distance: Optional[float] = None,
378        foreground_id: int = 1,
379        invert: bool = False,
380        func: Optional[Callable] = None,
381    ):
382        if sum((distances, directed_distances)) == 0:
383            raise ValueError("At least one of 'distances' or 'directed_distances' must be set to 'True'")
384        self.directed_distances = directed_distances
385        self.distances = distances
386        self.normalize = normalize
387        self.max_distance = max_distance
388        self.foreground_id = foreground_id
389        self.invert = invert
390        self.func = func
391
392    def _compute_distances(self, directed_distances):
393        distances = np.linalg.norm(directed_distances, axis=0)
394        if self.max_distance is not None:
395            distances = np.clip(distances, 0, self.max_distance)
396        if self.normalize:
397            distances /= (distances.max() + self.eps)
398        if self.invert:
399            distances = distances.max() - distances
400        if self.func is not None:
401            distances = self.func(distances)
402        return distances
403
404    def _compute_directed_distances(self, directed_distances):
405        if self.max_distance is not None:
406            directed_distances = np.clip(directed_distances, -self.max_distance, self.max_distance)
407        if self.normalize:
408            directed_distances /= (np.abs(directed_distances).max(axis=(1, 2), keepdims=True) + self.eps)
409        if self.invert:
410            directed_distances = directed_distances.max(axis=(1, 2), keepdims=True) - directed_distances
411        if self.func is not None:
412            directed_distances = self.func(directed_distances)
413        return directed_distances
414
415    def _get_distances_for_empty_labels(self, labels):
416        shape = labels.shape
417        fill_value = 0.0 if self.invert else np.sqrt(np.linalg.norm(list(shape)) ** 2 / 2)
418        data = np.full((labels.ndim,) + shape, fill_value)
419        return data
420
421    def __call__(self, labels: np.ndarray) -> np.ndarray:
422        """Compute the distances.
423
424        Args:
425            labels: The segmentation.
426
427        Returns:
428            The distances.
429        """
430        distance_mask = (labels == self.foreground_id).astype("uint32")
431        # the distances are not computed corrected if they are all zero
432        # so this case needs to be handled separately
433        if distance_mask.sum() == 0:
434            directed_distances = self._get_distances_for_empty_labels(labels)
435        else:
436            ndim = distance_mask.ndim
437            to_channel_first = (ndim,) + tuple(range(ndim))
438            directed_distances = vigra.filters.vectorDistanceTransform(distance_mask).transpose(to_channel_first)
439
440        if self.distances:
441            distances = self._compute_distances(directed_distances)
442
443        if self.directed_distances:
444            directed_distances = self._compute_directed_distances(directed_distances)
445
446        if self.distances and self.directed_distances:
447            return np.concatenate((distances[None], directed_distances), axis=0)
448        if self.distances:
449            return distances
450        if self.directed_distances:
451            return directed_distances
452
453
454class PerObjectDistanceTransform:
455    """Transformation to compute normalized distances per object in a segmentation.
456
457    Args:
458        distances: Whether to compute the undirected distances.
459        boundary_distances: Whether to compute the distances to the object boundaries.
460        directed_distances: Whether to compute the directed distances (vector distances).
461        foreground: Whether to return a foreground channel.
462        apply_label: Whether to apply connected components to the labels before computing distances.
463        correct_centers: Whether to correct centers that are not in the objects.
464        min_size: Minimal size of objects for distance calculdation.
465        distance_fill_value: Fill value for the distances outside of objects.
466    """
467    eps = 1e-7
468
469    def __init__(
470        self,
471        distances: bool = True,
472        boundary_distances: bool = True,
473        directed_distances: bool = False,
474        foreground: bool = True,
475        instances: bool = False,
476        apply_label: bool = True,
477        correct_centers: bool = True,
478        min_size: int = 0,
479        distance_fill_value: float = 1.0,
480    ):
481        if sum([distances, directed_distances, boundary_distances]) == 0:
482            raise ValueError("At least one of distances or directed distances has to be passed.")
483        self.distances = distances
484        self.boundary_distances = boundary_distances
485        self.directed_distances = directed_distances
486        self.foreground = foreground
487        self.instances = instances
488
489        self.apply_label = apply_label
490        self.correct_centers = correct_centers
491        self.min_size = min_size
492        self.distance_fill_value = distance_fill_value
493
494    def compute_normalized_object_distances(self, mask, boundaries, bb, center, distances):
495        """@private
496        """
497        # Crop the mask and generate array with the correct center.
498        cropped_mask = mask[bb]
499        cropped_center = tuple(ce - b.start for ce, b in zip(center, bb))
500
501        # The centroid might not be inside of the object.
502        # In this case we correct the center by taking the maximum of the distance to the boundary.
503        # Note: the centroid is still the best estimate for the center, as long as it's in the object.
504        correct_center = not cropped_mask[cropped_center]
505
506        # Compute the boundary distances if necessary.
507        # (Either if we need to correct the center, or compute the boundary distances anyways.)
508        if correct_center or self.boundary_distances:
509            # Crop the boundary mask and compute the boundary distances.
510            cropped_boundary_mask = boundaries[bb]
511            boundary_distances = vigra.filters.distanceTransform(cropped_boundary_mask)
512            boundary_distances[~cropped_mask] = 0
513            max_dist_point = np.unravel_index(np.argmax(boundary_distances), boundary_distances.shape)
514
515        # Set the crop center to the max dist point
516        if correct_center:
517            # Find the center (= maximal distance from the boundaries).
518            cropped_center = max_dist_point
519
520        cropped_center_mask = np.zeros_like(cropped_mask, dtype="uint32")
521        cropped_center_mask[cropped_center] = 1
522
523        # Compute the directed distances,
524        if self.distances or self.directed_distances:
525            this_distances = vigra.filters.vectorDistanceTransform(cropped_center_mask)
526        else:
527            this_distances = None
528
529        # Keep only the specified distances:
530        if self.distances and self.directed_distances:  # all distances
531            # Compute the undirected ditacnes from directed distances and concatenate,
532            undir = np.linalg.norm(this_distances, axis=-1, keepdims=True)
533            this_distances = np.concatenate([undir, this_distances], axis=-1)
534
535        elif self.distances:  # only undirected distances
536            # Compute the undirected distances from directed distances and keep only them.
537            this_distances = np.linalg.norm(this_distances, axis=-1, keepdims=True)
538
539        elif self.directed_distances:  # only directed distances
540            pass  # We don't have to do anything becasue the directed distances are already computed.
541
542        # Add an extra channel for the boundary distances if specified.
543        if self.boundary_distances:
544            boundary_distances = (boundary_distances[max_dist_point] - boundary_distances)[..., None]
545            if this_distances is None:
546                this_distances = boundary_distances
547            else:
548                this_distances = np.concatenate([this_distances, boundary_distances], axis=-1)
549
550        # Set distances outside of the mask to zero.
551        this_distances[~cropped_mask] = 0
552
553        # Normalize the distances.
554        spatial_axes = tuple(range(mask.ndim))
555        this_distances /= (np.abs(this_distances).max(axis=spatial_axes, keepdims=True) + self.eps)
556
557        # Set the distance values in the global result.
558        distances[bb][cropped_mask] = this_distances[cropped_mask]
559
560        return distances
561
562    def __call__(self, labels: np.ndarray) -> np.ndarray:
563        """Compute the per object distance transform.
564
565        Args:
566            labels: The segmentation
567
568        Returns:
569            The distances.
570        """
571        # Apply label (connected components) if specified.
572        if self.apply_label:
573            labels = skimage.measure.label(labels).astype("uint32")
574        else:  # Otherwise just relabel the segmentation.
575            labels = vigra.analysis.relabelConsecutive(labels)[0].astype("uint32")
576
577        # Filter out small objects if min_size is specified.
578        if self.min_size > 0:
579            ids, sizes = np.unique(labels, return_counts=True)
580            discard_ids = ids[sizes < self.min_size]
581            labels[np.isin(labels, discard_ids)] = 0
582            labels = vigra.analysis.relabelConsecutive(labels)[0].astype("uint32")
583
584        # Compute the boundaries. They will be used to determine the most central point,
585        # and if 'self.boundary_distances is True' to add the boundary distances.
586        boundaries = skimage.segmentation.find_boundaries(labels, mode="inner").astype("uint32")
587
588        # Compute region properties to derive bounding boxes and centers.
589        ndim = labels.ndim
590        props = skimage.measure.regionprops(labels)
591        bounding_boxes = {
592            prop.label: tuple(slice(prop.bbox[i], prop.bbox[i + ndim]) for i in range(ndim)) for prop in props
593        }
594
595        # Compute the object centers from centroids.
596        centers = {prop.label: np.round(prop.centroid).astype("int") for prop in props}
597
598        # Compute how many distance channels we have.
599        n_channels = 0
600        if self.distances:  # We need one channel for the overall distances.
601            n_channels += 1
602        if self.boundary_distances:  # We need one channel for the boundary distances.
603            n_channels += 1
604        if self.directed_distances:  # And ndim channels for directed distances.
605            n_channels += ndim
606
607        # Compute the per object distances.
608        distances = np.full(labels.shape + (n_channels,), self.distance_fill_value, dtype="float32")
609        for prop in props:
610            label_id = prop.label
611            mask = labels == label_id
612            distances = self.compute_normalized_object_distances(
613                mask, boundaries, bounding_boxes[label_id], centers[label_id], distances
614            )
615
616        # Bring the distance channel to the first dimension.
617        to_channel_first = (ndim,) + tuple(range(ndim))
618        distances = distances.transpose(to_channel_first)
619
620        # Add the foreground mask as first channel if specified.
621        if self.foreground:
622            binary_labels = (labels > 0).astype("float32")
623            distances = np.concatenate([binary_labels[None], distances], axis=0)
624
625        if self.instances:
626            distances = np.concatenate([labels[None], distances], axis=0)
627
628        return distances