elf.segmentation.features

  1import multiprocessing
  2from concurrent import futures
  3from typing import Dict, List, Optional, Tuple
  4
  5import numpy as np
  6import vigra
  7import nifty
  8import nifty.graph.rag as nrag
  9import nifty.ground_truth as ngt
 10try:
 11    import nifty.distributed as ndist
 12except ImportError:
 13    ndist = None
 14
 15try:
 16    import fastfilters as ff
 17except ImportError:
 18    import vigra.filters as ff
 19
 20from tqdm import tqdm
 21from .multicut import transform_probabilities_to_costs
 22
 23
 24#
 25# Region Adjacency Graph and Features
 26#
 27
 28def compute_rag(segmentation: np.ndarray, n_labels: Optional[int] = None, n_threads: Optional[int] = None):
 29    """Compute region adjacency graph of segmentation.
 30
 31    Args:
 32        segmentation: The segmentation.
 33        n_labels: The number of labels in the segmentation. If None, will be computed from the data.
 34        n_threads: The number of threads used, set to cpu count by default.
 35
 36    Returns:
 37        The region adjacency graph.
 38    """
 39    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
 40    n_labels = int(segmentation.max()) + 1 if n_labels is None else n_labels
 41    rag = nrag.gridRag(segmentation, numberOfLabels=n_labels, numberOfThreads=n_threads)
 42    return rag
 43
 44
 45def compute_boundary_features(
 46    rag, boundary_map: np.ndarray, min_value: float = 0.0, max_value: float = 1.0, n_threads: Optional[int] = None
 47) -> np.ndarray:
 48    """Compute edge features from boundary map.
 49
 50    Args:
 51        rag: The region adjacency graph.
 52        boundary_map:The boundary map.
 53        min_value: The minimum value used in accumulation.
 54        max_value: The maximum value used in accumulation.
 55        n_threads: The number of threads used, set to cpu count by default.
 56
 57    Returns:
 58        The edge features.
 59    """
 60    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
 61    if tuple(rag.shape) != boundary_map.shape:
 62        raise ValueError("Incompatible shapes: %s, %s" % (str(rag.shape), str(boundary_map.shape)))
 63    features = nrag.accumulateEdgeStandartFeatures(
 64        rag, boundary_map, min_value, max_value, numberOfThreads=n_threads
 65    )
 66    return features
 67
 68
 69def compute_affinity_features(
 70    rag,
 71    affinity_map: np.ndarray,
 72    offsets: List[List[int]],
 73    min_value: float = 0.0,
 74    max_value: float = 1.0,
 75    n_threads: Optional[int] = None
 76) -> np.ndarray:
 77    """Compute edge features from affinity map.
 78
 79    Args:
 80        rag: The region adjacency graph.
 81        affinity_map: The affinity map.
 82        min_value: The minimum value used in accumulation.
 83        max_value: The maximum value used in accumulation.
 84        n_threads: The umber of threads used, set to cpu count by default.
 85
 86    Returns:
 87        The edge features.
 88    """
 89    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
 90    if tuple(rag.shape) != affinity_map.shape[1:]:
 91        raise ValueError("Incompatible shapes: %s, %s" % (str(rag.shape), str(affinity_map.shape[1:])))
 92    if len(offsets) != affinity_map.shape[0]:
 93        raise ValueError("Incompatible number of channels and offsets: %i, %i" % (len(offsets),
 94                                                                                  affinity_map.shape[0]))
 95    features = nrag.accumulateAffinityStandartFeatures(
 96        rag, affinity_map, offsets, min_value, max_value, numberOfThreads=n_threads
 97    )
 98    return features
 99
100
101def compute_boundary_mean_and_length(rag, input_: np.ndarray, n_threads: Optional[int] = None) -> np.ndarray:
102    """Compute mean value and length of boundaries.
103
104    Args:
105        rag: The region adjacency graph.
106        input_: The input map.
107        n_threads: The number of threads used, set to cpu count by default.
108
109    Returns:
110        The edge features.
111    """
112    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
113    if tuple(rag.shape) != input_.shape:
114        raise ValueError("Incompatible shapes: %s, %s" % (str(rag.shape), str(input_.shape)))
115    features = nrag.accumulateEdgeMeanAndLength(rag, input_, numberOfThreads=n_threads)
116    return features
117
118
119# TODO generalize and move to elf.features.parallel
120def _filter_2d(input_, filter_name, sigma, n_threads):
121    filter_fu = getattr(ff, filter_name)
122
123    def _fz(inp):
124        response = filter_fu(inp, sigma)
125        # we add a channel last axis for 2d filter responses
126        if response.ndim == 2:
127            response = response[None, ..., None]
128        elif response.ndim == 3:
129            response = response[None]
130        else:
131            raise RuntimeError("Invalid filter response")
132        return response
133
134    with futures.ThreadPoolExecutor(n_threads) as tp:
135        tasks = [tp.submit(_fz, input_[z]) for z in range(input_.shape[0])]
136        response = [t.result() for t in tasks]
137
138    response = np.concatenate(response, axis=0)
139    return response
140
141
142def compute_boundary_features_with_filters(
143    rag,
144    input_: np.ndarray,
145    apply_2d: bool = False,
146    n_threads: Optional[int] = None,
147    filters: Dict[str, List[float]] = {"gaussianSmoothing": [1.6, 4.2, 8.3],
148                                       "laplacianOfGaussian": [1.6, 4.2, 8.3],
149                                       "hessianOfGaussianEigenvalues": [1.6, 4.2, 8.3]}
150) -> np.ndarray:
151    """Compute boundary features accumulated over filter responses on input.
152
153    Args:
154        rag: The region adjacency graph.
155        input_: The input data.
156        apply_2d: Whether to apply the filters in 2d for 3d input data.
157        n_threads: The number of threads.
158        filters: The filters to apply, expects a dictionary mapping filter names to sigma values.
159
160    Returns:
161        The edge filters.
162    """
163    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
164    features = []
165
166    # apply 2d: we compute filters and derived features in parallel per filter
167    if apply_2d:
168
169        def _compute_2d(filter_name, sigma):
170            response = _filter_2d(input_, filter_name, sigma, n_threads)
171            assert response.ndim == 4
172            n_channels = response.shape[-1]
173            features = []
174            for chan in range(n_channels):
175                chan_data = response[..., chan]
176                feats = compute_boundary_features(rag, chan_data,
177                                                  chan_data.min(), chan_data.max(), n_threads)
178                features.append(feats)
179
180            features = np.concatenate(features, axis=1)
181            assert len(features) == rag.numberOfEdges
182            return features
183
184        features = [_compute_2d(filter_name, sigma)
185                    for filter_name, sigmas in filters.items() for sigma in sigmas]
186
187    # apply 3d: we parallelize over the whole filter + feature computation
188    # this can be very memory intensive, and it would be better to parallelize inside
189    # of the loop, but 3d parallel filters in elf.parallel.filters are not working properly yet
190    else:
191
192        def _compute_3d(filter_name, sigma):
193            filter_fu = getattr(ff, filter_name)
194            response = filter_fu(input_, sigma)
195            if response.ndim == 3:
196                response = response[..., None]
197
198            n_channels = response.shape[-1]
199            features = []
200
201            for chan in range(n_channels):
202                chan_data = response[..., chan]
203                feats = compute_boundary_features(rag, chan_data,
204                                                  chan_data.min(), chan_data.max(),
205                                                  n_threads=1)
206                features.append(feats)
207            features = np.concatenate(features, axis=1)
208            assert len(features) == rag.numberOfEdges, f"{len(features), {rag.numberOfEdges}}"
209            return features
210
211        with futures.ThreadPoolExecutor(n_threads) as tp:
212            tasks = [tp.submit(_compute_3d, filter_name, sigma)
213                     for filter_name, sigmas in filters.items() for sigma in sigmas]
214            features = [t.result() for t in tasks]
215
216    features = np.concatenate(features, axis=1)
217    assert len(features) == rag.numberOfEdges
218    return features
219
220
221def compute_region_features(
222    uv_ids: np.ndarray,
223    input_map: np.ndarray,
224    segmentation: np.ndarray,
225    n_threads: Optional[int] = None
226) -> np.ndarray:
227    """Compute edge features from an input map accumulated over segmentation and mapped to edges.
228
229    Args:
230        uv_ids: The edge uv ids.
231        input_: The input data.
232        segmentation: The segmentation.
233        n_threads: The number of threads used, set to cpu count by default.
234
235    Returns:
236        The edge features.
237    """
238    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
239
240    # compute the node features
241    stat_feature_names = ["Count", "Kurtosis", "Maximum", "Minimum", "Quantiles",
242                          "RegionRadii", "Skewness", "Sum", "Variance"]
243    coord_feature_names = ["Weighted<RegionCenter>", "RegionCenter"]
244    feature_names = stat_feature_names + coord_feature_names
245    node_features = vigra.analysis.extractRegionFeatures(input_map, segmentation,
246                                                         features=feature_names)
247
248    # get the image statistics based features, that are combined via [min, max, sum, absdiff]
249    stat_features = [node_features[fname] for fname in stat_feature_names]
250    stat_features = np.concatenate([feat[:, None] if feat.ndim == 1 else feat
251                                    for feat in stat_features], axis=1)
252
253    # get the coordinate based features, that are combined via euclidean distance
254    coord_features = [node_features[fname] for fname in coord_feature_names]
255    coord_features = np.concatenate([feat[:, None] if feat.ndim == 1 else feat
256                                     for feat in coord_features], axis=1)
257
258    u, v = uv_ids[:, 0], uv_ids[:, 1]
259
260    # combine the stat features for all edges
261    feats_u, feats_v = stat_features[u], stat_features[v]
262    features = [np.minimum(feats_u, feats_v), np.maximum(feats_u, feats_v),
263                np.abs(feats_u - feats_v), feats_u + feats_v]
264
265    # combine the coord features for all edges
266    feats_u, feats_v = coord_features[u], coord_features[v]
267    features.append((feats_u - feats_v) ** 2)
268
269    features = np.nan_to_num(np.concatenate(features, axis=1))
270    assert len(features) == len(uv_ids)
271    return features
272
273
274#
275# Grid Graph and Features
276#
277
278def compute_grid_graph(shape: Tuple[int, ...]):
279    """Compute grid graph for the given shape.
280
281    Args:
282        shape: The shape of the data.
283
284    Returns:
285        The grid graph.
286    """
287    grid_graph = nifty.graph.undirectedGridGraph(shape)
288    return grid_graph
289
290
291def compute_grid_graph_image_features(
292    grid_graph,
293    image: np.ndarray,
294    mode: str,
295    offsets: Optional[List[List[int]]] = None,
296    strides: Optional[List[int]] = None,
297    randomize_strides: bool = False,
298) -> Tuple[np.ndarray, np.ndarray]:
299    """Compute edge features for image for the given grid_graph.
300
301    Args:
302        grid_graph: The grid graph
303        image: The image, from which the features will be derived.
304        mode: Feature accumulation method.
305        offsets: The offsets, which correspond to the affinity channels.
306            If none are given, the affinites for the nearest neighbor transitions are used.
307        strides: The strides used to subsample edges that are computed from offsets.
308        randomize_strides: Whether to subsample randomly instead of using regular strides.
309
310    Returns:
311        The uv ids of the edges.
312        The edge features.
313    """
314    gndim = len(grid_graph.shape)
315
316    if image.ndim == gndim:
317        if offsets is not None:
318            raise NotImplementedError
319        modes = ("l1", "l2", "min", "max", "sum", "prod", "interpixel")
320        if mode not in modes:
321            raise ValueError(f"Invalid feature mode {mode}, expect one of {modes}")
322        features = grid_graph.imageToEdgeMap(image, mode)
323        edges = grid_graph.uvIds()
324
325    elif image.ndim == gndim + 1:
326        modes = ("l1", "l2", "cosine")
327        if mode not in modes:
328            raise ValueError(f"Invalid feature mode {mode}, expect one of {modes}")
329
330        if offsets is None:
331            features = grid_graph.imageWithChannelsToEdgeMap(image, mode)
332            edges = grid_graph.uvIds()
333        else:
334            (n_edges,
335             edges,
336             features) = grid_graph.imageWithChannelsToEdgeMapWithOffsets(image, mode,
337                                                                          offsets=offsets,
338                                                                          strides=strides,
339                                                                          randomize_strides=randomize_strides)
340            edges, features = edges[:n_edges], features[:n_edges]
341
342    else:
343        msg = f"Invalid image dimension {image.ndim}, expect one of {gndim} or {gndim + 1}"
344        raise ValueError(msg)
345
346    return edges, features
347
348
349def compute_grid_graph_affinity_features(
350    grid_graph,
351    affinities: np.ndarray,
352    offsets: Optional[List[List[int]]] = None,
353    strides: Optional[List[int]] = None,
354    mask: Optional[np.ndarray] = None,
355    randomize_strides: bool = False,
356) -> Tuple[np.ndarray, np.ndarray]:
357    """Compute edge features from affinities for the given grid graph.
358
359    Args:
360        grid_graph: The grid graph
361        affinities: The affinity map.
362        offsets: The offsets, which correspond to the affinity channels.
363            If none are given, the affinites for the nearest neighbor transitions are used.
364        strides: The strides used to subsample edges that are computed from offsets.
365        mask: Mask to exclude from the edge and feature computation.
366        randomize_strides: Whether to subsample randomly instead of using regular strides.
367
368    Returns:
369        The uv ids of the edges.
370        The edge features.
371    """
372    gndim = len(grid_graph.shape)
373    if affinities.ndim != gndim + 1:
374        raise ValueError
375
376    if offsets is None:
377        assert affinities.shape[0] == gndim
378        assert strides is None
379        assert mask is None
380        features = grid_graph.affinitiesToEdgeMap(affinities)
381        edges = grid_graph.uvIds()
382    elif mask is not None:
383        assert strides is None and not randomize_strides, "Strides and mask cannot be used at the same time"
384        n_edges, edges, features = grid_graph.affinitiesToEdgeMapWithMask(affinities,
385                                                                          offsets=offsets,
386                                                                          mask=mask)
387        edges, features = edges[:n_edges], features[:n_edges]
388    else:
389        n_edges, edges, features = grid_graph.affinitiesToEdgeMapWithOffsets(affinities,
390                                                                             offsets=offsets,
391                                                                             strides=strides,
392                                                                             randomize_strides=randomize_strides)
393        edges, features = edges[:n_edges], features[:n_edges]
394
395    return edges, features
396
397
398def apply_mask_to_grid_graph_weights(
399    grid_graph,
400    mask: np.ndarray,
401    weights: np.ndarray,
402    masked_edge_weight: float = 0.0,
403    transition_edge_weight: float = 1.0,
404) -> np.ndarray:
405    """Mask edges in grid graph.
406
407    Set the weights derived from a grid graph to a fixed value, for edges that connect masked nodes
408    and edges that connect masked and unmasked nodes.
409
410    Args:
411        grid_graph: The grid graph.
412        mask: The binary mask, foreground (=non-masked) is True.
413        weights: The edge weights.
414        masked_edge_weight: The value for edges that connect two masked nodes.
415        transition_edge_weight: The value for edges that connect a masked with a non-masked node.
416
417    Returns:
418        The masked edge weights.
419    """
420    assert np.dtype(mask.dtype) == np.dtype("bool")
421    node_ids = grid_graph.projectNodeIdsToPixels()
422    assert node_ids.shape == mask.shape == tuple(grid_graph.shape), \
423        f"{node_ids.shape}, {mask.shape}, {grid_graph.shape}"
424    masked_ids = node_ids[~mask]
425
426    edges = grid_graph.uvIds()
427    assert len(edges) == len(weights)
428    edge_state = np.isin(edges, masked_ids).sum(axis=1)
429    masked_edges = edge_state == 2
430    transition_edges = edge_state == 1
431    weights[masked_edges] = masked_edge_weight
432    weights[transition_edges] = transition_edge_weight
433    return weights
434
435
436def apply_mask_to_grid_graph_edges_and_weights(
437    grid_graph, mask: np.ndarray, edges: np.ndarray, weights: np.ndarray, transition_edge_weight: float = 1.0
438) -> Tuple[np.ndarray, np.ndarray]:
439    """Remove uv ids that connect masked nodes and set weights that connect masked to non-masked nodes to a fixed value.
440
441    Args:
442        grid_graph: The grid graph.
443        mask: The binary mask, foreground (=non-masked) is True.
444        edges: The edges (uv-ids).
445        weights: The edge weights.
446        transition_edge_weight: The value for edges that connect a masked with a non-masked node.
447
448    Returns:
449        The edge uv-ids.
450        The edge weights.
451    """
452    assert np.dtype(mask.dtype) == np.dtype("bool")
453    node_ids = grid_graph.projectNodeIdsToPixels()
454    assert node_ids.shape == mask.shape == tuple(grid_graph.shape), \
455        f"{node_ids.shape}, {mask.shape}, {grid_graph.shape}"
456    masked_ids = node_ids[~mask]
457
458    edge_state = np.isin(edges, masked_ids).sum(axis=1)
459    keep_edges = edge_state != 2
460
461    edges, weights, edge_state = edges[keep_edges], weights[keep_edges], edge_state[keep_edges]
462    transition_edges = edge_state == 1
463    weights[transition_edges] = transition_edge_weight
464
465    return edges, weights
466
467
468#
469# Lifted Features
470#
471
472def lifted_edges_from_graph_neighborhood(graph, max_graph_distance):
473    """@private
474    """
475    if max_graph_distance < 2:
476        raise ValueError(f"Graph distance must be greater equal 2, got {max_graph_distance}")
477    if isinstance(graph, nifty.graph.UndirectedGraph):
478        objective = nifty.graph.opt.lifted_multicut.liftedMulticutObjective(graph)
479    else:
480        graph_ = nifty.graph.undirectedGraph(graph.numberOfNodes)
481        graph_.insertEdges(graph.uvIds())
482        objective = nifty.graph.opt.lifted_multicut.liftedMulticutObjective(graph_)
483    objective.insertLiftedEdgesBfs(max_graph_distance)
484    lifted_uvs = objective.liftedUvIds()
485    return lifted_uvs
486
487
488def feats_to_costs_default(lifted_labels, lifted_features):
489    """@private
490    """
491    # we assume that we only have different classes for a given lifted
492    # edge here (mode = "different") and then set all edges to be repulsive
493
494    # the higher the class probability, the more repulsive the edges should be,
495    # so we just multiply both probabilities
496    lifted_costs = lifted_features[:, 0] * lifted_features[:, 1]
497    lifted_costs = transform_probabilities_to_costs(lifted_costs)
498    return lifted_costs
499
500
501def lifted_problem_from_probabilities(
502    rag,
503    watershed: np.ndarray,
504    input_maps: List[np.ndarray],
505    assignment_threshold: float,
506    graph_depth: int,
507    feats_to_costs: callable = feats_to_costs_default,
508    mode: str = "different",
509    n_threads: Optional[int] = None,
510) -> Tuple[np.ndarray, np.ndarray]:
511    """Compute lifted problem from probability maps by mapping them to superpixels.
512
513    Example: compute a lifted problem from two attributions (axon, dendrite) that induce
514    repulsive edges between different attributions. The construction of lifted eges and
515    features can be customized using the `feats_to_costs` and `mode` arguments.
516    ```
517    lifted_uvs, lifted_costs = lifted_problem_from_probabilties(
518       rag, superpixels,
519       input_maps=[
520         axon_probabilities,  # probabilty map for axon attribution
521         dendrite_probabilities  # probability map for dendrite attributtion
522       ],
523       assignment_threshold=0.6,  # probability threshold to assign superpixels to a class
524       graph_depth=10,  # the max. graph depth along which lifted edges are introduced
525    )
526    ```
527
528    Args:
529        rag: The region adjacency graph.
530        watershed: The watershed over-segmentation.
531        input_maps: List of probability maps. Each map must have the same shape as the watersheds
532            and each map is treated as the probability to correspond to a different class.
533        assignment_threshold: Minimal expression level to assign a class to a graph node (= watershed segment).
534        graph_depth: Maximal graph depth up to which lifted edges will be included.
535        feats_to_costs: Function to calculate the lifted costs from the class assignment probabilities.
536            The input to the function are `lifted_labels`, which stores the two classes assigned to a lifted edge,
537            and `lifted_features`, which stores the two assignment probabilities.
538        mode: The mode for insertion of lifted edges. One of:
539            "all" - lifted edges will be inserted in between all nodes with attribution.
540            "different" - lifted edges will only be inserted in between nodes attributed to different classes.
541            "same" - lifted edges will only be inserted in between nodes attribted to the same class.
542        n_threads: The number of threads used for the calculation.
543
544    Returns:
545        The lifted uv ids (= superpixel ids connected by the lifted edge).
546        The lifted costs (= cost associated with each lifted edge).
547    """
548    assert ndist is not None, "Need nifty.distributed package"
549
550    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
551    # validate inputs
552    assert isinstance(input_maps, (list, tuple))
553    assert all(isinstance(inp, np.ndarray) for inp in input_maps)
554    shape = watershed.shape
555    assert all(inp.shape == shape for inp in input_maps)
556
557    # map the probability maps to superpixels - we only map to superpixels which
558    # have a larger mean expression than `assignment_threshold`
559
560    # TODO handle the dtype conversion for vigra gracefully somehow ...
561    # think about supporting uint8 input and normalizing
562
563    # TODO how do we handle cases where the same superpixel is mapped to
564    # more than one class ?
565
566    n_nodes = int(watershed.max()) + 1
567    node_labels = np.zeros(n_nodes, dtype="uint64")
568    node_features = np.zeros(n_nodes, dtype="float32")
569    # TODO we could allow for more features that could then be used for the cost estimation
570    for class_id, inp in enumerate(input_maps):
571        mean_prob = vigra.analysis.extractRegionFeatures(inp, watershed, features=["mean"])["mean"]
572        # we can in principle map multiple classes here, and right now will just override
573        class_mask = mean_prob > assignment_threshold
574        node_labels[class_mask] = class_id
575        node_features[class_mask] = mean_prob[class_mask]
576
577    # find all lifted edges up to the graph depth between mapped nodes
578    # NOTE we need to convert to the different graph type for now, but
579    # it would be nice to support all nifty graphs at some type
580    uv_ids = rag.uvIds()
581    g_temp = ndist.Graph(uv_ids)
582
583    lifted_uvs = ndist.liftedNeighborhoodFromNodeLabels(g_temp, node_labels, graph_depth, mode=mode,
584                                                        numberOfThreads=n_threads, ignoreLabel=0)
585    lifted_labels = node_labels[lifted_uvs]
586    lifted_features = node_features[lifted_uvs]
587
588    lifted_costs = feats_to_costs(lifted_labels, lifted_features)
589    return lifted_uvs, lifted_costs
590
591
592# TODO support setting costs proportional to overlaps
593def lifted_problem_from_segmentation(
594    rag,
595    watershed: np.ndarray,
596    input_segmentation: np.ndarray,
597    overlap_threshold: float,
598    graph_depth: int,
599    same_segment_cost: float,
600    different_segment_cost: float,
601    mode: str = "all",
602    n_threads: Optional[int] = None,
603) -> Tuple[np.ndarray, np.ndarray]:
604    """Compute lifted problem from segmentation by mapping segments to superpixels.
605
606    Args:
607        rag: The region adjacency graph.
608        watershed: The watershed over-segmentation.
609        input_segmentation: The segmentation used to determine node attribution.
610        overlap_threshold: The minimal overlap to assign a segment id to node.
611        graph_depth: The maximal graph depth up to which lifted edges will be included
612        same_segment_cost: The cost for edges between nodes with same segment id attribution.
613        different_segment_cost: The cost for edges between nodes with different segment id attribution.
614        mode: The mode for insertion of lifted edges. One of:
615            "all" - lifted edges will be inserted in between all nodes with attribution.
616            "different" - lifted edges will only be inserted in between nodes attributed to different classes.
617            "same" - lifted edges will only be inserted in between nodes attribted to the same class.
618        n_threads: The number of threads used for the calculation.
619
620    Returns:
621        The lifted uv ids (= superpixel ids connected by the lifted edge).
622        The lifted costs (= cost associated with each lifted edge).
623    """
624    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
625    assert input_segmentation.shape == watershed.shape
626
627    # compute the overlaps
628    ovlp_comp = ngt.overlap(watershed, input_segmentation)
629    ws_ids = np.unique(watershed)
630    n_labels = int(ws_ids[-1]) + 1
631    assert n_labels == rag.numberOfNodes, "%i, %i" % (n_labels, rag.numberOfNodes)
632
633    # initialise the arrays for node labels, to be
634    # dense in the watershed id space (even if some ws-ids are not present)
635    node_labels = np.zeros(n_labels, dtype="uint64")
636
637    # extract the overlap values and node labels from the overlap
638    # computation results
639    overlaps = [ovlp_comp.overlapArraysNormalized(ws_id, sorted=False)
640                for ws_id in ws_ids]
641    node_label_vals = np.array([ovlp[0][0] for ovlp in overlaps])
642    overlap_values = np.array([ovlp[1][0] for ovlp in overlaps])
643    node_label_vals[overlap_values < overlap_threshold] = 0
644    assert len(node_label_vals) == len(ws_ids)
645    node_labels[ws_ids] = node_label_vals
646
647    # find all lifted edges up to the graph depth between mapped nodes
648    # NOTE we need to convert to the different graph type for now, but
649    # it would be nice to support all nifty graphs at some type
650    uv_ids = rag.uvIds()
651    g_temp = ndist.Graph(uv_ids)
652
653    lifted_uvs = ndist.liftedNeighborhoodFromNodeLabels(g_temp, node_labels, graph_depth, mode=mode,
654                                                        numberOfThreads=n_threads, ignoreLabel=0)
655    # make sure that the lifted uv ids are in range of the node labels
656    assert lifted_uvs.max() < rag.numberOfNodes, "%i, %i" % (int(lifted_uvs.max()),
657                                                             rag.numberOfNodes)
658    lifted_labels = node_labels[lifted_uvs]
659    lifted_costs = np.zeros(len(lifted_labels), dtype="float64")
660
661    same_mask = lifted_labels[:, 0] == lifted_labels[:, 1]
662    lifted_costs[same_mask] = same_segment_cost
663    lifted_costs[~same_mask] = different_segment_cost
664
665    return lifted_uvs, lifted_costs
666
667
668#
669# Misc
670#
671
672def get_stitch_edges(
673    rag,
674    seg: np.ndarray,
675    block_shape: Tuple[int, ...],
676    n_threads: Optional[int] = None,
677    verbose: bool = False
678) -> np.ndarray:
679    """Get the edges between blocks.
680
681    Args:
682        rag: The region adjacency graph.
683        seg: The segmentation underlying the rag.
684        block_shape: The shape of the blocking.
685        n_threads: The number of threads used for the calculation.
686        verbose: Whether to be verbose.
687
688    Returns:
689        The edge mask indicating edges between blocks.
690    """
691    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
692    ndim = seg.ndim
693    blocking = nifty.tools.blocking([0] * ndim, seg.shape, block_shape)
694
695    def find_stitch_edges(block_id):
696        stitch_edges = []
697        block = blocking.getBlock(block_id)
698        for axis in range(ndim):
699            if blocking.getNeighborId(block_id, axis, True) == -1:
700                continue
701            face_a = tuple(
702                beg if d == axis else slice(beg, end)
703                for d, beg, end in zip(range(ndim), block.begin, block.end)
704            )
705            face_b = tuple(
706                beg - 1 if d == axis else slice(beg, end)
707                for d, beg, end in zip(range(ndim), block.begin, block.end)
708            )
709
710            labels_a = seg[face_a].ravel()
711            labels_b = seg[face_b].ravel()
712
713            uv_ids = np.concatenate(
714                [labels_a[:, None], labels_b[:, None]],
715                axis=1
716            )
717            uv_ids = np.unique(uv_ids, axis=0)
718
719            edge_ids = rag.findEdges(uv_ids)
720            edge_ids = edge_ids[edge_ids != -1]
721            stitch_edges.append(edge_ids)
722
723        if stitch_edges:
724            stitch_edges = np.concatenate(stitch_edges)
725            stitch_edges = np.unique(stitch_edges)
726        else:
727            stitch_edges = None
728        return stitch_edges
729
730    with futures.ThreadPoolExecutor(n_threads) as tp:
731        if verbose:
732            stitch_edges = list(tqdm(
733                tp.map(find_stitch_edges, range(blocking.numberOfBlocks)),
734                total=blocking.numberOfBlocks
735            ))
736        else:
737            stitch_edges = tp.map(find_stitch_edges, range(blocking.numberOfBlocks))
738
739    stitch_edges = np.concatenate([st for st in stitch_edges if st is not None])
740    stitch_edges = np.unique(stitch_edges)
741    full_edges = np.zeros(rag.numberOfEdges, dtype="bool")
742    full_edges[stitch_edges] = 1
743    return full_edges
744
745
746def project_node_labels_to_pixels(rag, node_labels: np.ndarray, n_threads: Optional[int] = None) -> np.ndarray:
747    """Project label values for graph nodes back to pixels to obtain segmentation.
748
749    Args:
750        rag: The region adjacency graph.
751        node_labels: The array with node labels.
752        n_threads: The number of threads used, set to cpu count by default.
753
754    Returns:
755        The segmentation.
756    """
757    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
758    if len(node_labels) != rag.numberOfNodes:
759        raise ValueError("Incompatible number of node labels: %i, %i" % (len(node_labels), rag.numberOfNodes))
760    seg = nrag.projectScalarNodeDataToPixels(rag, node_labels, numberOfThreads=n_threads)
761    return seg
762
763
764def compute_z_edge_mask(rag, watershed: np.ndarray) -> np.ndarray:
765    """Compute edge mask of in-between plane edges for flat superpixels.
766
767    Flat superpixels are volumetric superpixels that are independent across slices.
768    This function does not check wether the input watersheds are actually flat.
769
770    Args:
771        rag: The region adjacency graph.
772        watershed: The underlying watershed over-segmentation (superpixels).
773
774    Returns:
775        The edge mask indicating in-between slice edges.
776    """
777    node_z_coords = np.zeros(rag.numberOfNodes, dtype="uint32")
778    for z in range(watershed.shape[0]):
779        node_z_coords[watershed[z]] = z
780    uv_ids = rag.uvIds()
781    z_edge_mask = node_z_coords[uv_ids[:, 0]] != node_z_coords[uv_ids[:, 1]]
782    return z_edge_mask