torch_em.transform.label

  1from typing import Callable, List, Optional, Sequence, Union, Tuple
  2
  3import numpy as np
  4import skimage.measure
  5import skimage.segmentation
  6import bioimage_cpp as bic
  7
  8from ..util import ensure_array, ensure_spatial_array
  9
 10try:
 11    from bioimage_cpp.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 = bic.segmentation.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 = bic.segmentation.relabel_sequential(labels)[0]
 59    else:
 60        if 0 in labels:
 61            labels += 1
 62        seg = bic.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, _, _ = bic.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[int] = 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 compute_affinities: 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, ignore_label=self.ignore_label)
304        # we use the "disaffinity" convention for training; i.e. 1 means repulsive, 0 attractive
305        affs = 1. - affs
306
307        # remove transitions to the ignore label from the mask
308        if self.ignore_label is not None and self.include_ignore_transitions:
309            affs, mask = self.add_ignore_transitions(affs, mask, labels)
310
311        if self.add_binary_target:
312            binary = labels_to_binary(labels)[None].astype(affs.dtype)
313            assert binary.ndim == affs.ndim
314            affs = np.concatenate([binary, affs], axis=0)
315
316        if self.add_mask:
317            if self.add_binary_target:
318                if self.ignore_label is None:
319                    mask_for_bin = np.ones((1,) + labels.shape, dtype=mask.dtype)
320                else:
321                    mask_for_bin = (labels != self.ignore_label)[None].astype(mask.dtype)
322                assert mask.ndim == mask_for_bin.ndim
323                mask = np.concatenate([mask_for_bin, mask], axis=0)
324            assert affs.shape == mask.shape
325            affs = np.concatenate([affs, mask.astype(affs.dtype)], axis=0)
326
327        return affs
328
329
330class OneHotTransform:
331    """Transformations to compute one-hot labels from a semantic segmentation.
332
333    Args:
334        class_ids: The class ids to convert to one-hot labels.
335    """
336    def __init__(self, class_ids: Optional[Union[int, Sequence[int]]] = None):
337        self.class_ids = list(range(class_ids)) if isinstance(class_ids, int) else class_ids
338
339    def __call__(self, labels: np.ndarray) -> np.ndarray:
340        """Compute the one hot transformation.
341
342        Args:
343            labels: The segmentation.
344
345        Returns:
346            The one-hot transformation.
347        """
348        class_ids = np.unique(labels).tolist() if self.class_ids is None else self.class_ids
349        n_classes = len(class_ids)
350        one_hot = np.zeros((n_classes,) + labels.shape, dtype="float32")
351        for i, class_id in enumerate(class_ids):
352            one_hot[i][labels == class_id] = 1.0
353        return one_hot
354
355
356class DistanceTransform:
357    """Transformation to compute distances to foreground in the labels.
358
359    Args:
360        distances: Whether to compute the absolute distances.
361        directed_distances: Whether to compute the directed distances (vector distances).
362        normalize: Whether to normalize the computed distances.
363        max_distance: Maximal distance at which to threshold the distances.
364        foreground_id: Label id to which the distance is compute.
365        invert: Whether to invert the distances.
366        func: Normalization function for the distances.
367    """
368    eps = 1e-7
369
370    def __init__(
371        self,
372        distances: bool = True,
373        directed_distances: bool = False,
374        normalize: bool = True,
375        max_distance: Optional[float] = None,
376        foreground_id: int = 1,
377        invert: bool = False,
378        func: Optional[Callable] = None,
379    ):
380        if sum((distances, directed_distances)) == 0:
381            raise ValueError("At least one of 'distances' or 'directed_distances' must be set to 'True'")
382        self.directed_distances = directed_distances
383        self.distances = distances
384        self.normalize = normalize
385        self.max_distance = max_distance
386        self.foreground_id = foreground_id
387        self.invert = invert
388        self.func = func
389
390    def _compute_distances(self, directed_distances):
391        distances = np.linalg.norm(directed_distances, axis=0)
392        if self.max_distance is not None:
393            distances = np.clip(distances, 0, self.max_distance)
394        if self.normalize:
395            distances /= (distances.max() + self.eps)
396        if self.invert:
397            distances = distances.max() - distances
398        if self.func is not None:
399            distances = self.func(distances)
400        return distances
401
402    def _compute_directed_distances(self, directed_distances):
403        if self.max_distance is not None:
404            directed_distances = np.clip(directed_distances, -self.max_distance, self.max_distance)
405        if self.normalize:
406            directed_distances /= (np.abs(directed_distances).max(axis=(1, 2), keepdims=True) + self.eps)
407        if self.invert:
408            directed_distances = directed_distances.max(axis=(1, 2), keepdims=True) - directed_distances
409        if self.func is not None:
410            directed_distances = self.func(directed_distances)
411        return directed_distances
412
413    def _get_distances_for_empty_labels(self, labels):
414        shape = labels.shape
415        fill_value = 0.0 if self.invert else np.sqrt(np.linalg.norm(list(shape)) ** 2 / 2)
416        data = np.full((labels.ndim,) + shape, fill_value)
417        return data
418
419    def __call__(self, labels: np.ndarray) -> np.ndarray:
420        """Compute the distances.
421
422        Args:
423            labels: The segmentation.
424
425        Returns:
426            The distances.
427        """
428        distance_mask = (labels == self.foreground_id).astype("uint32")
429        # the distances are not computed corrected if they are all zero
430        # so this case needs to be handled separately
431        if distance_mask.sum() == 0:
432            directed_distances = self._get_distances_for_empty_labels(labels)
433        else:
434            ndim = distance_mask.ndim
435            to_channel_first = (ndim,) + tuple(range(ndim))
436            directed_distances = bic.distance.vector_difference_transform(
437                distance_mask == 0
438            ).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        instances: Whether to append the original labels as an extra channel to the target.
463        apply_label: Whether to apply connected components to the labels before computing distances.
464        correct_centers: Whether to correct centers that are not in the objects.
465        min_size: Minimal size of objects for distance calculdation.
466        distance_fill_value: Fill value for the distances outside of objects.
467        sampling: The spacing of the distance transform. This is especially relevant for anisotropic data;
468            for which it is recommended to use a sampling of (ANISOTROPY_FACTOR, 1, 1).
469    """
470    eps = 1e-7
471
472    def __init__(
473        self,
474        distances: bool = True,
475        boundary_distances: bool = True,
476        directed_distances: bool = False,
477        foreground: bool = True,
478        instances: bool = False,
479        apply_label: bool = True,
480        correct_centers: bool = True,
481        min_size: int = 0,
482        distance_fill_value: float = 1.0,
483        sampling: Optional[Tuple[float, ...]] = None
484    ):
485        if sum([distances, directed_distances, boundary_distances]) == 0:
486            raise ValueError("At least one of distances or directed distances has to be passed.")
487        self.distances = distances
488        self.boundary_distances = boundary_distances
489        self.directed_distances = directed_distances
490        self.foreground = foreground
491        self.instances = instances
492
493        self.apply_label = apply_label
494        self.correct_centers = correct_centers
495        self.min_size = min_size
496        self.distance_fill_value = distance_fill_value
497        self.sampling = sampling
498
499    def compute_normalized_object_distances(self, mask, boundaries, bb, center, distances):
500        """@private
501        """
502        # Crop the mask and generate array with the correct center.
503        cropped_mask = mask[bb]
504        cropped_center = tuple(ce - b.start for ce, b in zip(center, bb))
505
506        # The centroid might not be inside of the object.
507        # In this case we correct the center by taking the maximum of the distance to the boundary.
508        # Note: the centroid is still the best estimate for the center, as long as it's in the object.
509        correct_center = not cropped_mask[cropped_center]
510
511        # Compute the boundary distances if necessary.
512        # (Either if we need to correct the center, or compute the boundary distances anyways.)
513        if correct_center or self.boundary_distances:
514            # Crop the boundary mask and compute the boundary distances.
515            cropped_boundary_mask = boundaries[bb]
516            boundary_distances = bic.distance.distance_transform(cropped_boundary_mask == 0, sampling=self.sampling)
517            boundary_distances[~cropped_mask] = 0
518            max_dist_point = np.unravel_index(np.argmax(boundary_distances), boundary_distances.shape)
519
520        # Set the crop center to the max dist point
521        if correct_center:
522            # Find the center (= maximal distance from the boundaries).
523            cropped_center = max_dist_point
524
525        cropped_center_mask = np.zeros_like(cropped_mask, dtype="uint32")
526        cropped_center_mask[cropped_center] = 1
527
528        # Compute the directed distances,
529        if self.distances or self.directed_distances:
530            this_distances = bic.distance.vector_difference_transform(cropped_center_mask == 0, sampling=self.sampling)
531        else:
532            this_distances = None
533
534        # Keep only the specified distances:
535        if self.distances and self.directed_distances:  # all distances
536            # Compute the undirected distances from directed distances and concatenate,
537            undir = np.linalg.norm(this_distances, axis=-1, keepdims=True)
538            this_distances = np.concatenate([undir, this_distances], axis=-1)
539
540        elif self.distances:  # only undirected distances
541            # Compute the undirected distances from directed distances and keep only them.
542            this_distances = np.linalg.norm(this_distances, axis=-1, keepdims=True)
543
544        elif self.directed_distances:  # only directed distances
545            pass  # We don't have to do anything becasue the directed distances are already computed.
546
547        # Add an extra channel for the boundary distances if specified.
548        if self.boundary_distances:
549            boundary_distances = (boundary_distances[max_dist_point] - boundary_distances)[..., None]
550            if this_distances is None:
551                this_distances = boundary_distances
552            else:
553                this_distances = np.concatenate([this_distances, boundary_distances], axis=-1)
554
555        # Set distances outside of the mask to zero.
556        this_distances[~cropped_mask] = 0
557
558        # Normalize the distances.
559        spatial_axes = tuple(range(mask.ndim))
560        this_distances /= (np.abs(this_distances).max(axis=spatial_axes, keepdims=True) + self.eps)
561
562        # Set the distance values in the global result.
563        distances[bb][cropped_mask] = this_distances[cropped_mask]
564
565        return distances
566
567    def __call__(self, labels: np.ndarray) -> np.ndarray:
568        """Compute the per object distance transform.
569
570        Args:
571            labels: The segmentation
572
573        Returns:
574            The distances.
575        """
576        # Apply label (connected components) if specified.
577        if self.apply_label:
578            labels = bic.segmentation.label(labels).astype("uint32")
579        else:  # Otherwise just relabel the segmentation.
580            labels = bic.segmentation.relabel_sequential(labels)[0].astype("uint32")
581
582        # Filter out small objects if min_size is specified.
583        if self.min_size > 0:
584            ids, sizes = np.unique(labels, return_counts=True)
585            discard_ids = ids[sizes < self.min_size]
586            labels[np.isin(labels, discard_ids)] = 0
587            labels = bic.segmentation.relabel_sequential(labels)[0].astype("uint32")
588
589        # Compute the boundaries. They will be used to determine the most central point,
590        # and if 'self.boundary_distances is True' to add the boundary distances.
591        boundaries = skimage.segmentation.find_boundaries(labels, mode="inner").astype("uint32")
592
593        # Compute region properties to derive bounding boxes and centers.
594        ndim = labels.ndim
595        props = skimage.measure.regionprops(labels)
596        bounding_boxes = {
597            prop.label: tuple(slice(prop.bbox[i], prop.bbox[i + ndim]) for i in range(ndim)) for prop in props
598        }
599
600        # Compute the object centers from centroids.
601        centers = {prop.label: np.round(prop.centroid).astype("int") for prop in props}
602
603        # Compute how many distance channels we have.
604        n_channels = 0
605        if self.distances:  # We need one channel for the overall distances.
606            n_channels += 1
607        if self.boundary_distances:  # We need one channel for the boundary distances.
608            n_channels += 1
609        if self.directed_distances:  # And ndim channels for directed distances.
610            n_channels += ndim
611
612        # Compute the per object distances.
613        distances = np.full(labels.shape + (n_channels,), self.distance_fill_value, dtype="float32")
614        for prop in props:
615            label_id = prop.label
616            mask = labels == label_id
617            distances = self.compute_normalized_object_distances(
618                mask, boundaries, bounding_boxes[label_id], centers[label_id], distances
619            )
620
621        # Bring the distance channel to the first dimension.
622        to_channel_first = (ndim,) + tuple(range(ndim))
623        distances = distances.transpose(to_channel_first)
624
625        # Add the foreground mask as first channel if specified.
626        if self.foreground:
627            binary_labels = (labels > 0).astype("float32")
628            distances = np.concatenate([binary_labels[None], distances], axis=0)
629
630        if self.instances:
631            distances = np.concatenate([labels[None], distances], axis=0)
632
633        return distances