def lifted_problem_from_segmentation(rag, watershed, input_segmentation, overlap_threshold, graph_depth, same_segment_cost, different_segment_cost, mode='all', n_threads=None): """ Compute lifted problem from segmentation by mapping segments to watershed superpixels. Arguments: rag [RegionAdjacencyGraph] - the region adjacency graph watershed [np.ndarray] - the watershed over segmentation input_segmentation [np.ndarray] - Segmentation used to determine node attribution. overlap_threshold [float] - minimal overlap to assign a segment id to node graph_depth [int] - maximal graph depth up to which lifted edges will be included same_segment_cost [float] - costs for edges between nodes with same segment id attribution different_segment_cost [float] - costs for edges between nodes with different segment id attribution mode [str] - mode for insertion of lifted edges. Can be "all" - lifted edges will be inserted in between all nodes with attribution "different" - lifted edges will only be inserted in between nodes attributed to different classes "same" - lifted edges will only be inserted in between nodes attribted to the same class (default: "different") n_threads [int] - number of threads used for the calculation (default: None) """ n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads assert input_segmentation.shape == watershed.shape # compute the overlaps ovlp_comp = ngt.overlap(watershed, input_segmentation) ws_ids = np.unique(watershed) n_labels = ws_ids[-1] + 1 assert n_labels == rag.numberOfNodes, "%i, %i" % (n_labels, rag.numberOfNodes) # initialise the arrays for node labels, to be # dense in the watershed id space (even if some ws-ids are not present) node_labels = np.zeros(n_labels, dtype='uint64') # extract the overlap values and node labels from the overlap # computation results overlaps = [ ovlp_comp.overlapArraysNormalized(ws_id, sorted=False) for ws_id in ws_ids ] node_label_vals = np.array([ovlp[0][0] for ovlp in overlaps]) overlap_values = np.array([ovlp[1][0] for ovlp in overlaps]) node_label_vals[overlap_values < overlap_threshold] = 0 assert len(node_label_vals) == len(ws_ids) node_labels[ws_ids] = node_label_vals # find all lifted edges up to the graph depth between mapped nodes # NOTE we need to convert to the different graph type for now, but # it would be nice to support all nifty graphs at some type uv_ids = rag.uvIds() g_temp = ndist.Graph(uv_ids) lifted_uvs = ndist.liftedNeighborhoodFromNodeLabels( g_temp, node_labels, graph_depth, mode=mode, numberOfThreads=n_threads, ignoreLabel=0) # make sure that the lifted uv ids are in range of the node labels assert lifted_uvs.max() < rag.numberOfNodes, "%i, %i" % (int( lifted_uvs.max()), rag.numberOfNodes) lifted_labels = node_labels[lifted_uvs] lifted_costs = np.zeros_like(lifted_labels, dtype='float32') same_mask = lifted_labels[:, 0] == lifted_labels[:, 1] lifted_costs[same_mask] = same_segment_cost lifted_costs[~same_mask] = different_segment_cost return lifted_uvs, lifted_costs
def lifted_problem_from_probabilities(rag, watershed, input_maps, assignment_threshold, graph_depth, feats_to_costs=feats_to_costs_default, mode='different', n_threads=None): """ Compute lifted problem from probability maps by mapping them to superpixels. Arguments: rag [RegionAdjacencyGraph] - the region adjacency graph watershed [np.ndarray] - the watershed over segmentation input_maps [list[np.ndarray]] - list of probability maps. Each map must have the same shape as the watersheds and each map is treated as the probability to correspond to a different class. assignment_threshold [float] - minimal expression level to assign a class to a graph node (= watershed segment) graph_depth [int] - maximal graph depth up to which lifted edges will be included feats_to_costs [callable] - function to calculate the lifted costs from the class assignment probabilities. This becomes as inputs 'lifted_labels', which stores the two classes assigned to a lifted edge, and `lifted_features`, which stores the two assignment probabilities. (default: feats_to_costs_default). mode [str] - mode for insertion of lifted edges. Can be "all" - lifted edges will be inserted in between all nodes with attribution "different" - lifted edges will only be inserted in between nodes attributed to different classes "same" - lifted edges will only be inserted in between nodes attribted to the same class (default: "different") n_threads [int] - number of threads used for the calculation (default: None) """ n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads # validate inputs assert isinstance(input_maps, (list, tuple)) assert all(isinstance(inp, np.ndarray) for inp in input_maps) shape = watershed.shape assert all(inp.shape == shape for inp in input_maps) # map the probability maps to superpixels - we only map to superpixels which # have a larger mean expression than `assignment_threshold` # TODO handle the dtype conversion for vigra gracefully somehow ... # think about supporting uint8 input and normalizing # TODO how do we handle cases where the same superpixel is mapped to # more than one class ? n_nodes = int(watershed.max()) + 1 node_labels = np.zeros(n_nodes, dtype='uint64') node_features = np.zeros(n_nodes, dtype='float32') # TODO we could allow for more features that could then be used for the cost estimation for class_id, inp in enumerate(input_maps): mean_prob = vigra.analysis.extractRegionFeatures(inp, watershed, features=['mean' ])['mean'] # we can in principle map multiple classes here, and right now will just override class_mask = mean_prob > assignment_threshold node_labels[class_mask] = class_id node_features[class_mask] = mean_prob[class_mask] # find all lifted edges up to the graph depth between mapped nodes # NOTE we need to convert to the different graph type for now, but # it would be nice to support all nifty graphs at some type uv_ids = rag.uvIds() g_temp = ndist.Graph(uv_ids) lifted_uvs = ndist.liftedNeighborhoodFromNodeLabels( g_temp, node_labels, graph_depth, mode=mode, numberOfThreads=n_threads, ignoreLabel=0) lifted_labels = node_labels[lifted_uvs] lifted_features = node_features[lifted_uvs] lifted_costs = feats_to_costs(lifted_labels, lifted_features) return lifted_uvs, lifted_costs