def segment_volume(self, pmaps):
if self.ws_2D:
# WS in 2D
ws = self.ws_dt_2D(pmaps)
else:
# WS in 3D
ws, _ = distance_transform_watershed(pmaps,
self.ws_threshold,
self.ws_sigma,
sigma_weights=self.ws_w_sigma,
min_size=self.ws_minsize)
rag = compute_rag(ws, 1)
# Computing edge features
features = nrag.accumulateEdgeMeanAndLength(
rag, pmaps, numberOfThreads=1) # DO NOT CHANGE numberOfThreads
probs = features[:, 0] # mean edge prob
edge_sizes = features[:, 1]
# Prob -> edge costs
costs = transform_probabilities_to_costs(probs,
edge_sizes=edge_sizes,
beta=self.beta)
# Creating graph
graph = nifty.graph.undirectedGraph(rag.numberOfNodes)
graph.insertEdges(rag.uvIds())
# Solving Multicut
node_labels = multicut_kernighan_lin(graph, costs)
return nifty.tools.take(node_labels, ws)
def segment_volume_mc(pmaps,
threshold=0.4,
sigma=2.0,
beta=0.6,
ws=None,
sp_min_size=100):
if ws is None:
ws = distance_transform_watershed(pmaps,
threshold,
sigma,
min_size=sp_min_size)[0]
rag = compute_rag(ws, 1)
features = nrag.accumulateEdgeMeanAndLength(rag, pmaps, numberOfThreads=1)
probs = features[:, 0] # mean edge prob
edge_sizes = features[:, 1]
costs = transform_probabilities_to_costs(probs,
edge_sizes=edge_sizes,
beta=beta)
graph = nifty.graph.undirectedGraph(rag.numberOfNodes)
graph.insertEdges(rag.uvIds())
node_labels = multicut_kernighan_lin(graph, costs)
return nifty.tools.take(node_labels, ws)
def segment_mc(pred, seg, delta):
rag = feats.compute_rag(seg)
edge_probs = embed.edge_probabilities_from_embeddings(
pred, seg, rag, delta)
edge_sizes = feats.compute_boundary_mean_and_length(rag, pred[0])[:, 1]
costs = mc.transform_probabilities_to_costs(edge_probs,
edge_sizes=edge_sizes)
mc_seg = mc.multicut_kernighan_lin(rag, costs)
mc_seg = feats.project_node_labels_to_pixels(rag, mc_seg)
return mc_seg
def multicut_from_probas(segmentation, edges, edge_weights):
rag = compute_rag(segmentation)
edge_dict = dict(zip(list(map(tuple, edges)), edge_weights))
costs = np.empty(len(edge_weights))
for i, neighbors in enumerate(rag.uvIds()):
if tuple(neighbors) in edge_dict:
costs[i] = edge_dict[tuple(neighbors)]
else:
costs[i] = edge_dict[(neighbors[1], neighbors[0])]
costs = transform_probabilities_to_costs(costs)
node_labels = multicut_kernighan_lin(rag, costs)
return project_node_labels_to_pixels(rag, node_labels).squeeze()
def supervoxel_merging(mem, sv, beta=0.5, verbose=False):
rag = feats.compute_rag(sv)
costs = feats.compute_boundary_features(rag, mem)[:, 0]
edge_sizes = feats.compute_boundary_mean_and_length(rag, mem)[:, 1]
costs = mc.transform_probabilities_to_costs(costs,
edge_sizes=edge_sizes,
beta=beta)
node_labels = mc.multicut_kernighan_lin(rag, costs)
segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
return segmentation
def segment_volume_lmc_from_seg(boundary_pmaps,
nuclei_seg,
threshold=0.4,
sigma=2.0,
sp_min_size=100):
watershed = distance_transform_watershed(boundary_pmaps,
threshold,
sigma,
min_size=sp_min_size)[0]
# compute the region adjacency graph
rag = compute_rag(watershed)
# compute the edge costs
features = compute_boundary_mean_and_length(rag, boundary_pmaps)
costs, sizes = features[:, 0], features[:, 1]
# transform the edge costs from [0, 1] to [-inf, inf], which is
# necessary for the multicut. This is done by intepreting the values
# as probabilities for an edge being 'true' and then taking the negative log-likelihood.
# in addition, we weight the costs by the size of the corresponding edge
# we choose a boundary bias smaller than 0.5 in order to
# decrease the degree of over segmentation
boundary_bias = .6
costs = transform_probabilities_to_costs(costs,
edge_sizes=sizes,
beta=boundary_bias)
max_cost = np.abs(np.max(costs))
lifted_uvs, lifted_costs = lifted_problem_from_segmentation(
rag,
watershed,
nuclei_seg,
overlap_threshold=0.2,
graph_depth=4,
same_segment_cost=5 * max_cost,
different_segment_cost=-5 * max_cost)
# solve the full lifted problem using the kernighan lin approximation introduced in
# http://openaccess.thecvf.com/content_iccv_2015/html/Keuper_Efficient_Decomposition_of_ICCV_2015_paper.html
node_labels = lmc.lifted_multicut_kernighan_lin(rag, costs, lifted_uvs,
lifted_costs)
lifted_segmentation = project_node_labels_to_pixels(rag, node_labels)
return lifted_segmentation
def map_to_lifted_costs(path, exp_path, semantic_label_key, lifted_nh_key, lifted_cost_key, weight=1.5):
with z5py.File(exp_path, 'r') as f:
lifted_uvs = f[lifted_nh_key][:]
ves_feats = f['region_features/vesicles'][:, 1]
den_feats = f['region_features/dendrites'][:, 1]
with z5py.File(path, 'r') as f:
node_labels = f[semantic_label_key][:]
node_feats = np.zeros(len(node_labels), dtype='float32')
node_feats[node_labels == 1] = ves_feats[node_labels == 1]
node_feats[node_labels == 2] = den_feats[node_labels == 2]
lifted_costs = node_feats[lifted_uvs[:, 0]] * node_feats[lifted_uvs[:, 1]]
lifted_costs = transform_probabilities_to_costs(lifted_costs)
lifted_costs *= weight
with z5py.File(exp_path, 'a') as f:
ds = f.require_dataset(lifted_cost_key, shape=lifted_costs.shape, compression='gzip',
chunks=lifted_costs.shape, dtype=lifted_costs.dtype)
ds[:] = lifted_costs
def test_transform_probabilities_to_costs(self):
from elf.segmentation.multicut import transform_probabilities_to_costs
probs = np.random.rand(100).astype('float32')
costs = transform_probabilities_to_costs(probs)
self.assertTrue(np.isfinite(costs).all())
affs = np.transpose(affs.cpu().numpy(), (1, 0, 2, 3))
gt_affs = np.transpose(gt_affs.cpu().numpy(), (1, 0, 2, 3))
seg = seg.cpu().numpy()
gt_seg = gt_seg.cpu().numpy()
boundary_input = np.mean(affs, axis=0)
gt_boundary_input = np.mean(gt_affs, axis=0)
rag = feats.compute_rag(seg)
# edges rag.uvIds() [[1, 2], ...]
costs = feats.compute_affinity_features(rag, affs, offsets)[:, 0]
gt_costs = calculate_gt_edge_costs(rag.uvIds(), seg.squeeze(),
gt_seg.squeeze())
edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
gt_edge_sizes = feats.compute_boundary_mean_and_length(rag,
gt_boundary_input)[:, 1]
costs = mc.transform_probabilities_to_costs(costs, edge_sizes=edge_sizes)
gt_costs = mc.transform_probabilities_to_costs(gt_costs, edge_sizes=edge_sizes)
node_labels = mc.multicut_kernighan_lin(rag, costs)
gt_node_labels = mc.multicut_kernighan_lin(rag, gt_costs)
segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
gt_segmentation = feats.project_node_labels_to_pixels(rag, gt_node_labels)
plt.imshow(
np.concatenate(
(gt_segmentation.squeeze(), segmentation.squeeze(), seg.squeeze()),
axis=1))
plt.show()
def probs_to_costs(job_id, config_path):
fu.log("start processing job %i" % job_id)
fu.log("reading config from %s" % config_path)
with open(config_path, 'r') as f:
config = json.load(f)
input_path = config['input_path']
input_key = config['input_key']
output_path = config['output_path']
output_key = config['output_key']
features_path = config['features_path']
features_key = config['features_key']
# config for cost transformations
invert_inputs = config.get('invert_inputs', False)
transform_to_costs = config.get('transform_to_costs', True)
weight_edges = config.get('weight_edges', False)
weighting_exponent = config.get('weighting_exponent', 1.)
beta = config.get('beta', 0.5)
# additional node labels
node_labels = config.get('node_labels', None)
n_threads = config['threads_per_job']
fu.log("reading input from %s:%s" % (input_path, input_key))
with vu.file_reader(input_path) as f:
ds = f[input_key]
ds.n_threads = n_threads
# we might have 1d or 2d inputs, depending on input from features or random forest
slice_ = slice(None) if ds.ndim == 1 else (slice(None), slice(0, 1))
costs = ds[slice_].squeeze()
# normalize to range 0, 1
min_, max_ = costs.min(), costs.max()
fu.log('input-range: %f %f' % (min_, max_))
fu.log('%f +- %f' % (costs.mean(), costs.std()))
if invert_inputs:
fu.log("inverting probability inputs")
costs = 1. - costs
if transform_to_costs:
fu.log("converting probability inputs to costs")
if weight_edges:
fu.log("weighting edges by size")
# the edge sizes are at the last feature index
with vu.file_reader(features_path) as f:
ds = f[features_key]
n_features = ds.shape[1]
ds.n_threads = n_threads
edge_sizes = ds[:, n_features - 1:n_features].squeeze()
else:
fu.log("no edge weighting")
edge_sizes = None
costs = transform_probabilities_to_costs(
costs,
beta=beta,
edge_sizes=edge_sizes,
weighting_exponent=weighting_exponent)
# adjust edges of nodes with labels if given
if node_labels is not None:
fu.log("have node labels")
max_repulsive = 5 * costs.min()
max_attractive = 5 * costs.max()
fu.log("maximally attractive edge weight %f" % max_attractive)
fu.log("maximally repulsive edge weight %f" % max_repulsive)
with vu.file_reader(features_path, 'r') as f:
ds = f['s0/graph/edges']
ds.n_threads = n_threads
uv_ids = ds[:]
for mode, path_key in node_labels.items():
path, key = path_key
fu.log("applying node labels with mode %s from %s:%s" %
(mode, path, key))
with vu.file_reader(path, 'r') as f:
ds = f[key]
ds.n_threads = n_threads
labels = ds[:]
costs = _apply_node_labels(costs, uv_ids, mode, labels,
max_repulsive, max_attractive)
with vu.file_reader(output_path) as f:
ds = f[output_key]
ds.n_threads = n_threads
ds[:] = costs
fu.log_job_success(job_id)
def stitching_multicut(job_id, config_path):
fu.log("start processing job %i" % job_id)
fu.log("reading config from %s" % config_path)
with open(config_path, 'r') as f:
config = json.load(f)
problem_path = config['problem_path']
graph_key = config['graph_key']
features_key = config['features_key']
edge_key = config['edge_key']
output_path = config['output_path']
output_key = config['output_key']
beta1 = config['beta1']
beta2 = config['beta2']
n_threads = config.get('threads_per_job', 1)
time_limit = config.get('time_limit_solver', None)
agglomerator_key = config.get('agglomerator', 'kernighan-lin')
# load edges and features
fu.log("Loading features and edges")
with vu.file_reader(problem_path, 'r') as f:
ds = f[features_key]
ds.n_threads = n_threads
feats = ds[:]
ds = f[edge_key]
ds.n_threads = n_threads
stitch_edges = ds[:].astype('bool')
g = f[graph_key]
n_nodes = g.attrs['nodeMaxId'] + 1
ds = g['edges']
ds.n_threads = n_threads
uv_ids = ds[:]
assert uv_ids.max() < n_nodes, "%i, %i" % (int(uv_ids.max()), n_nodes)
assert len(stitch_edges) == len(uv_ids), "%i, %i" % (len(stitch_edges),
len(uv_ids))
assert len(feats) == len(uv_ids), "%i, %i" % (len(feats), len(uv_ids))
# load graph
fu.log("Building graph")
graph = nifty.graph.undirectedGraph(n_nodes)
graph.insertEdges(uv_ids)
n_edges = graph.numberOfEdges
# compute costs
fu.log("Computing costs")
feats, sizes = feats[:, 0], feats[:, -1]
costs = np.zeros(n_edges, dtype='float32')
costs[stitch_edges] = transform_probabilities_to_costs(
feats[stitch_edges], edge_sizes=sizes[stitch_edges], beta=beta1)
costs[~stitch_edges] = transform_probabilities_to_costs(
feats[~stitch_edges], edge_sizes=sizes[~stitch_edges], beta=beta2)
# solve multicut
solver = get_multicut_solver(agglomerator_key)
fu.log("Start multicut with solver %s" % agglomerator_key)
if time_limit is not None:
fu.log("With time limit %i s" % time_limit)
node_labels = solver(graph,
costs,
n_threads=n_threads,
time_limit=time_limit)
vigra.analysis.relabelConsecutive(node_labels,
out=node_labels,
keep_zeros=True,
start_label=1)
fu.log("Multicut done")
# write result
with vu.file_reader(output_path) as f:
chunks = (min(int(1e6), len(node_labels)), )
ds = f.require_dataset(output_key,
shape=node_labels.shape,
chunks=chunks,
compression='gzip',
dtype=node_labels.dtype)
ds.n_threads = n_threads
ds[:] = node_labels
# log success
fu.log_job_success(job_id)