def compute_state(affs, seg, offsets, n_attractive): # with affogato TODO debug this # FIXME the uv ids don't make sense! # grid_graph = compute_grid_graph(segmentation.shape) # uvs, weights, attractive = grid_graph.compute_state_for_segmentation(affs, segmentation, offsets, # n_attractive_channels=3, # ignore_label=False) # weights[np.logical_not(attractive)] *= -1 # state = (uvs, weights) # with nifty rag = nrag.gridRag(seg, numberOfLabels=int(seg.max() + 1), numberOfThreads=1) uv_ids = rag.uvIds() affs_attractive = affs[:n_attractive] # -2 corresponds to max value weights_attractive = nrag.accumulateAffinityStandartFeatures( rag, affs_attractive, offsets, numberOfThreads=1)[:, -2] affs_repulsive = np.require(affs[n_attractive:], requirements='C') weights_repulsive = nrag.accumulateAffinityStandartFeatures( rag, affs_repulsive, offsets, numberOfThreads=1)[:, -2] weights = weights_attractive repulsive = weights_repulsive > weights_attractive weights[repulsive] = -1 * weights_repulsive[repulsive] return uv_ids, weights
def compute_state(affs, seg, offsets, n_attractive): # with affogato TODO debug this # FIXME the uv ids don't make sense! # grid_graph = compute_grid_graph(segmentation.shape) # uvs, weights, attractive = grid_graph.compute_state_for_segmentation(affs, segmentation, offsets, # n_attractive_channels=3, # ignore_label=False) # weights[np.logical_not(attractive)] *= -1 # state = (uvs, weights) # with nifty rag = nrag.gridRag(seg, numberOfLabels=int(seg.max() + 1), numberOfThreads=1) uv_ids = rag.uvIds() affs_attractive = affs[:n_attractive] # -2 corresponds to max value weights_attractive = nrag.accumulateAffinityStandartFeatures(rag, affs_attractive, offsets, numberOfThreads=1)[:, -2] affs_repulsive = np.require(affs[n_attractive:], requirements='C') weights_repulsive = nrag.accumulateAffinityStandartFeatures(rag, affs_repulsive, offsets, numberOfThreads=1)[:, -2] weights = weights_attractive repulsive = weights_repulsive > weights_attractive weights[repulsive] = -1*weights_repulsive[repulsive] return uv_ids, weights
def compute_affinity_features(rag, affinity_map, offsets, min_value=0, max_value=1, n_threads=None): """ Compute edge features from affinity map. Arguments: rag [RegionAdjacencyGraph] - region adjacency graph boundary_map [np.ndarray] - boundary map. min_value [float] - minimum value used in accumulation (default: 0) max_value [float] - maximum value used in accumulation (default: 1) n_threads [int] - number of threads used, set to cpu count by default. (default: None) """ n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads if tuple(rag.shape) != affinity_map.shape[1:]: raise ValueError("Incompatible shapes: %s, %s" % (str(rag.shape), str(affinity_map.shape[1:]))) if len(offsets) != affinity_map.shape[0]: raise ValueError( "Incompatible number of channels and offsets: %i, %i" % (len(offsets), affinity_map.shape[0])) features = nrag.accumulateAffinityStandartFeatures( rag, affinity_map, offsets, min_value, max_value, numberOfThreads=n_threads) return features
def compute_mcrf_segments(ws, affs, n_labels, glia, offsets, rf, block_id): rag = nrag.gridRag(ws, numberOfLabels=int(n_labels), numberOfThreads=1) uv_ids = rag.uvIds() # we can have a single over-segment id for (small border) blocks # resulting in 0 edges if(uv_ids.shape[0] == 0): print("WARNING:, block", block_id, "contains only a single id, but is not masked") print("This is may be caused by an incorrect mask") return ws # TODO split features over different affinity ranges ? feats = np.concatenate([nrag.accumulateAffinityStandartFeatures(rag, affs, offsets, numberOfThreads=1), glia_features(rag, ws, glia)], axis=1) probs = rf.predict_proba(feats)[:, 1] mc = cseg.Multicut('kernighan-lin', weight_edges=False) costs = mc.probabilities_to_costs(probs) ignore_edges = (uv_ids == 0).any(axis=1) costs[ignore_edges] = 5 * costs.min() graph = nifty.graph.undirectedGraph(n_labels) graph.insertEdges(uv_ids) node_labels = mc(graph, costs) return nrag.projectScalarNodeDataToPixels(rag, node_labels)
def _mc_impl(ws, affs, offsets, weight_mulitcut_edges=False, weighting_exponent=1): ws = ws.astype('uint32') n_labels = int(ws.max()) + 1 rag = nrag.gridRag(ws, numberOfLabels=n_labels, numberOfThreads=8) uv_ids = rag.uvIds() # compute the features and get edge probabilities (from mean affinities) # and edge sizes features = nrag.accumulateAffinityStandartFeatures(rag, affs, offsets, numberOfThreads=8) probs = features[:, 0] sizes = features[:, -1] graph = nifty.graph.undirectedGraph(n_labels) graph.insertEdges(uv_ids) # compute multicut edge results node_labels, _ = multicut.run_mc( graph, probs, uv_ids, with_ignore_edges=True, edge_sizes=sizes if weight_mulitcut_edges else None, weighting_exponent=weighting_exponent) return node_labels
def compute_mc(ws, affs, offsets, n_labels, weight_mulitcut_edges, weighting_exponent): # compute the region adjacency graph rag = nrag.gridRag(ws, numberOfLabels=n_labels, numberOfThreads=1) uv_ids = rag.uvIds() if uv_ids.size == 0: return None, None, None # compute the features and get edge probabilities (from mean affinities) # and edge sizes features = nrag.accumulateAffinityStandartFeatures(rag, affs, offsets, numberOfThreads=1) probs = features[:, 0] sizes = features[:, -1] graph = nifty.graph.undirectedGraph(n_labels) graph.insertEdges(uv_ids) # compute multicut edge results _, merge_indicator = run_mc( graph, probs, uv_ids, with_ignore_edges=True, edge_sizes=sizes if weight_mulitcut_edges else None, weighting_exponent=weighting_exponent) return uv_ids, merge_indicator, sizes
def _compute_edge_probabilities(self, input_, fragments=None): assert input_.ndim == 4 assert input_.shape[0] == len( self.offsets), "%i, %i" % (input_.shape[0], len(self.offsets)) features = nrag.accumulateAffinityStandartFeatures( self.rag, input_, self.offsets) return features[:, self.stat_index]
def segment_block(ds_ws, ds_affs, blocking, block_id, offsets): # load the segmentation block = blocking.getBlock(block_id) bb = tuple(slice(beg, end) for beg, end in zip(block.begin, block.end)) ws = ds_ws[bb] # if this block only contains a single segment id (usually 0 = ignore label) continue ws_ids = np.unique(ws) if len(ws_ids) == 1: return None max_id = ws_ids[-1] # TODO should we do this ? # map to a consecutive segmentation to speed up graph computations # ws, max_id, mapping = vigra.analysis.relabelConsecutive(ws, keep_zeros=True, start_label=1) # load the affinities n_channels = len(offsets) bb_affs = (slice(0, n_channels), ) + bb affs = ds_affs[bb_affs] # convert affinities to float and invert them # to get boundary probabilities if affs.dtype == np.dtype('uint8'): affs = affs.astype('float32') / 255. affs = 1. - affs # compute the region adjacency graph n_labels = int(max_id) + 1 rag = nrag.gridRag(ws, numberOfLabels=n_labels, numberOfThreads=1) uv_ids = rag.uvIds() # compute the features and get edge probabilities (from mean affinities) # and edge sizes features = nrag.accumulateAffinityStandartFeatures(rag, affs, offsets, numberOfThreads=1) probs = features[:, 0] sizes = features[:, -1].astype('uint64') # compute multicut mc = cseg.Multicut('kernighan-lin', weight_edges=False) # transform probabilities to costs costs = mc.probabilities_to_costs(probs) # set edges connecting to 0 (= ignore label) to repulsive ignore_edges = (uv_ids == 0).any(axis=1) costs[ignore_edges] = -100 # solve the mc problem graph = nifty.graph.undirectedGraph(n_labels) graph.insertEdges(uv_ids) node_labels = mc(graph, costs) # get indicators for merge ! # and return uv-ids, edge indicators and edge sizes edge_indicator = ( node_labels[uv_ids[:, 0]] == node_labels[uv_ids[:, 1]]).astype('uint8') return uv_ids, edge_indicator, sizes
def separate_channel_features(rag, affs): offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1], [-2, 0, 0], [0, -3, 0], [0, 0, -3], [-3, 0, 0], [0, -9, 0], [0, 0, -9], [-4, 0, 0], [0, -27, 0], [0, 0, -27]] assert len(affs) == len(offsets) features = [] for channel_id in range(len(affs)): features.append( nrag.accumulateAffinityStandartFeatures( rag, affs[channel_id:channel_id + 1], [offsets[channel_id]])) features = np.concatenate(features, axis=1) return np.nan_to_num(features)
def preagglomerate(affs, ws, n_labels, thresh=.95): offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1], [-2, 0, 0], [0, -3, 0], [0, 0, -3], [-3, 0, 0], [0, -9, 0], [0, 0, -9], [-4, 0, 0], [0, -27, 0], [0, 0, -27]] rag = nrag.gridRag(ws, numberOfLabels=int(n_labels), numberOfThreads=1) probs = nrag.accumulateAffinityStandartFeatures(rag, affs, offsets, numberOfThreads=1)[:, 0] agglomerator = cseg.MalaClustering(thresh) g = nifty.graph.undirectedGraph(rag.numberOfNodes) g.insertEdges(rag.uvIds()) node_labels = agglomerator(g, probs) return nrag.projectScalarNodeDataToPixels(rag, node_labels), int( node_labels.max())
def compute_mc_nodes(ws, affs, n_labels, offsets, glia=None): rag = nrag.gridRag(ws, numberOfLabels=int(n_labels), numberOfThreads=1) probs = nrag.accumulateAffinityStandartFeatures(rag, affs, offsets, numberOfThreads=1)[:, 0] uv_ids = rag.uvIds() mc = cseg.Multicut('kernighan-lin', weight_edges=False) costs = mc.probabilities_to_costs(probs) ignore_edges = (uv_ids == 0).any(axis=1) costs[ignore_edges] = - 100 # run multicut graph = nifty.graph.undirectedGraph(n_labels) graph.insertEdges(uv_ids) node_labels = mc(graph, costs) # find the pairs of merged nodes from the multicut # node labeling return get_merged_nodes(uv_ids, node_labels)
def segment_mcrf(ws, affs, n_labels, offsets, rf, return_merged_nodes=False): rag = nrag.gridRag(ws, numberOfLabels=int(n_labels), numberOfThreads=1) feats = nrag.accumulateAffinityStandartFeatures(rag, affs, offsets, numberOfThreads=1) probs = rf.predict_proba(feats)[:, 1] uv_ids = rag.uvIds() mc = cseg.Multicut('kernighan-lin', weight_edges=False) costs = mc.probabilities_to_costs(probs) ignore_edges = (uv_ids == 0).any(axis=1) costs[ignore_edges] = 5 * costs.min() graph = nifty.graph.undirectedGraph(n_labels) graph.insertEdges(uv_ids) node_labels = mc(graph, costs) if return_merged_nodes: return get_merged_nodes(uv_ids, node_labels) else: return nrag.projectScalarNodeDataToPixels(rag, node_labels)
def compute_rag_feats(): path = '/home/papec/mnt/papec/Work/neurodata_hdd/cremi_warped/sampleA+.n5' print("Loading seg") seg = z5py.File(path)['segmentations/watershed'][:] print(seg.shape) print("Computing rag") rag = nrag.gridRag(seg, numberOfLabels=int(seg.max() + 1)) print("Loading affs") affs = z5py.File(path)['predictions/affs_glia'][0:12, :] offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1], [-2, 0, 0], [0, -3, 0], [0, 0, -3], [-3, 0, 0], [0, -9, 0], [0, 0, -9], [-4, 0, 0], [0, -27, 0], [0, 0, -27]] assert len(offsets) == len(affs) print("Computing feats") feats = nrag.accumulateAffinityStandartFeatures(rag, affs, offsets) f_save = z5py.File('./costs_tmp.n5') ds = f_save.create_dataset('feats', dtype='float32', shape=feats.shape, chunks=(feats.shape[0], 1), compression='gzip') ds[:] = feats
def compute_mcrf_nodes(ws, affs, n_labels, offsets, glia, rf): rag = nrag.gridRag(ws, numberOfLabels=int(n_labels), numberOfThreads=1) # TODO split features over different affinity ranges ? feats = np.concatenate([nrag.accumulateAffinityStandartFeatures(rag, affs, offsets, numberOfThreads=1), glia_features(rag, ws, glia)], axis=1) probs = rf.predict_proba(feats)[:, 1] uv_ids = rag.uvIds() mc = cseg.Multicut('kernighan-lin', weight_edges=False) costs = mc.probabilities_to_costs(probs) ignore_edges = (uv_ids == 0).any(axis=1) costs[ignore_edges] = -100 # run multicut graph = nifty.graph.undirectedGraph(n_labels) graph.insertEdges(uv_ids) node_labels = mc(graph, costs) # find the pairs of merged nodes from the multicut # node labeling return get_merged_nodes(uv_ids, node_labels)
def affinity_features(rag, affs): offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1], [-2, 0, 0], [0, -3, 0], [0, 0, -3], [-3, 0, 0], [0, -9, 0], [0, 0, -9], [-4, 0, 0], [0, -27, 0], [0, 0, -27]] features = nrag.accumulateAffinityStandartFeatures(rag, affs, offsets) return np.nan_to_num(features)
def compute_lmc_learned(ws, affs, glia, offsets, n_labels, lifted_rf, lifted_nh, weight_mulitcut_edges, weighting_exponent): # compute the region adjacency graph rag = nrag.gridRag(ws, numberOfLabels=n_labels, numberOfThreads=1) uv_ids = rag.uvIds() # compute the features and get edge probabilities (from mean affinities) # and edge sizes features = nrag.accumulateAffinityStandartFeatures(rag, affs, offsets, numberOfThreads=1) local_probs = features[:, 0] sizes = features[:, -1].astype('uint64') # remove all edges connecting to the ignore label, because # they introduce short-cut lifted edges valid_edges = (uv_ids != 0).all(axis=1) uv_ids = uv_ids[valid_edges] # if we only had a single edge to ignore label, we can end up # with empty uv-ids at this point. # if so, return None if uv_ids.size == 0: return None, None, None local_probs = local_probs[valid_edges] sizes = sizes[valid_edges] # build the original graph and lifted objective # with lifted uv-ids lifted_uv_ids = feat.make_filtered_lifted_nh(rag, n_labels, uv_ids, lifted_nh) graph = nifty.graph.undirectedGraph(n_labels) graph.insertEdges(uv_ids) # we may not get any lifted edges, in this case, fall back to normal multicut if lifted_uv_ids.size == 0: _, merge_indicator = run_mc( graph, local_probs, uv_ids, weighting_exponent=weighting_exponent, edge_sizes=sizes if weight_mulitcut_edges else None) return uv_ids, merge_indicator, sizes # get features for the lifted edges lifted_feats = np.concatenate( [ # feat.ucm_features(n_labels, lifted_objective, local_probs), feat.clustering_features(graph, local_probs, lifted_uv_ids), feat.ucm_features(n_labels, uv_ids, lifted_uv_ids, local_probs), feat.region_features(ws, lifted_uv_ids, glia) ], axis=1) lifted_probs = lifted_rf.predict_proba(lifted_feats)[:, 1] _, merge_indicator = run_lmc( graph, uv_ids, lifted_uv_ids, local_probs, lifted_probs, local_sizes=sizes if weight_mulitcut_edges else None, weighting_exponent=weighting_exponent, with_ignore_edges=False) # we don't weight, because we might just have few lifted edges # and this would downvote the local edges significantly # weight the costs # n_local, n_lifted = len(uv_ids), len(lifted_uv_ids) # total = float(n_lifted) + n_local # local_costs *= (n_lifted / total) # lifted_costs *= (n_local / total) return uv_ids, merge_indicator, sizes
def extract_feats_and_labels(path, aff_key, ws_key, gt_key, mask_key, lifted_nh, offsets=[[-1, 0, 0], [0, -1, 0], [0, 0, -1]], n_threads=40): f = z5py.File(path) # load the watershed segmentation and compute rag ds_seg = f[ws_key] ds_seg.n_threads = n_threads seg = ds_seg[:] print(seg.shape) n_labels = int(seg.max()) + 1 rag = nrag.gridRag(seg, numberOfLabels=n_labels, numberOfThreads=n_threads) # load affinities and glia channel ds_affs = f[aff_key] ds_affs.n_threads = n_threads aff_slice = slice(0, len(offsets)) affs = ds_affs[aff_slice] if affs.dtype == np.dtype('uint8'): affs = affs.astype('float32') / 255. affs = 1. - affs n_chans = ds_affs.shape[0] glia_slice = slice(n_chans - 1, n_chans) glia = ds_affs[glia_slice] if glia.dtype == np.dtype('uint8'): glia = glia.astype('float32') / 255. # compute local probs from affinities print("Computing local probabilities") probs = nrag.accumulateAffinityStandartFeatures( rag, affs, offsets, numberOfThreads=n_threads)[:, 0] probs = np.nan_to_num(probs) # remove zero-label (== ignore label) from the graph, because it short-circuits # lifted edges uv_ids = rag.uvIds() valid_edges = (uv_ids != 0).all(axis=1) uv_ids = uv_ids[valid_edges] probs = probs[valid_edges] # compute the lifted graph and lifted features print("Computing lifted objective") lifted_uv_ids = feat.make_filtered_lifted_nh(rag, n_labels, uv_ids, lifted_nh) graph = nifty.graph.undirectedGraph(n_labels) graph.insertEdges(uv_ids) # TODO parallelize some of these print("Computing lifted features") features = np.concatenate( [ # feat.ucm_features(n_labels, lifted_objective, probs), feat.clustering_features(graph, probs, lifted_uv_ids), feat.ucm_features(n_labels, uv_ids, lifted_uv_ids, probs), feat.region_features(seg, lifted_uv_ids, glia) ], axis=1) # load mask and groundtruth ds_mask = f[mask_key] ds_mask.n_threads = n_threads mask = ds_mask[:] ds_gt = f[gt_key] ds_gt.n_threads = n_threads gt = ds_gt[:] gt[np.logical_not(mask)] = 0 # compute the edge labels and valid edges node_labels = nrag.gridRagAccumulateLabels(rag, gt) labels = (node_labels[lifted_uv_ids[:, 0]] != node_labels[lifted_uv_ids[:, 1]]).astype('uint8') valid_edges = (node_labels[lifted_uv_ids] != 0).all(axis=1) print(np.sum(valid_edges), "edges of", len(lifted_uv_ids), "are valid") assert features.shape[0] == labels.shape[0] # just for temporary inspection, deactivate ! import vigra vigra.writeHDF5(features, './feats_tmp.h5', 'data', chunks=True) vigra.writeHDF5(labels, './labs_tmp.h5', 'data', chunks=True) return features[valid_edges], labels[valid_edges]
def _agglomerate_block(blocking, block_id, ds_in, ds_out, config): fu.log("start processing block %i" % block_id) have_ignore_label = config['have_ignore_label'] use_mala_agglomeration = config.get('use_mala_agglomeration', True) threshold = config.get('threshold', 0.9) size_regularizer = config.get('size_regularizer', .5) invert_inputs = config.get('invert_inputs', False) offsets = config.get('offsets', None) bb = vu.block_to_bb(blocking.getBlock(block_id)) # load the segmentation / output seg = ds_out[bb] # check if this block is empty if np.sum(seg) == 0: fu.log_block_success(block_id) return # load the input data ndim_in = ds_in.ndim if ndim_in == 4: assert offsets is not None assert len(offsets) <= ds_in.shape[0] bb_in = (slice(0, len(offsets)),) + bb input_ = vu.normalize(ds_in[bb_in]) else: assert offsets is None input_ = vu.normalize(ds_in[bb]) if invert_inputs: input_ = 1. - input_ id_offset = int(seg[seg != 0].min()) # relabel the segmentation _, max_id, _ = relabelConsecutive(seg, out=seg, keep_zeros=True, start_label=1) seg = seg.astype('uint32') # construct rag rag = nrag.gridRag(seg, numberOfLabels=max_id + 1, numberOfThreads=1) # extract edge features if offsets is None: edge_features = nrag.accumulateEdgeMeanAndLength(rag, input_, numberOfThreads=1) else: edge_features = nrag.accumulateAffinityStandartFeatures(rag, input_, offsets, numberOfThreads=1) edge_features, edge_sizes = edge_features[:, 0], edge_features[:, -1] uv_ids = rag.uvIds() # set edges to ignore label to be maximally repulsive if have_ignore_label: ignore_mask = (uv_ids == 0).any(axis=1) edge_features[ignore_mask] = 1 # build undirected graph n_nodes = rag.numberOfNodes graph = nifty.graph.undirectedGraph(n_nodes) graph.insertEdges(uv_ids) if use_mala_agglomeration: node_labels = mala_clustering(graph, edge_features, edge_sizes, threshold) else: node_ids, node_sizes = np.unique(seg, return_counts=True) if node_ids[0] != 0: node_sizes = np.concatenate([np.array([0]), node_sizes]) n_stop = int(threshold * n_nodes) node_labels = agglomerative_clustering(graph, edge_features, node_sizes, edge_sizes, n_stop, size_regularizer) # run clusteting node_labels, max_id, _ = relabelConsecutive(node_labels, start_label=1, keep_zeros=True) fu.log("reduced number of labels from %i to %i" % (n_nodes, max_id + 1)) # project node labels back to segmentation seg = nrag.projectScalarNodeDataToPixels(rag, node_labels, numberOfThreads=1) seg = seg.astype('uint64') # add offset back to segmentation seg[seg != 0] += id_offset ds_out[bb] = seg # log block success fu.log_block_success(block_id)
def nearest_features(rag, affs): offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1]] affs_ = affs[:3] probs = nrag.accumulateAffinityStandartFeatures(rag, affs_, offsets)[:, 0] return probs