torch_em.metric.instance_segmentation_metric

  1from functools import partial
  2from typing import List, Optional
  3
  4import numpy as np
  5import elf.evaluation as elfval
  6import elf.segmentation as elfseg
  7import elf.segmentation.embeddings as elfemb
  8import torch
  9import torch.nn as nn
 10import bioimage_cpp as bic
 11from skimage.measure import regionprops
 12from elf.segmentation.watershed import apply_size_filter
 13
 14
 15class BaseInstanceSegmentationMetric(nn.Module):
 16    """@private
 17    """
 18    def __init__(self, segmenter, metric, to_numpy=True):
 19        super().__init__()
 20        self.segmenter = segmenter
 21        self.metric = metric
 22        self.to_numpy = to_numpy
 23
 24    def forward(self, input_, target):
 25        if self.to_numpy:
 26            input_ = input_.detach().cpu().numpy().astype("float32")
 27            target = target.detach().cpu().numpy()
 28        assert input_.ndim == target.ndim
 29        assert len(input_) == len(target)
 30        scores = []
 31        # compute the metric per batch
 32        for pred, trgt in zip(input_, target):
 33            seg = self.segmenter(pred)
 34            # by convention we assume that the segmentation channel is always in the last channel of trgt
 35            scores.append(self.metric(seg, trgt[-1].astype("uint32")))
 36        return torch.tensor(scores).mean()
 37
 38
 39#
 40# Segmenters
 41#
 42
 43def filter_sizes(seg, min_seg_size, hmap=None):
 44    """@private
 45    """
 46    seg_ids, counts = np.unique(seg, return_counts=True)
 47    if hmap is None:
 48        bg_ids = seg_ids[counts < min_seg_size]
 49        seg[np.isin(seg, bg_ids)] = 0
 50    else:
 51        ndim = seg.ndim
 52        hmap_ = hmap if hmap.ndim == ndim else np.max(hmap, axis=0)
 53        seg, _ = apply_size_filter(seg, hmap_, min_seg_size)
 54    return seg
 55
 56
 57class MWS:
 58    """@private
 59    """
 60    def __init__(self, offsets, with_background, min_seg_size, strides=None):
 61        self.offsets = offsets
 62        self.with_background = with_background
 63        self.min_seg_size = min_seg_size
 64        if strides is None:
 65            strides = [4] * len(offsets[0])
 66        assert len(strides) == len(offsets[0])
 67        self.strides = strides
 68
 69    def __call__(self, affinities):
 70        if self.with_background:
 71            assert len(affinities) == len(self.offsets) + 1
 72            mask, affinities = affinities[0], affinities[1:]
 73            # bioimage_cpp's mutex_watershed requires a boolean mask; treat the foreground
 74            # channel as a probability and threshold it (matching torch_em.util.segmentation).
 75            mask = mask >= 0.5
 76        else:
 77            assert len(affinities) == len(self.offsets)
 78            mask = None
 79        seg = elfseg.mutex_watershed.mutex_watershed(affinities, self.offsets, self.strides,
 80                                                     randomize_strides=True, mask=mask).astype("uint32")
 81        if self.min_seg_size > 0:
 82            seg = filter_sizes(seg, self.min_seg_size,
 83                               hmap=None if self.with_background else affinities)
 84        return seg
 85
 86
 87class EmbeddingMWS:
 88    """@private
 89    """
 90    def __init__(self, delta, offsets, with_background, min_seg_size, strides=None):
 91        self.delta = delta
 92        self.offsets = offsets
 93        self.with_background = with_background
 94        self.min_seg_size = min_seg_size
 95        if strides is None:
 96            strides = [4] * len(offsets[0])
 97        assert len(strides) == len(offsets[0])
 98        self.strides = strides
 99
100    def merge_background(self, seg, embeddings):
101        seg += 1
102        seg_ids, counts = np.unique(seg, return_counts=True)
103        bg_seg = seg_ids[np.argmax(counts)]
104        n_labels = int(seg.max()) + 1
105        mean_embeddings = []
106        for emb in embeddings:
107            means = np.zeros(n_labels, dtype="float32")
108            for prop in regionprops(seg, intensity_image=np.asarray(emb)):
109                means[prop.label] = prop.mean_intensity
110            mean_embeddings.append(means[None])
111        mean_embeddings = np.concatenate(mean_embeddings, axis=0)
112        bg_embed = mean_embeddings[:, bg_seg][:, None]
113        bg_probs = elfemb._embeddings_to_probabilities(mean_embeddings, bg_embed, self.delta, 0)
114        bg_ids = np.where(bg_probs > 0.5)
115        seg[np.isin(seg, bg_ids)] = 0
116        seg, _, _ = bic.segmentation.relabel_sequential(seg)
117        return seg
118
119    def __call__(self, embeddings):
120        weight = partial(elfemb.discriminative_loss_weight, delta=self.delta)
121        seg = elfemb.segment_embeddings_mws(
122            embeddings, "l2", self.offsets, strides=self.strides, weight_function=weight
123        ).astype("uint32")
124        if self.with_background:
125            seg = self.merge_background(seg, embeddings)
126        if self.min_seg_size > 0:
127            seg = filter_sizes(seg, self.min_seg_size)
128        return seg
129
130
131class Multicut:
132    """@private
133    """
134    def __init__(self, min_seg_size, anisotropic=False, dt_threshold=0.25, sigma_seeds=2.0, solver="decomposition"):
135        self.min_seg_size = min_seg_size
136        self.anisotropic = anisotropic
137        self.dt_threshold = dt_threshold
138        self.sigma_seeds = sigma_seeds
139        self.solver = solver
140
141    def __call__(self, boundaries):
142        if boundaries.shape[0] == 1:
143            boundaries = boundaries[0]
144        assert boundaries.ndim in (2, 3), f"{boundaries.ndim}"
145        if self.anisotropic and boundaries.ndim == 3:
146            ws, max_id = elfseg.stacked_watershed(boundaries, threshold=self.dt_threshold,
147                                                  sigma_seed=self.sigma_seeds,
148                                                  sigma_weights=self.sigma_seeds,
149                                                  n_threads=1)
150        else:
151            ws, max_id = elfseg.distance_transform_watershed(boundaries, threshold=self.dt_threshold,
152                                                             sigma_seeds=self.sigma_seeds,
153                                                             sigma_weights=self.sigma_seeds)
154        rag = elfseg.compute_rag(ws, max_id + 1, n_threads=1)
155        feats = elfseg.compute_boundary_mean_and_length(rag, ws, boundaries, n_threads=1)[:, 0]
156        costs = elfseg.compute_edge_costs(feats)
157        solver = elfseg.get_multicut_solver(self.solver)
158        node_labels = solver(rag, costs, n_threads=1)
159        seg = elfseg.project_node_labels_to_pixels(rag, ws, node_labels, n_threads=1).astype("uint32")
160        if self.min_seg_size > 0:
161            seg = filter_sizes(seg, self.min_seg_size, hmap=boundaries)
162        return seg
163
164
165class HDBScan:
166    """@private
167    """
168    def __init__(self, min_size, eps, remove_largest):
169        self.min_size = min_size
170        self.eps = eps
171        self.remove_largest = remove_largest
172
173    def __call__(self, embeddings):
174        return elfemb.segment_hdbscan(embeddings, self.min_size, self.eps, self.remove_largest)
175
176
177#
178# Metrics
179#
180
181class IOUError:
182    """@private
183    """
184    def __init__(self, threshold=0.5, metric="precision"):
185        self.threshold = threshold
186        self.metric = metric
187
188    def __call__(self, seg, target):
189        score = 1.0 - elfval.matching(seg, target, threshold=self.threshold)[self.metric]
190        return score
191
192
193class VariationOfInformation:
194    """@private
195    """
196    def __call__(self, seg, target):
197        vis, vim = elfval.variation_of_information(seg, target)
198        return vis + vim
199
200
201class AdaptedRandError:
202    """@private
203    """
204    def __call__(self, seg, target):
205        are, _ = elfval.rand_index(seg, target)
206        return are
207
208
209class SymmetricBestDice:
210    """@private
211    """
212    def __call__(self, seg, target):
213        score = 1.0 - elfval.symmetric_best_dice_score(seg, target)
214        return score
215
216
217#
218# Prefab Full Metrics
219#
220
221
222class EmbeddingMWSIOUMetric(BaseInstanceSegmentationMetric):
223    """Intersection over union metric based on mutex watershed computed from embedding-derived affinites.
224
225    This class can be used as validation metric when training a network for instance segmentation.
226
227    Args:
228        delta: The hinge distance of the contrastive loss for training the embeddings.
229        offsets: The offsets for deriving the affinities from the embeddings.
230        min_seg_size: Size for filtering the segmentation objects.
231        iou_threshold: Threshold for the intersection over union metric.
232        strides: The strides for the mutex watershed.
233    """
234    def __init__(
235        self,
236        delta: float,
237        offsets: List[List[int]],
238        min_seg_size: int,
239        iou_threshold: float = 0.5,
240        strides: Optional[List[int]] = None,
241    ):
242        segmenter = EmbeddingMWS(delta, offsets, with_background=True, min_seg_size=min_seg_size)
243        metric = IOUError(iou_threshold)
244        super().__init__(segmenter, metric)
245        self.init_kwargs = {"delta": delta, "offsets": offsets, "min_seg_size": min_seg_size,
246                            "iou_threshold": iou_threshold, "strides": strides}
247
248
249class EmbeddingMWSSBDMetric(BaseInstanceSegmentationMetric):
250    """Symmetric best dice metric based on mutex watershed computed from embedding-derived affinites.
251
252    This class can be used as validation metric when training a network for instance segmentation.
253
254    Args:
255        delta: The hinge distance of the contrastive loss for training the embeddings.
256        offsets: The offsets for deriving the affinities from the embeddings.
257        min_seg_size: Size for filtering the segmentation objects.
258        strides: The strides for the mutex watershed.
259    """
260    def __init__(self, delta: float, offsets: List[List[int]], min_seg_size: int, strides: Optional[List[int]] = None):
261        segmenter = EmbeddingMWS(delta, offsets, with_background=True, min_seg_size=min_seg_size)
262        metric = SymmetricBestDice()
263        super().__init__(segmenter, metric)
264        self.init_kwargs = {"delta": delta, "offsets": offsets, "min_seg_size": min_seg_size, "strides": strides}
265
266
267class EmbeddingMWSVOIMetric(BaseInstanceSegmentationMetric):
268    """Variation of inofrmation metric based on mutex watershed computed from embedding-derived affinites.
269
270    This class can be used as validation metric when training a network for instance segmentation.
271
272    Args:
273        delta: The hinge distance of the contrastive loss for training the embeddings.
274        offsets: The offsets for deriving the affinities from the embeddings.
275        min_seg_size: Size for filtering the segmentation objects.
276        strides: The strides for the mutex watershed.
277    """
278    def __init__(self, delta: float, offsets: List[List[int]], min_seg_size: int, strides: Optional[List[int]] = None):
279        segmenter = EmbeddingMWS(delta, offsets, with_background=False, min_seg_size=min_seg_size)
280        metric = VariationOfInformation()
281        super().__init__(segmenter, metric)
282        self.init_kwargs = {"delta": delta, "offsets": offsets, "min_seg_size": min_seg_size, "strides": strides}
283
284
285class EmbeddingMWSRandMetric(BaseInstanceSegmentationMetric):
286    """Rand index metric based on mutex watershed computed from embedding-derived affinites.
287
288    This class can be used as validation metric when training a network for instance segmentation.
289
290    Args:
291        delta: The hinge distance of the contrastive loss for training the embeddings.
292        offsets: The offsets for deriving the affinities from the embeddings.
293        min_seg_size: Size for filtering the segmentation objects.
294        strides: The strides for the mutex watershed.
295    """
296    def __init__(self, delta: float, offsets: List[List[int]], min_seg_size: int, strides: Optional[List[int]] = None):
297        segmenter = EmbeddingMWS(delta, offsets, with_background=False, min_seg_size=min_seg_size)
298        metric = AdaptedRandError()
299        super().__init__(segmenter, metric)
300        self.init_kwargs = {"delta": delta, "offsets": offsets, "min_seg_size": min_seg_size, "strides": strides}
301
302
303class HDBScanIOUMetric(BaseInstanceSegmentationMetric):
304    """Intersection over union metric based on HDBScan computed from embeddings.
305
306    This class can be used as validation metric when training a network for instance segmentation.
307
308    Args:
309        min_size: The minimal segment size.
310        eps: The epsilon value for HDBScan.
311        iou_threshold: The threshold for the intersection over union value.
312    """
313    def __init__(self, min_size: int, eps: float, iou_threshold: float = 0.5):
314        segmenter = HDBScan(min_size=min_size, eps=eps, remove_largest=True)
315        metric = IOUError(iou_threshold)
316        super().__init__(segmenter, metric)
317        self.init_kwargs = {"min_size": min_size, "eps": eps, "iou_threshold": iou_threshold}
318
319
320class HDBScanSBDMetric(BaseInstanceSegmentationMetric):
321    """Symmetric best dice metric based on HDBScan computed from embeddings.
322
323    This class can be used as validation metric when training a network for instance segmentation.
324
325    Args:
326        min_size: The minimal segment size.
327        eps: The epsilon value for HDBScan.
328    """
329    def __init__(self, min_size: int, eps: float):
330        segmenter = HDBScan(min_size=min_size, eps=eps, remove_largest=True)
331        metric = SymmetricBestDice()
332        super().__init__(segmenter, metric)
333        self.init_kwargs = {"min_size": min_size, "eps": eps}
334
335
336class HDBScanRandMetric(BaseInstanceSegmentationMetric):
337    """Rand index metric based on HDBScan computed from embeddings.
338
339    This class can be used as validation metric when training a network for instance segmentation.
340
341    Args:
342        min_size: The minimal segment size.
343        eps: The epsilon value for HDBScan.
344    """
345    def __init__(self, min_size: int, eps: float):
346        segmenter = HDBScan(min_size=min_size, eps=eps, remove_largest=True)
347        metric = AdaptedRandError()
348        super().__init__(segmenter, metric)
349        self.init_kwargs = {"min_size": min_size, "eps": eps}
350
351
352class HDBScanVOIMetric(BaseInstanceSegmentationMetric):
353    """Variation of information metric based on HDBScan computed from embeddings.
354
355    This class can be used as validation metric when training a network for instance segmentation.
356
357    Args:
358        min_size: The minimal segment size.
359        eps: The epsilon value for HDBScan.
360    """
361    def __init__(self, min_size: int, eps: float):
362        segmenter = HDBScan(min_size=min_size, eps=eps, remove_largest=True)
363        metric = VariationOfInformation()
364        super().__init__(segmenter, metric)
365        self.init_kwargs = {"min_size": min_size, "eps": eps}
366
367
368class MulticutVOIMetric(BaseInstanceSegmentationMetric):
369    """Variation of information metric based on a multicut computed from boundary predictions.
370
371    This class can be used as validation metric when training a network for instance segmentation.
372
373    Args:
374        min_seg_size: The minimal segment size.
375        anisotropic: Whether to compute the watersheds in 2d for volumetric data.
376        dt_threshold: The threshold to apply to the boundary predictions before computing the distance transform.
377        sigma_seeds: The sigma value for smoothing the distance transform before computing seeds.
378    """
379    def __init__(
380        self, min_seg_size: int, anisotropic: bool = False, dt_threshold: float = 0.25, sigma_seeds: float = 2.0
381    ):
382        segmenter = Multicut(dt_threshold, anisotropic, sigma_seeds)
383        metric = VariationOfInformation()
384        super().__init__(segmenter, metric)
385        self.init_kwargs = {"anisotropic": anisotropic, "min_seg_size": min_seg_size,
386                            "dt_threshold": dt_threshold, "sigma_seeds": sigma_seeds}
387
388
389class MulticutRandMetric(BaseInstanceSegmentationMetric):
390    """Rand index metric based on a multicut computed from boundary predictions.
391
392    This class can be used as validation metric when training a network for instance segmentation.
393
394    Args:
395        min_seg_size: The minimal segment size.
396        anisotropic: Whether to compute the watersheds in 2d for volumetric data.
397        dt_threshold: The threshold to apply to the boundary predictions before computing the distance transform.
398        sigma_seeds: The sigma value for smoothing the distance transform before computing seeds.
399    """
400    def __init__(
401        self, min_seg_size: int, anisotropic: bool = False, dt_threshold: float = 0.25, sigma_seeds: float = 2.0
402    ):
403        segmenter = Multicut(dt_threshold, anisotropic, sigma_seeds)
404        metric = AdaptedRandError()
405        super().__init__(segmenter, metric)
406        self.init_kwargs = {"anisotropic": anisotropic, "min_seg_size": min_seg_size,
407                            "dt_threshold": dt_threshold, "sigma_seeds": sigma_seeds}
408
409
410class MWSIOUMetric(BaseInstanceSegmentationMetric):
411    """Intersection over union metric based on a mutex watershed computed from affinity predictions.
412
413    This class can be used as validation metric when training a network for instance segmentation.
414
415    Args:
416        offsets: The offsets corresponding to the affinity channels.
417        min_seg_size: The minimal segment size.
418        iou_threshold: The threshold for the intersection over union value.
419        strides: The strides for the mutex watershed.
420    """
421    def __init__(
422        self,
423        offsets: List[List[int]],
424        min_seg_size: int,
425        iou_threshold: float = 0.5,
426        strides: Optional[List[int]] = None
427    ):
428        segmenter = MWS(offsets, with_background=True, min_seg_size=min_seg_size, strides=strides)
429        metric = IOUError(iou_threshold)
430        super().__init__(segmenter, metric)
431        self.init_kwargs = {"offsets": offsets, "min_seg_size": min_seg_size,
432                            "iou_threshold": iou_threshold, "strides": strides}
433
434
435class MWSSBDMetric(BaseInstanceSegmentationMetric):
436    """Symmetric best dice score metric based on a mutex watershed computed from affinity predictions.
437
438    This class can be used as validation metric when training a network for instance segmentation.
439
440    Args:
441        offsets: The offsets corresponding to the affinity channels.
442        min_seg_size: The minimal segment size.
443        strides: The strides for the mutex watershed.
444    """
445    def __init__(self, offsets: List[List[int]], min_seg_size: int, strides: Optional[List[int]] = None):
446        segmenter = MWS(offsets, with_background=True, min_seg_size=min_seg_size, strides=strides)
447        metric = SymmetricBestDice()
448        super().__init__(segmenter, metric)
449        self.init_kwargs = {"offsets": offsets, "min_seg_size": min_seg_size, "strides": strides}
450
451
452class MWSVOIMetric(BaseInstanceSegmentationMetric):
453    """Variation of information metric based on a mutex watershed computed from affinity predictions.
454
455    This class can be used as validation metric when training a network for instance segmentation.
456
457    Args:
458        offsets: The offsets corresponding to the affinity channels.
459        min_seg_size: The minimal segment size.
460        strides: The strides for the mutex watershed.
461    """
462    def __init__(self, offsets: List[List[int]], min_seg_size: int, strides: Optional[List[int]] = None):
463        segmenter = MWS(offsets, with_background=False, min_seg_size=min_seg_size, strides=strides)
464        metric = VariationOfInformation()
465        super().__init__(segmenter, metric)
466        self.init_kwargs = {"offsets": offsets, "min_seg_size": min_seg_size, "strides": strides}
467
468
469class MWSRandMetric(BaseInstanceSegmentationMetric):
470    """Rand index metric based on a mutex watershed computed from affinity predictions.
471
472    This class can be used as validation metric when training a network for instance segmentation.
473
474    Args:
475        offsets: The offsets corresponding to the affinity channels.
476        min_seg_size: The minimal segment size.
477        strides: The strides for the mutex watershed.
478    """
479    def __init__(self, offsets: List[List[int]], min_seg_size: int, strides: Optional[List[int]] = None):
480        segmenter = MWS(offsets, with_background=False, min_seg_size=min_seg_size, strides=strides)
481        metric = AdaptedRandError()
482        super().__init__(segmenter, metric)
483        self.init_kwargs = {"offsets": offsets, "min_seg_size": min_seg_size, "strides": strides}